Enable temporary shardy roudtrip import pass

tenstorrent · Feb 25, 2025 · 07b387a · 07b387a
1 parent a8769d7
commit 07b387a
Show file tree

Hide file tree

Showing 4 changed files with 502 additions and 0 deletions.
diff --git a/src/common/CMakeLists.txt b/src/common/CMakeLists.txt
@@ -18,6 +18,7 @@ add_dependencies(TTPJRTCommonDylibPlatform loguru)
 add_library(TTPJRTCommon
     "platform.cc"
     "module_builder.cc"
+    "shardy_roundtrip_pass.cc"
     "plugin_attributes.cc"
     "pjrt_implementation/api_bindings.cc"
     "pjrt_implementation/buffer_instance.cc"

diff --git a/src/common/module_builder.cc b/src/common/module_builder.cc
@@ -29,6 +29,10 @@
 #include "stablehlo/dialect/Version.h"
 #include "stablehlo/transforms/Passes.h"
 
+// shardy includes
+#include "common/shardy_roundtrip_pass.h"
+#include "shardy/dialect/sdy/ir/register.h"
+
 // tt-mlir includes
 #define TTMLIR_ENABLE_STABLEHLO
 #include "tt/runtime/runtime.h"
@@ -56,6 +60,7 @@ ModuleBuilder::ModuleBuilder()
 
   mlir::tt::registerAllDialects(registry);
   mlir::stablehlo::registerAllDialects(registry);
+  mlir::sdy::registerAllDialects(registry);
 
   mlir::func::registerAllExtensions(registry);
   mlir::tt::registerAllExtensions(registry);
@@ -130,6 +135,18 @@ void ModuleBuilder::convertFromVHLOToSHLO(
     return;
   }
 
+  // This is a temporary solution for the "roundtrip" mlir from openXLA.
+  // Once openXLA natively supports Shardy, we can remvoe following import
+  // passes.
+  mlir::PassManager shardy_pm(mlir_module.get()->getName());
+  tt::pjrt::sdy::addSdyRoundTripImportPipeline(shardy_pm);
+
+  if (mlir::failed(shardy_pm.run(mlir_module.get()))) {
+    DLOG_F(ERROR, "Failed to convert from Shardy roundtrip import pass module");
+    m_status = tt_pjrt_status::kInternal;
+    return;
+  }
+
   DLOG_F(LOG_DEBUG, "SHLO Module:");
   printModule(mlir_module);
 }