[mlir][math] Simplify pow(x, 0.75) into sqrt(sqrt(x)) * sqrt(x).

Trivial simplification for CPU2017/503.bwaves resulting in 3.89%
speed-up on icelake.

Differential Revision: https://reviews.llvm.org/D137351

Author: vzakhari

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

@@ -109,6 +109,15 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
     return success();
   }
 
+  // Replace `pow(x, 0.75)` with `sqrt(sqrt(x)) * sqrt(x)`.
+  if (isExponentValue(0.75)) {
+    Value pow_half = rewriter.create<math::SqrtOp>(op.getLoc(), x);
+    Value pow_quarter = rewriter.create<math::SqrtOp>(op.getLoc(), pow_half);
+    rewriter.replaceOpWithNewOp<arith::MulFOp>(
+        op, ValueRange{pow_half, pow_quarter});
+    return success();
+  }
+
   return failure();
 }
 

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

@@ -74,6 +74,22 @@ func.func @pow_rsqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>)
   return %0, %1 : f32, vector<4xf32>
 }
 
+// CHECK-LABEL: @pow_0_75
+func.func @pow_0_75(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
+  // CHECK: %[[SQRT1S:.*]] = math.sqrt %arg0
+  // CHECK: %[[SQRT2S:.*]] = math.sqrt %[[SQRT1S]]
+  // CHECK: %[[SCALAR:.*]] = arith.mulf %[[SQRT1S]], %[[SQRT2S]]
+  // CHECK: %[[SQRT1V:.*]] = math.sqrt %arg1
+  // CHECK: %[[SQRT2V:.*]] = math.sqrt %[[SQRT1V]]
+  // CHECK: %[[VECTOR:.*]] = arith.mulf %[[SQRT1V]], %[[SQRT2V]]
+  // CHECK: return %[[SCALAR]], %[[VECTOR]]
+  %c = arith.constant 0.75 : f32
+  %v = arith.constant dense <0.75> : vector<4xf32>
+  %0 = math.powf %arg0, %c : f32
+  %1 = math.powf %arg1, %v : vector<4xf32>
+  return %0, %1 : f32, vector<4xf32>
+}
+
 // CHECK-LABEL: @ipowi_zero_exp(
 // CHECK-SAME: %[[ARG0:.+]]: i32
 // CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>