Adding input sharding collection in ModuleBuilder #289

ajakovljevicTT · 2025-02-28T15:24:55Z

As part of adding support for multichip, adding collection of the information about the device sharding of the inputs of the the StableHLO graph during compilation. That way, we will know how to create MultiDevice input tensors when executing.

Edit: Also added output sharding collection.

github-actions · 2025-02-28T15:38:32Z

	Tests	Passed ✅	Skipped ⚠️	Failed
TT-XLA Tests	603 ran	432 passed	171 skipped	0 failed

Test	Result
No test annotations available

src/common/module_builder.cc

src/common/module_builder.h

src/common/module_builder.cc

codecov-commenter · 2025-03-04T09:53:16Z

Codecov Report

Attention: Patch coverage is 76.27119% with 14 lines in your changes missing coverage. Please review.

Project coverage is 78.65%. Comparing base (5f758e5) to head (b5925e5).
Report is 14 commits behind head on main.

✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/common/module_builder.cc	71.42%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #289      +/-   ##
==========================================
+ Coverage   77.97%   78.65%   +0.67%     
==========================================
  Files          21       21              
  Lines         990     1059      +69     
==========================================
+ Hits          772      833      +61     
- Misses        218      226       +8

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

kmitrovicTT · 2025-03-05T15:55:31Z

src/common/module_builder.cc

+void ModuleBuilder::collectInputShardings(
+    const mlir::OwningOpRef<mlir::ModuleOp> &module) {
+  DLOG_F(LOG_DEBUG, "ModuleBuilder::collectInputShardings");
+  m_input_shardings.clear();
+
+  module.get().walk([&](mlir::Operation *op) {
+    mlir::func::FuncOp funcOp = mlir::dyn_cast<mlir::func::FuncOp>(op);
+    mlir::ModuleOp moduleOp = mlir::dyn_cast<mlir::ModuleOp>(op);
+
+    if (!funcOp || !funcOp.isPublic()) {
+      return;
+    }
+
+    for (unsigned i = 0; i < funcOp.getNumArguments(); ++i) {
+
+      mlir::tt::sharding_utils::MeshSharding meshSharding;
+
+      auto shardingAttr = llvm::dyn_cast_if_present<mlir::StringAttr>(
+          funcOp.getArgAttr(i, mlir::tt::sharding_utils::kXlaShardingAttr));
+
+      if (!shardingAttr) {
+        m_input_shardings.push_back(meshSharding);
+        continue;
+      }
+
+      mlir::LogicalResult conversionResult =
+          fillMeshShardingFromGSPMDString(shardingAttr, meshSharding);
+
+      if (conversionResult.failed()) {
+        m_status = tt_pjrt_status::kInternal;
+        return;
+      }
+      m_input_shardings.push_back(meshSharding);
+    }
+  });
+}
+
+void ModuleBuilder::collectOutputShardings(
+    const mlir::OwningOpRef<mlir::ModuleOp> &module) {
+  DLOG_F(LOG_DEBUG, "ModuleBuilder::collectOutputShardings");
+  m_output_shardings.clear();
+
+  module.get().walk([&](mlir::Operation *op) {
+    mlir::func::FuncOp funcOp = mlir::dyn_cast<mlir::func::FuncOp>(op);
+    mlir::ModuleOp moduleOp = mlir::dyn_cast<mlir::ModuleOp>(op);
+
+    if (!funcOp || !funcOp.isPublic()) {
+      return;
+    }
+
+    for (unsigned i = 0; i < funcOp.getNumResults(); ++i) {
+
+      mlir::tt::sharding_utils::MeshSharding meshSharding;
+
+      auto shardingAttr = llvm::dyn_cast_if_present<mlir::StringAttr>(
+          funcOp.getResultAttr(i, mlir::tt::sharding_utils::kXlaShardingAttr));
+
+      if (!shardingAttr) {
+        m_output_shardings.push_back(meshSharding);
+        continue;
+      }
+
+      mlir::LogicalResult conversionResult =
+          fillMeshShardingFromGSPMDString(shardingAttr, meshSharding);
+
+      if (conversionResult.failed()) {
+        m_status = tt_pjrt_status::kInternal;
+        return;
+      }
+      m_output_shardings.push_back(meshSharding);
+    }
+  });
+}


