Fix a bug in algebraic simplification, and enable the tests.

[bakhtiyarneyman]

Fix a typo.

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

PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
   }
 
   // Replace `pow(x, 2.0)` with `x * x`.
   if (isExponentValue(2.0)) {
     rewriter.replaceOpWithNewOp<MulFOp>(op, ValueRange({x, x}));
     return success();
   }
 
-  // Replace `pow(x, 2.0)` with `x * x * x`.
+  // Replace `pow(x, 3.0)` with `x * x * x`.
   if (isExponentValue(3.0)) {
     Value square = rewriter.create<MulFOp>(op.getLoc(), ValueRange({x, x}));
     rewriter.replaceOpWithNewOp<MulFOp>(op, ValueRange({x, square}));
     return success();
   }
 
   // Replace `pow(x, -1.0)` with `1.0 / x`.
   if (isExponentValue(-1.0)) {
     Value one = rewriter.create<ConstantOp>(
         loc, rewriter.getFloatAttr(getElementTypeOrSelf(op.getType()), 1.0));
     rewriter.replaceOpWithNewOp<DivFOp>(op, ValueRange({bcast(one), x}));
     return success();
   }
 
-  // Replace `pow(x, -2.0)` with `sqrt(x)`.
-  if (isExponentValue(-1.0)) {
+  // Replace `pow(x, 0.5)` with `sqrt(x)`.
+  if (isExponentValue(0.5)) {
     rewriter.replaceOpWithNewOp<math::SqrtOp>(op, x);
     return success();
   }
 
+  // Replace `pow(x, -0.5)` with `sqrt(x)`.
+  if (isExponentValue(-0.5)) {
+    rewriter.replaceOpWithNewOp<math::RsqrtOp>(op, x);
+    return success();
+  }
+
   return failure();
 }
 
 //----------------------------------------------------------------------------//
 
 void mlir::populateMathAlgebraicSimplificationPatterns(
     RewritePatternSet &patterns) {
   patterns.add<PowFStrengthReduction>(patterns.getContext());
 }

mlir/test/Dialect/Math/algebraic-simplification.mlir

func @pow_recip(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
   // CHECK: %[[VECTOR:.*]] = divf %[[CST_V]], %arg1
   // CHECK: return %[[SCALAR]], %[[VECTOR]]
   %c = constant -1.0 : f32
   %v = constant dense <-1.0> : vector<4xf32>
   %0 = math.powf %arg0, %c : f32
   %1 = math.powf %arg1, %v : vector<4xf32>
   return %0, %1 : f32, vector<4xf32>
 }
+
+// CHECK-LABEL: @pow_sqrt
+func @pow_sqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
+  // CHECK: %[[SCALAR:.*]] = math.sqrt %arg0
+  // CHECK: %[[VECTOR:.*]] = math.sqrt %arg1
+  // CHECK: return %[[SCALAR]], %[[VECTOR]]
+  %c = constant 0.5 : f32
+  %v = constant dense <0.5> : vector<4xf32>
+  %0 = math.powf %arg0, %c : f32
+  %1 = math.powf %arg1, %v : vector<4xf32>
+  return %0, %1 : f32, vector<4xf32>
+}
+
+// CHECK-LABEL: @pow_rsqrt
+func @pow_rsqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
+  // CHECK: %[[SCALAR:.*]] = math.rsqrt %arg0
+  // CHECK: %[[VECTOR:.*]] = math.rsqrt %arg1
+  // CHECK: return %[[SCALAR]], %[[VECTOR]]
+  %c = constant -0.5 : f32
+  %v = constant dense <-0.5> : vector<4xf32>
+  %0 = math.powf %arg0, %c : f32
+  %1 = math.powf %arg1, %v : vector<4xf32>
+  return %0, %1 : f32, vector<4xf32>
+}