feat: abstraction of xla::OpSharding proto using wrapper class #9467

kvshbg-aws · 2025-07-10T00:29:51Z

This PR includes the changes related to abstracting xla::OpSharidng proto object into a torch_xla::OpSharding wrapper class.

This new class object will not have the requirements of xla::OpSharding (however, it will be an extension xla::OpSharding proto defined over here).
We have defined the wrapper class in torch/xla which will construct an xla::OpSharding object with additional fields such as global_device_ids/global_tile_assignment and will have forwarded/proxy functions to xla::OpSharding . These forwarded functions will help user still make use of the same xla::OpSharding APIs as they normally would. We can also define torch_xla specific functions in this wrapper class to further use the extra fields that were stored during the initialization of the OpSharding object. This approach also allows the flexibility of converting the torch_xla::OpSharding object back to xla::OpSharding while lowering into HLO, thus, giving user the flexibility to use the abstracted class (and other additional fields stored) anywhere in the code base as needed, this is particularly useful since the XLA's HLOs are 0th indexed, hence we need to use the normalized_device_ids (starting from index 0) when lowering the program into the HLO, whereas we can still use the denormalized/global_device_ids in other places such as inside pjrt client to set the device_assignment using the user specified device_ids.

Component diagram for reference -

Ref issue - #9390

…_assignment() is empty

…o_data

…ment

pgmoka

Left comments. Overall, the PR looks very good.

pgmoka · 2025-08-19T16:53:45Z

torch_xla/csrc/ir.h

+  const std::vector<std::shared_ptr<torch_xla::OpSharding>> GetShardings()
+      const {
+    return output_shardings_;
+  }


I would keep GetShardings consistend with SetSharding and have its implementation be in ir.cpp

pgmoka · 2025-08-19T17:05:43Z

torch_xla/csrc/lowering_context.cpp

+    std::vector<int64_t> denormalized_tile_assignment =
+        sharding->GetDenormalizedTileAssignment();
+    if (!denormalized_tile_assignment.empty()) {
+      denormalized_tile_assignments_.push_back(
+          sharding->GetDenormalizedTileAssignment());
+    }


NIT: Mixing declaration types here is a bit confusing to read. I would consider:

In line 247, using:

denormalized_tile_assignments_.push_back(denormalized_tile_assignment);

OR
2) Remove denormalized_tile_assignment, and just check !sharding->GetDenormalizedTileAssignment().empty()

pgmoka · 2025-08-19T17:47:56Z

torch_xla/csrc/xla_graph_executor.cpp

+  std::vector<std::vector<int64_t>> denormalized_tile_assignments;
+  for (const auto* node : po_data->post_order) {
+    const XlaNode* const casted = dynamic_cast<const XlaNode*>(node);
+    auto shardings = casted->GetShardings();
+    if (!shardings.empty()) {
+      for (auto sharding : shardings) {
+        std::vector<int64_t> denormalized_tile_assignment =
+            sharding->GetDenormalizedTileAssignment();
+        if (!denormalized_tile_assignment.empty()) {
+          denormalized_tile_assignments.push_back(
+              sharding->GetDenormalizedTileAssignment());
+        }
+      }
+    }
+  }


For posterity, a quick comment to outline what this is doing might be useful.

pgmoka · 2025-08-19T18:00:32Z

torch_xla/csrc/xla_sharding_util.cpp

+  if (sharding.tile_assignment_devices().empty() && !tile_assignment.empty()) {
+    // Convert the Python list tile_assignment to a flattened vector for
+    // denormalized assignment
+    xla::Array<int64_t> tile_array = TileListToArray(tile_assignment);
+    denormalized_tile_assignment.assign(tile_array.begin(), tile_array.end());
+  } else {
+    // Use the tile_assignment_devices from the XLA OpSharding object
+    denormalized_tile_assignment.assign(
+        sharding.tile_assignment_devices().begin(),
+        sharding.tile_assignment_devices().end());
+  }


NIT suggestion: we can remove if branching here by doing something like:

// Use the tile_assignment_devices from the XLA OpSharding object by default denormalized_tile_assignment.assign( sharding.tile_assignment_devices().begin(), sharding.tile_assignment_devices().end()); if (sharding.tile_assignment_devices().empty() && !tile_assignment.empty()) { // Convert the Python list tile_assignment to a flattened vector for // denormalized assignment xla::Array<int64_t> tile_array = TileListToArray(tile_assignment); denormalized_tile_assignment.assign(tile_array.begin(), tile_array.end()); }

pgmoka · 2025-08-19T21:17:53Z

torch_xla/csrc/xla_sharding_util.cpp

    }
  }
-  if (input_shardings.size() == 0) {
+  if (xla_input_shardings.size() == 0) {
    TF_VLOG(3) << "ReshardParamters... skip with empty input_shardings.";


NIT: "ReshardParamters... skip with empty xla_input_shardings."

pgmoka · 2025-08-19T21:25:54Z

torch_xla/csrc/xla_sharding_util.cpp

+        sharding_spec->sharding.GetDenormalizedTileAssignment();
+  }
+  for (const auto& sharding : xla_input_shardings) {
+    if (denormalized_tile_assignment.size() > 0) {


I think this generates an odd use case. I am not sure if it would ever get triggered, but if (*tensors)[0]->sharding_spec() does not exist, denormalized_tile_assignment never gets created, which might generate a error here. The way to ensure this is to move the for loop into the if (sharding_spec) check.

pgmoka · 2025-08-19T22:04:29Z

test/cpp/test_xla_sharding.cpp

+  auto xla_sharding = xla::HloSharding::Tile({
+                                                 {0, 1, 2, 3},
+                                                 {4, 5, 6, 7},
+                                             })
+                          .ToProto();
+  std::vector<int64_t> denormalized_tile_assignment = {0, 1, 2, 3, 4, 5, 6, 7};
+  torch_xla::OpSharding sharding(xla_sharding, denormalized_tile_assignment);
+  XLATensor::ShardingSpec tiled_2d(sharding, tensor_shape);
+  xla_sharding =
+      xla::HloSharding::Tile({{{0, 1}, {2, 3}, {4, 5}, {6, 7}}}).ToProto();
+  sharding = torch_xla::OpSharding(xla_sharding, denormalized_tile_assignment);
+  XLATensor::ShardingSpec tiled_3d(sharding, tensor_shape);
+  xla_sharding = xla::HloSharding::Replicate().ToProto();
+  sharding = torch_xla::OpSharding(xla_sharding, denormalized_tile_assignment);
+  XLATensor::ShardingSpec replicated(sharding, tensor_shape);


This is a bit confusing. Specifically, xla_sharding is used and redefined multiple times through the test. I would perhaps separate these into separate instances. It is less efficient, but I would favor readability for the test

pgmoka · 2025-08-19T22:11:20Z

torch_xla/csrc/runtime/ifrt_computation_client.h

+    IfrtComputation(
+        xla::XlaComputation computation, std::vector<std::string> devices,
+        std::shared_ptr<xla::ifrt::LoadedExecutable> executable,
+        std::optional<std::vector<int64_t>> denormalized_tile_assignment)
        : Computation(std::move(computation), std::move(devices)),
-          executable(std::move(executable)) {
-      output_shardings_ = this->executable->GetOutputShardings();
+          executable(std::move(executable)),
+          denormalized_tile_assignment_(std::move(
+              denormalized_tile_assignment.value_or(std::vector<int64_t>{}))) {
+      xla_output_shardings_ = this->executable->GetOutputShardings();
+      if (xla_output_shardings_.has_value()) {
+        output_shardings_ = std::vector<torch_xla::OpSharding>{};
+        output_shardings_->reserve(xla_output_shardings_.value().size());
+        for (const auto& sharding : xla_output_shardings_.value()) {
+          // convert each into torch_xla::OpSharding object
+          torch_xla::OpSharding torch_xla_op_sharding(
+              sharding, denormalized_tile_assignment_);
+          output_shardings_.value().push_back(torch_xla_op_sharding);
+        }
+      } else {
+        output_shardings_ = std::nullopt;
+      }


It would be good to move the constructor to ifrt_computation_client.cpp. If you can do here, great. If you don't want to do it, please create an issue and link a TODO from here for posterity.

pgmoka · 2025-08-19T22:17:08Z

torch_xla/csrc/runtime/ifrt_computation_client.h

+      if (xla_output_shardings_.has_value()) {
+        output_shardings_ = std::vector<torch_xla::OpSharding>{};
+        output_shardings_->reserve(xla_output_shardings_.value().size());
+        for (const auto& sharding : xla_output_shardings_.value()) {
+          // convert each into torch_xla::OpSharding object
+          torch_xla::OpSharding torch_xla_op_sharding(
+              sharding, denormalized_tile_assignment_);
+          output_shardings_.value().push_back(torch_xla_op_sharding);
+        }
+      } else {
+        output_shardings_ = std::nullopt;
+      }


We can simplify this if statement tree by setting output_shardings_ = std::nullopt; as default. Ex.:

output_shardings_ = std::nullopt; if (xla_output_shardings_.has_value()) { output_shardings_ = std::vector<torch_xla::OpSharding>{}; output_shardings_->reserve(xla_output_shardings_.value().size()); for (const auto& sharding : xla_output_shardings_.value()) { // convert each into torch_xla::OpSharding object torch_xla::OpSharding torch_xla_op_sharding( sharding, denormalized_tile_assignment_); output_shardings_.value().push_back(torch_xla_op_sharding); } }

pgmoka · 2025-08-19T22:21:50Z

torch_xla/csrc/runtime/pjrt_computation_client.h

+     * @param executable The compiled PJRT executable
+     * @param denormalized_tile_assignment Optional tile assignment for sharding
+     */
+    PjRtComputation(


Same notes as ifrt_computation_client.cpp

kvshbg-aws force-pushed the kvshbg-aws/local-spmd-abstraction branch from d0502ab to 7fc15ea Compare July 10, 2025 18:34

qihqi requested review from rpsilva-aws and pgmoka July 11, 2025 04:20

kvshbg-aws force-pushed the kvshbg-aws/local-spmd-abstraction branch 4 times, most recently from 7c4a3cd to 1d55ae9 Compare July 16, 2025 23:57

kvshbg-aws force-pushed the kvshbg-aws/local-spmd-abstraction branch 2 times, most recently from 1ddbb1b to 2756c1a Compare July 25, 2025 23:53

kvshbg-aws mentioned this pull request Jul 28, 2025

feat: infer device_ids and normalize tile assignment #9514

Open

kvshbg-aws added 7 commits July 30, 2025 04:13

feat: abstraction of xla::OpSharding proto using wrapper class

50faf0d

fix for failing ci/cd tests

1cfe785

small fix to remove unwanted var

940a902

bug fix: set value of denormalized_tile_assignment when sharding.tile…

9782a99

…_assignment() is empty

use tensors to get denormalized_tile_assignment directly instead of p…

1805c50

…o_data

use tensors instead of paramters_data to get denormalized_tile_assign…

2d1f60e

…ment

linter fix

4556faa

kvshbg-aws force-pushed the kvshbg-aws/local-spmd-abstraction branch from 2756c1a to 4556faa Compare July 30, 2025 04:14

fix: use lowering_cntxt to save/get denormalized_tile_assignment

4752558

pgmoka reviewed Aug 19, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: abstraction of xla::OpSharding proto using wrapper class #9467

feat: abstraction of xla::OpSharding proto using wrapper class #9467

Uh oh!

kvshbg-aws commented Jul 10, 2025

Uh oh!

pgmoka left a comment

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

pgmoka Aug 19, 2025

Uh oh!

Uh oh!

feat: abstraction of xla::OpSharding proto using wrapper class #9467

Are you sure you want to change the base?

feat: abstraction of xla::OpSharding proto using wrapper class #9467

Uh oh!

Conversation

kvshbg-aws commented Jul 10, 2025

Uh oh!

pgmoka left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!