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

[vzakhari]

Summary

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

I am not sure about this change. The backends (e.g. LLVM) may be better place for this,
especially if in some target implementations of libm pow(x, 0.75) is faster.
Please let me know what you think.

[kiranchandramohan]

I can confirm similar speedups are obtained on the AArch64 platform.

I think we will get the same transformation once the fast-math and related flags are enabled. I could see the transformation happening in classic-flang with Ofast and I see that the code is in Selection Dag https://github.com/llvm/llvm-project/blob/0589038a6ff12be452be2a2011ac23d1c0d8e7dc/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp#L15922

I guess since the code is not in a Target Specific portion of LLVM this might be an OK simplification to have here. Although LLVM users might not really require it (since the transformation is in LLVM), non-LLVM users (if any) might find it useful.

So, I am +1 for having it here.

[vzakhari]

Kiran, thank you for pointing out the code in selection dag! I should have checked it :)

Edit: I meant to say that the selection dag is not target dependent, so such a canonicalization should be always profitable and it makes sense to have it in MLIR.

[kiranchandramohan]

OK. LGTM.

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

PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
   }
 
   // Replace `pow(x, -0.5)` with `rsqrt(x)`.
   if (isExponentValue(-0.5)) {
     rewriter.replaceOpWithNewOp<math::RsqrtOp>(op, x);
     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();
 }
 
 //----------------------------------------------------------------------------//
 // FPowIOp/IPowIOp strength reduction.
 //----------------------------------------------------------------------------//
 
 namespace {

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

func.func @pow_rsqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
   // CHECK: return %[[SCALAR]], %[[VECTOR]]
   %c = arith.constant -0.5 : f32
   %v = arith.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_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>
 // CHECK-SAME: -> (i32, vector<4xi32>) {
 func.func @ipowi_zero_exp(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>) {
   // CHECK: %[[CST_S:.*]] = arith.constant 1 : i32
   // CHECK: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
   // CHECK: return %[[CST_S]], %[[CST_V]]