Merge pull request #32 from Xilinx/christopher.FXML-1991_reciprocal_constant_folding

[FXML-1991] Reciprocal Constant Folding For a Splat Tensor

Author: cmcgirr-amd

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

@@ -1057,6 +1057,8 @@ def Tosa_ReciprocalOp : Tosa_Op<"reciprocal", [
   let results = (outs
     Tosa_Tensor:$output
   );
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

@@ -630,6 +630,26 @@ OpFoldResult MulOp::fold(ArrayRef<Attribute> operands) {
   return mulBinaryFolder(lhsAttr, rhsAttr, lhsTy, getShift());
 }
 
+OpFoldResult ReciprocalOp::fold(ArrayRef<Attribute> operands) {
+  auto constantAttr = dyn_cast_or_null<DenseElementsAttr>(operands[0]);
+  auto lhsTy = dyn_cast<RankedTensorType>(getInput1().getType());
+
+  if (!lhsTy || !constantAttr) {
+    return {};
+  }
+
+  if (!constantAttr.isSplat()) {
+    return {};
+  }
+
+  auto floatVal = constantAttr.getSplatValue<llvm::APFloat>();
+
+  auto recipAttr = FloatAttr::get(lhsTy.getElementType(), 1.0);
+  APFloat recip = recipAttr.getValue();
+  recip.divide(floatVal, APFloat::rmNearestTiesToEven);
+  return DenseElementsAttr::get(lhsTy, recip);
+}
+
 OpFoldResult SubOp::fold(ArrayRef<Attribute> operands) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();

mlir/test/Dialect/Tosa/constant-op-fold.mlir

@@ -328,6 +328,29 @@ func.func @fold_mul_splat_f32() -> tensor<10xf32> {
 
 // -----
 
+// CHECK-LABEL: @fold_reciprocal_splat_f32
+func.func @fold_reciprocal_splat_f32() -> tensor<f32> {
+  %half = "tosa.const"() {value = dense<0.5> : tensor<f32>} : () -> tensor<f32>
+  %recp = "tosa.reciprocal"(%half) : (tensor<f32>) -> tensor<f32>
+  // CHECK: %[[CST:.*]] = "tosa.const"() {value = dense<2.000000e+00> : tensor<f32>}
+  // CHECK: return %[[CST]]
+  return %recp : tensor<f32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_reciprocal_splat_zero_f32
+func.func @fold_reciprocal_splat_zero_f32() -> tensor<f32> {
+  %zero = "tosa.const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
+  %recp = "tosa.reciprocal"(%zero) : (tensor<f32>) -> tensor<f32>
+  // 0x7F800000 represents +inf as we have computed 1/0
+  // CHECK: %[[CST:.*]] = "tosa.const"() {value = dense<0x7F800000> : tensor<f32>}
+  // CHECK: return %[[CST]]
+  return %recp : tensor<f32>
+}
+
+// -----
+
 // CHECK-LABEL: @fold_sub_zero_rhs_f32
 func.func @fold_sub_zero_rhs_f32(%arg0: tensor<f32>) -> tensor<f32> {
   %zero = "tosa.const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>