From 2a79743e8614ed842c8c31324675eaabbd2b9da1 Mon Sep 17 00:00:00 2001
From: Yan Xu <yancey1989@gmail.com>
Date: Fri, 15 Sep 2023 16:20:19 +0800
Subject: [PATCH] support stablehlo in torch-mlir-opt binary

---
 pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp | 5 +++++
 1 file changed, 5 insertions(+)

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
index bbc6c82f6ee..41878401108 100644
--- 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);