These two look quite similar. The only differences are m_input_shardings vs m_output_shardings and getArgAttr instead of getResultAttr. Can we somehow extract it in a common function if it is not too cumbersome (like passing function pointers or something, let's not do that, it's an overkill)?

The issue I see here is with the getArgAttr and getResultAttr, I can use templates or just pass a bollean if it makes more sense? Not sure of the tradeoff between clarity and the code size here.

Yeah, that's what I though. Leave it like this.

You can do something like:

void ModuleBuilder::collectInputShardings( const mlir::OwningOpRef<mlir::ModuleOp> &module) { DLOG_F(LOG_DEBUG, "ModuleBuilder::collectInputShardings"); m_input_shardings.clear(); std::vector<mlir::StringAttr> gspmd_attributes; for (unsigned i = 0; i < funcOp.getNumArguments(); ++i) { gspmd_attributes.push_back(llvm::dyn_cast_if_present<mlir::StringAttr>( funcOp.getArgAttr(i, mlir::tt::sharding_utils::kXlaShardingAttr))); } mlir::LogicalResult result = createShardingsFromGSPMD(gspmd_attributes, m_input_shardings); if (!tt_pjrt_status_is_ok(result)) { m_status = tt_pjrt_status::kInternal; } } void ModuleBuilder::collectOutputShardings( const mlir::OwningOpRef<mlir::ModuleOp> &module) { DLOG_F(LOG_DEBUG, "ModuleBuilder::collectOutputShardings"); m_output_shardings.clear(); std::vector<mlir::StringAttr> gspmd_attributes; for (unsigned i = 0; i < funcOp.getNumResults(); ++i) { gspmd_attributes.push_back(llvm::dyn_cast_if_present<mlir::StringAttr>( funcOp.getResultAttr(i, mlir::tt::sharding_utils::kXlaShardingAttr))); } mlir::LogicalResult result = createShardingsFromGSPMD(gspmd_attributes, m_output_shardings); if (!tt_pjrt_status_is_ok(result)) { m_status = tt_pjrt_status::kInternal; } } mlir::LogicalResult ModuleBuilder::createShardingsFromGSPMD( const std::vector<mlir::StringAttr> &gspmd_attributes, std::vector<mlir::tt::sharding_utils::MeshSharding> &shardings) { for (mlir::StringAttr gspmd_attr : gspmd_attributes) { mlir::tt::sharding_utils::MeshSharding meshSharding; if (!gspmd_attr) { shardings.push_back(meshSharding); continue; } mlir::LogicalResult conversionResult = fillMeshShardingFromGSPMDString(gspmd_attr, meshSharding); if (conversionResult.failed()) { return llvm::LogicalResult::failure(); } shardings.push_back(meshSharding); } return llvm::LogicalResult::success(); }

src/common/module_builder.h

mrakitaTT · 2025-03-07T15:38:21Z

inc/common/pjrt_implementation/executable_image.h

@@ -16,6 +16,10 @@
 // tt-mlir includes
 #include "tt/runtime/types.h"

+// tt-mlir includes


merge with the group above, like:

// tt-mlir includes #define TTMLIR_ENABLE_STABLEHLO 1 #include "tt/runtime/types.h" #include "ttmlir/Conversion/StableHLOToTTIR/ShardingUtils.h"

mrakitaTT · 2025-03-07T15:39:52Z

src/common/module_builder.cc

@@ -28,10 +32,12 @@
 #include "stablehlo/dialect/Register.h"
 #include "stablehlo/dialect/Version.h"
 #include "stablehlo/transforms/Passes.h"
+#include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"


move into group below

mrakitaTT · 2025-03-07T15:41:59Z

src/common/module_builder.cc


 // tt-mlir includes
-#define TTMLIR_ENABLE_STABLEHLO
+#define TTMLIR_ENABLE_STABLEHLO 1


shouldn't be needed since you've already defined it in module_builder.h

mrakitaTT · 2025-03-07T15:43:24Z

src/common/module_builder.h

