From 714270a922e81b5092d24782d76fa3d54d1d6dc5 Mon Sep 17 00:00:00 2001
From: Yuanqiang Liu <liuyuanqiang.yqliu@bytedance.com>
Date: Wed, 17 Jul 2024 00:05:11 +0800
Subject: [PATCH] [Stablehlo] legalize deprecated ops to stablehlo ops (#3543)

---
 lib/Dialect/TorchConversion/Transforms/Passes.cpp | 7 +++++++
 lib/InitAll.cpp                                   | 2 ++
 2 files changed, 9 insertions(+)

diff --git a/lib/Dialect/TorchConversion/Transforms/Passes.cpp b/lib/Dialect/TorchConversion/Transforms/Passes.cpp
index 4cdadb578..42ec495d9 100644
--- a/lib/Dialect/TorchConversion/Transforms/Passes.cpp
+++ b/lib/Dialect/TorchConversion/Transforms/Passes.cpp
@@ -166,5 +166,12 @@ void TorchConversion::createTorchBackendToStablehloBackendPipeline(
   pm.addNestedPass<func::FuncOp>(
       stablehlo::createStablehloCanonicalizeDynamismPass());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+
+  // Legalize deprecated ops to Stablehlo ops
+  stablehlo::StablehloLegalizeDeprecatedOpsPassOptions stablehloOptions;
+  stablehloOptions.failOnUnusedOps = false;
+  pm.addNestedPass<func::FuncOp>(
+      stablehlo::createStablehloLegalizeDeprecatedOpsPass(stablehloOptions));
+  pm.addPass(createCanonicalizerPass());
 }
 #endif
diff --git a/lib/InitAll.cpp b/lib/InitAll.cpp
index 7ade22b05..c9638c835 100644
--- a/lib/InitAll.cpp
+++ b/lib/InitAll.cpp
@@ -67,6 +67,8 @@ void mlir::torch::registerAllPasses() {
   mlir::stablehlo::registerStablehloAggressiveSimplificationPass();
   mlir::stablehlo::registerStablehloRefineShapesPass();
   mlir::stablehlo::registerStablehloConvertToSignlessPass();
+  mlir::stablehlo::registerShapeLegalizeToStablehloPass();
+  mlir::stablehlo::registerStablehloLegalizeDeprecatedOpsPass();
 #endif
 
 #ifdef TORCH_MLIR_ENABLE_REFBACKEND