[NFC] Add individual switches to polynomial approximation (#19697)

Give more freedom when woking with PolynomialApproximationPass, by adding individual switches to each of the approximation ops so they can be fine tuned for different architectures. --------- Co-authored-by: Jakub Kuderski <[email protected]>
iree-org · Feb 4, 2025 · 69b3ebb · 69b3ebb
1 parent a808900
commit 69b3ebb
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 15 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.td b/compiler/src/iree/compiler/Codegen/Common/Passes.td
@@ -529,6 +529,12 @@ def PadDynamicAllocPass :
 def PolynomialApproximationPass :
     Pass<"iree-codegen-polynomial-approximation", ""> {
   let summary = "Convert math operations to their polynomial approximation";
+  let options = [
+    ListOption<"noApproxOps", "no-approx-ops", "std::string",
+           [{List of operations that should not be approximated.\n"
+           "As of now, possible options are:\n"
+           "\ttan, sinh, cosh, asinh, acosh, atanh, powf, fpowf, erf\n}]>,
+  ];
 }
 
 def PropagateDispatchSizeBoundsPass :

diff --git a/compiler/src/iree/compiler/Codegen/Common/PolynomialApproximationPass.cpp b/compiler/src/iree/compiler/Codegen/Common/PolynomialApproximationPass.cpp
@@ -7,6 +7,7 @@
 #include "iree/compiler/Codegen/Common/Passes.h"
 #include "mlir/Dialect/Math/Transforms/Approximation.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
@@ -25,28 +26,47 @@ static llvm::cl::opt<bool> clNativeMathPrecision(
 
 namespace {
 
+static void populateErfPattern(RewritePatternSet &patterns) {
+  if (clNativeMathPrecision) {
+    patterns.add<math::ErfPolynomialApproximation>(patterns.getContext());
+  } else {
+    populateExpandExp2FPattern(patterns);
+    populateMathPolynomialApproximationPatterns(patterns);
+    populateExpandRoundEvenPattern(patterns);
+  }
+}
+
 /// math dialect elementry functions -> polynomial form.
 class PolynomialApproximationPass final
     : public impl::PolynomialApproximationPassBase<
           PolynomialApproximationPass> {
+public:
+  using Base::Base;
+
   void runOnOperation() override {
+    using PatternFunction = llvm::function_ref<void(RewritePatternSet &)>;
+    // Order matters here.
+    llvm::SmallVector<std::pair<StringRef, PatternFunction>> patternMap = {
+        {"tan", populateExpandTanPattern},
+        {"sinh", populateExpandSinhPattern},
+        {"cosh", populateExpandCoshPattern},
+        {"asinh", populateExpandAsinhPattern},
+        {"acosh", populateExpandAcoshPattern},
+        {"atanh", populateExpandAtanhPattern},
+        {"powf", populateExpandPowFPattern},
+        {"fpowi", populateExpandFPowIPattern},
+        {"erf", populateErfPattern},
+    };
+
     RewritePatternSet mathPatterns(&getContext());
-    populateExpandTanPattern(mathPatterns);
-    populateExpandSinhPattern(mathPatterns);
-    populateExpandCoshPattern(mathPatterns);
-    populateExpandAsinhPattern(mathPatterns);
-    populateExpandAcoshPattern(mathPatterns);
-    populateExpandAtanhPattern(mathPatterns);
-    populateExpandPowFPattern(mathPatterns);
-    populateExpandFPowIPattern(mathPatterns);
-
-    if (clNativeMathPrecision) {
-      mathPatterns.add<math::ErfPolynomialApproximation>(&getContext());
-    } else {
-      populateExpandExp2FPattern(mathPatterns);
-      populateMathPolynomialApproximationPatterns(mathPatterns);
-      populateExpandRoundEvenPattern(mathPatterns);
+
+    for (const auto &[fnName, populateFn] : patternMap) {
+      // Skip any ops in the "do not convert" list.
+      if (!llvm::is_contained(noApproxOps, fnName)) {
+        populateFn(mathPatterns);
+      }
     }
+
     if (failed(
             applyPatternsGreedily(getOperation(), std::move(mathPatterns)))) {
       return signalPassFailure();