@@ -52,6 +65,12 @@ class ModuleBuilder {
  // scalar or not.
  void collectOutputTypes(const mlir::OwningOpRef<mlir::ModuleOp> &module);

+  // Collects the information about the sharding of specifc inputs.


typo specifc, below too

mrakitaTT · 2025-03-07T15:47:11Z

src/common/module_builder.cc

+  auto error =
+      meshSharding.convertGSPMDShardingToMeshSharding(shardingStr.getValue());
+  if (auto e = error.takeError()) {


Is auto necessary in these places?

mrakitaTT · 2025-03-07T15:49:34Z

src/common/module_builder.cc

+      return;
+    }
+
+    for (unsigned i = 0; i < funcOp.getNumArguments(); ++i) {


Specify the full type instead of unsigned

mrakitaTT · 2025-03-07T15:54:40Z

src/common/module_builder.cc

+          funcOp.getArgAttr(i, mlir::tt::sharding_utils::kXlaShardingAttr));
+
+      if (!shardingAttr) {
+        m_input_shardings.push_back(meshSharding);


Is this okay that we just put the empty sharding here?

If it is, please add a comment explaining why

mrakitaTT · 2025-03-07T16:00:19Z

src/common/module_builder.cc

@@ -134,6 +142,92 @@ void ModuleBuilder::convertFromVHLOToSHLO(
  printModule(mlir_module);
 }

+mlir::LogicalResult ModuleBuilder::fillMeshShardingFromGSPMDString(


Move this function implementation below the implementations for the collectInputShardings and collectOutputShardings since we followed top-down order in this file

mrakitaTT · 2025-03-07T16:06:05Z

src/common/module_builder.h

+  // Fills sharding_utils::MeshSharding object with sharding info stored in a
+  // StringAttribute.
+  mlir::LogicalResult fillMeshShardingFromGSPMDString(
+      mlir::StringAttr shardingStr,


In PJRT we followed a slight change from tt-mlir guidelines, we use the snake case for variables, please update in all changes in ModuleBuilder

ajakovljevicTT self-assigned this Feb 28, 2025

ajakovljevicTT requested review from kmitrovicTT and mrakitaTT as code owners February 28, 2025 15:24

kmitrovicTT reviewed Mar 3, 2025

View reviewed changes

ajakovljevicTT force-pushed the ajakovljevic/adding_input_sharding_collection branch 3 times, most recently from 260fb4d to 3f348ad Compare March 4, 2025 09:13

kmitrovicTT reviewed Mar 5, 2025

View reviewed changes

kmitrovicTT approved these changes Mar 6, 2025

View reviewed changes

Adding input sharding support

b5925e5

ajakovljevicTT force-pushed the ajakovljevic/adding_input_sharding_collection branch from 3f348ad to b5925e5 Compare March 6, 2025 16:00

mrakitaTT reviewed Mar 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adding input sharding collection in ModuleBuilder #289

Adding input sharding collection in ModuleBuilder #289

ajakovljevicTT commented Feb 28, 2025 •

edited

Loading

github-actions bot commented Feb 28, 2025 •

edited

Loading

codecov-commenter commented Mar 4, 2025 •

edited

Loading

kmitrovicTT Mar 5, 2025

ajakovljevicTT Mar 5, 2025 •

edited

Loading

kmitrovicTT Mar 6, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025 •

edited

Loading

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

mrakitaTT Mar 7, 2025

Adding input sharding collection in ModuleBuilder #289

Are you sure you want to change the base?

Adding input sharding collection in ModuleBuilder #289

Conversation

ajakovljevicTT commented Feb 28, 2025 • edited Loading

github-actions bot commented Feb 28, 2025 • edited Loading

codecov-commenter commented Mar 4, 2025 • edited Loading

Codecov Report

Choose a reason for hiding this comment

ajakovljevicTT Mar 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mrakitaTT Mar 7, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ajakovljevicTT commented Feb 28, 2025 •

edited

Loading

github-actions bot commented Feb 28, 2025 •

edited

Loading

codecov-commenter commented Mar 4, 2025 •

edited

Loading

ajakovljevicTT Mar 5, 2025 •

edited

Loading

mrakitaTT Mar 7, 2025 •

edited

Loading