Add missing expansion patterns for math.powf (#14614)

Following the same approach as for other math ops. This will prevent the powf operation to be scalarized on CPU. Fixes #13547
iree-org · Sep 6, 2023 · c2f69fd · c2f69fd
1 parent 19a28bd
commit c2f69fd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/PolynomialApproximationPass.cpp b/compiler/src/iree/compiler/Codegen/Common/PolynomialApproximationPass.cpp
@@ -30,6 +30,7 @@ class PolynomialApproximationPass
     RewritePatternSet mathPatterns(&getContext());
     populateExpandTanPattern(mathPatterns);
     populateExpandExp2FPattern(mathPatterns);
+    populateExpandPowFPattern(mathPatterns);
 
     if (clNativeMathPrecision) {
       mathPatterns.add<math::ErfPolynomialApproximation>(&getContext());

diff --git a/compiler/src/iree/compiler/Codegen/Common/test/polynomial_approximation.mlir b/compiler/src/iree/compiler/Codegen/Common/test/polynomial_approximation.mlir
@@ -1,8 +1,17 @@
-// RUN: iree-opt --iree-codegen-polynomial-approximation %s | FileCheck %s
+// RUN: iree-opt --iree-codegen-polynomial-approximation --split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: @polynomial_tan
 func.func @polynomial_tan(%arg0: f32) -> f32 {
   // CHECK-NOT: math.tan
-  %0 = math.tan %arg0 : f32 
+  %0 = math.tan %arg0 : f32
+  return %0 : f32
+}
+
+// -----
+
+// CHECK-LABEL: @expanded_pow
+func.func @expanded_pow(%arg0: f32, %arg1: f32) -> f32 {
+  // CHECK-NOT: math.pow
+  %0 = math.powf %arg0, %arg1 : f32
   return %0 : f32
 }