[FXML-2939] Fix rsqrt folder for bfloat16 (#81)

Xilinx · Sep 14, 2023 · dd37d22 · dd37d22
1 parent faf28e5
commit dd37d22
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
@@ -588,6 +588,11 @@ struct TosaFoldConstantRSQRT
     auto floatVal = apFloatVal.convertToFloat();
     auto sqrtVal = std::sqrt(floatVal);
     APFloat apSqrtVal(sqrtVal);
+    // We fold only float32 and bfloat16, so we do not expect any precision loss
+    // for float32 and the tosa spec explicitly allows to implement bfloat16 as
+    // float32, so any precision loss on the conversion back is fine.
+    bool losesInfo = false;
+    apSqrtVal.convert(apFloatVal.getSemantics(), tosaRoundingMode, &losesInfo);
 
     // Compute the reciprocal
     return computeReciprocal(apSqrtVal, floatTy);
@@ -600,7 +605,7 @@ struct TosaFoldConstantRSQRT
   }
 
   bool isSupportedElementType(Type type) const {
-    return type.isBF16() || type.isF16() || type.isF32();
+    return type.isBF16() || type.isF32();
   }
 };
 

diff --git a/mlir/test/Dialect/Tosa/constant-rsqrt-opt.mlir b/mlir/test/Dialect/Tosa/constant-rsqrt-opt.mlir
@@ -123,6 +123,16 @@ func.func @rsqrt_fold() -> tensor<4x6xf32> {
   return %1 : tensor<4x6xf32>
 }
 
+// CHECK-LABEL: @rsqrt_fold_single_valued_bf16
+func.func @rsqrt_fold_single_valued_bf16() -> tensor<bf16> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}2.890630e-01{{.*}}tensor<bf16>
+  // CHECK-NOT: tosa.rsqrt
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<12.0> : tensor<bf16>} : () -> tensor<bf16>
+  %1 = "tosa.rsqrt"(%0) : (tensor<bf16>) -> tensor<bf16>
+  return %1 : tensor<bf16>
+}
+
 // CHECK-LABEL: @rsqrt_of_const_sparse
 // Sparse tensors are currently not supported
 func.func @rsqrt_of_const_sparse() -> tensor<32xbf16> {