support stablehlo in torch-mlir-opt binary

alibaba · Sep 15, 2023 · 2a79743 · 2a79743
1 parent da990f4
commit 2a79743
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp b/pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp
@@ -24,14 +24,19 @@
 #include "torch-mlir/Conversion/MhloPasses.h"
 #include "torch-mlir/InitAll.h"
 
+#include "mhlo/IR/hlo_ops.h"
+#include "mhlo/transforms/passes.h"
+#include "stablehlo/dialect/Register.h"
 using namespace mlir;
 
 int main(int argc, char** argv) {
   mlir::torch::registerTorchToMhloPasses();
   registerAllPasses();
   mlir::torch::registerAllPasses();
+  mlir::mhlo::registerStablehloLegalizeToHloPass();
 
   DialectRegistry registry;
+  mlir::stablehlo::registerAllDialects(registry);
   registerAllDialects(registry);
   mlir::torch::registerAllDialects(registry);