support bfloat16

EnzymeAD · Mar 12, 2024 · 222ffa0 · 222ffa0
1 parent bd82051
commit 222ffa0
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -629,6 +629,12 @@ DenseElementsAttr fromTensor(stablehlo::Tensor inp) {
   auto type = inp.getType();
   auto elemType = type.getElementType();
 
+  if (elemType.isBF16()) {
+    auto floatValues =
+        ArrayRef((char *)inp.getData(), 2 * inp.getNumElements());
+    return DenseFPElementsAttr::getFromRawBuffer(type, floatValues);
+  }
+
   if (elemType.isF32()) {
     auto floatValues = ArrayRef((float *)inp.getData(), inp.getNumElements());
     return DenseFPElementsAttr::get(type, floatValues);

diff --git a/test/lit_tests/convertconst.mlir b/test/lit_tests/convertconst.mlir
@@ -0,0 +1,14 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+  func.func @main() -> tensor<4xbf16> {
+    %concat = stablehlo.constant dense<3.140000e+00> : tensor<4xf32>
+    %conv = stablehlo.convert %concat : (tensor<4xf32>) -> tensor<4xbf16>
+    return %conv : tensor<4xbf16>
+  }
+}
+
+// CHECK:  func.func @main() -> tensor<4xbf16> {
+// CHECK-NEXT:    %0 = stablehlo.constant dense<3.140630e+00> : tensor<4xbf16>
+// CHECK-NEXT:    return %0 : tensor<4xbf16>
+// CHECK-NEXT:  }