Make math.fpowi rewrite more systematic and test it. (iree-org#20293)

This is a follow-up to iree-org#20281 which lacked a test, and had left the c++
logic in a state that was prone to unintentional regression. Also, this
ensures that the `math.fpowi` rewrite happens even with
`-iree-codegen-gpu-native-math-precision`, which is intentional as
`math.fpowi` is not really a math function, like the comment explains.

Signed-off-by: Benoit Jacob <[email protected]>

Author: bjacob

compiler/src/iree/compiler/Codegen/Common/MathTransformPass.cpp

@@ -63,6 +63,14 @@ static void populateMathFunctionsRewritePatterns(
 
 static bool predicateRewrite(StringRef name,
                              IREE::HAL::ExecutableTargetAttr target) {
+  if (name == math::FPowIOp::getOperationName()) {
+    // math.fpowi is a special op: it isn't really a "math function", rather
+    // it is generally used with a constant exponent that is a small integer,
+    // and it is then a shorthand for a few multiplications. That rewrite needs
+    // to happen to prevent falling back on a more expensive, more general
+    // implementation like math.powf.
+    return true;
+  }
   if (clNativeMathPrecision) { // Legacy.
     if (name == math::Exp2Op::getOperationName() ||
         name == math::RoundEvenOp::getOperationName()) {
@@ -72,9 +80,7 @@ static bool predicateRewrite(StringRef name,
   if (isROCMBackend(target)) {
     // On ROCm, we do not need most rewrites as we can generally bottom out on
     // either device library functions, or handling of intrinsics in AMDGPU.
-    // An important exception is math.fpowi, which needs to get rewritten to
-    // multiplications.
-    return llvm::is_contained({math::FPowIOp::getOperationName()}, name);
+    return false;
   }
   // Currently enable all non-approximative rewrites.
   return true;

compiler/src/iree/compiler/Codegen/Common/test/math_transform.mlir

@@ -91,3 +91,32 @@ func.func @no_approx_on_rocm(%arg0: f16) -> f16 attributes {
   %11 = math.erf %10 : f16
   return %11 : f16
 }
+
+// -----
+
+// CHECK-LABEL: @rewrite_fpowi_const_exponent_on_llvmcpu
+func.func @rewrite_fpowi_const_exponent_on_llvmcpu(%arg0: f32) -> f32 attributes {
+  hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz"}>
+} {
+  // math.fpowi with constant exponent should always be rewritten to muls.
+  // CHECK-NOT:    math.fpowi
+  // CHECK:        arith.mulf {{.*}} : f32
+  %c2 = arith.constant 2 : i32
+  %0 = math.fpowi %arg0, %c2 : f32, i32
+  return %0 : f32
+}
+
+// -----
+
+// CHECK-LABEL: @rewrite_fpowi_const_exponent_on_rocm
+func.func @rewrite_fpowi_const_exponent_on_rocm(%arg0: f32) -> f32 attributes {
+  hal.executable.target =  #hal.executable.target<"rocm", "rocm-hsaco-fb", {}>
+} {
+  // math.fpowi with constant exponent should always be rewritten to muls.
+  // CHECK-NOT:    math.fpowi
+  // CHECK:        arith.mulf {{.*}} : f32
+  // CHECK:        arith.mulf {{.*}} : f32
+  %c2 = arith.constant 3 : i32
+  %0 = math.fpowi %arg0, %c2 : f32, i32
+  return %0 : f32
+}