Reconcile math.rsqrt expansion with upstream (#222)

Author: cferry-AMD

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

@@ -485,9 +485,8 @@ static LogicalResult convertRsqrtOp(math::RsqrtOp op,
 
   Location loc = op->getLoc();
   auto constOneFloat = createFloatConst(loc, operandTy, 1.0, rewriter);
-  auto sqrtOp = rewriter.create<math::SqrtOp>(loc, op->getOperand(0));
-  rewriter.replaceOpWithNewOp<arith::DivFOp>(op, operandTy,
-                                             ValueRange{constOneFloat, sqrtOp});
+  auto sqrtOp = rewriter.create<math::SqrtOp>(loc, operand);
+  rewriter.replaceOpWithNewOp<arith::DivFOp>(op, constOneFloat, sqrtOp);
   return success();
 }
 

mlir/test/Dialect/Math/expand-math.mlir

@@ -515,20 +515,48 @@ func.func @roundeven16(%arg: f16) -> f16 {
 
 // -----
 
+// CHECK-LABEL:   func.func @rsqrt
+// CHECK-SAME:     (%[[ARG:.*]]: f16)
+// CHECK-SAME:    -> f16
+// CHECK-DAG:     %[[CST:.*]] = arith.constant 1.000000e+00 : f16
+// CHECK-DAG:     %[[SQRT:.*]] = math.sqrt %[[ARG]] : f16
+// CHECK-DAG:     %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : f16
+// CHECK:         return %[[DIV]] : f16
+func.func @rsqrt16(%float: f16) -> (f16)  {
+  %float_result = math.rsqrt %float : f16
+  return %float_result : f16
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @rsqrt
 // CHECK-SAME:     (%[[ARG:.*]]: f32)
 // CHECK-SAME:    -> f32
 // CHECK-DAG:     %[[CST:.*]] = arith.constant 1.000000e+00 : f32
 // CHECK-DAG:     %[[SQRT:.*]] = math.sqrt %[[ARG]] : f32
 // CHECK-DAG:     %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : f32
 // CHECK:         return %[[DIV]] : f32
-func.func @rsqrt(%float: f32) -> (f32)  {
+func.func @rsqrt32(%float: f32) -> (f32)  {
   %float_result = math.rsqrt %float : f32
   return %float_result : f32
 }
 
 // -----
 
+// CHECK-LABEL:   func.func @rsqrt
+// CHECK-SAME:     (%[[ARG:.*]]: f64)
+// CHECK-SAME:    -> f64
+// CHECK-DAG:     %[[CST:.*]] = arith.constant 1.000000e+00 : f64
+// CHECK-DAG:     %[[SQRT:.*]] = math.sqrt %[[ARG]] : f64
+// CHECK-DAG:     %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : f64
+// CHECK:         return %[[DIV]] : f64
+func.func @rsqrt64(%float: f64) -> (f64)  {
+  %float_result = math.rsqrt %float : f64
+  return %float_result : f64
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @rsqrt_vec
 // CHECK-SAME:     (%[[ARG:.*]]: vector<5xf32>)
 // CHECK-SAME:    -> vector<5xf32>