Add missing pass class definition

openxla · Dec 17, 2024 · a56b4ed · a56b4ed
1 parent 3004928
commit a56b4ed
Showing 1 changed file with 37 additions and 0 deletions.
diff --git a/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/transforms/StablehloComplexMathExpander.cpp
@@ -66,6 +66,43 @@ static Value getConstantLikeSmallestNormalizedValue(OpBuilder &b, Location loc,
       val);
 }
 
+//===----------------------------------------------------------------------===//
+// Pass
+//===----------------------------------------------------------------------===//
+
+struct StablehloComplexMathExpanderPass
+    : public impl::StablehloComplexMathExpanderPassBase<
+          StablehloComplexMathExpanderPass> {
+  StablehloComplexMathExpanderPass()
+      : StablehloComplexMathExpanderPassBase<
+            StablehloComplexMathExpanderPass>() {}
+ public:
+  LogicalResult initialize(MLIRContext *context) override {
+    // auto targetVersion = validateTargetVersion(targetVersionOption);
+
+    config.useTopDownTraversal = true;
+    RewritePatternSet patterns_(context);
+    populateStablehloComplexMathExpanderPatterns(&patterns_, context/*,
+                                                                        targetVersion*/);
+    patterns = std::move(patterns_);
+    return success();
+  }
+
+  void runOnOperation() override {
+    auto func = getOperation();
+    if (failed(applyPatternsAndFoldGreedily(func, patterns, config))) {
+      func.emitError(
+          "Failed to converge StableHLOComplexMathExpanderPass in ")
+          << config.maxIterations << " iterations";
+      signalPassFailure();
+    }
+  }
+
+ private:
+  FrozenRewritePatternSet patterns;
+  GreedyRewriteConfig config;
+};
+
 #include "stablehlo/transforms/StablehloComplexMathExpanderPatterns.h.inc"
 
 }  // namespace