[mpmd] Support replicated and unreduced axes in GenerateSdyMeshesFromTopologyPass

PiperOrigin-RevId: 810029732

Author: petebu

shardy/dialect/mpmd/transforms/import/generate_sdy_meshes_from_topology_pass.cc

@@ -111,10 +111,9 @@ class GenerateSdyMeshesFromTopologyPass
             sharding.getUnreducedAxes());
       }
       StringRef mesh_name;
-      SmallVector<sdy::DimensionShardingAttr> dim_shardings;
-      for (auto dim_sharding : sharding.getDimShardings()) {
-        SmallVector<sdy::AxisRefAttr> axes;
-        for (sdy::AxisRefAttr axis : dim_sharding.getAxes()) {
+      auto rename_axes = [&mesh_name](ArrayRef<sdy::AxisRefAttr> axes) {
+        SmallVector<sdy::AxisRefAttr> new_axes;
+        for (sdy::AxisRefAttr axis : axes) {
           auto [prefix, axis_name] = axis.getName().split(kMeshAxisSeparator);
           SDY_CHECK(!axis_name.empty())
               << "Axis name does not contain '" << kMeshAxisSeparator << "'";
@@ -124,19 +123,23 @@ class GenerateSdyMeshesFromTopologyPass
                 << prefix.str();
           }
           mesh_name = prefix;
-          axes.push_back(sdy::AxisRefAttr::get(
-              module_op.getContext(), axis_name, axis.getSubAxisInfo()));
+          new_axes.push_back(sdy::AxisRefAttr::get(axis.getContext(), axis_name,
+                                                   axis.getSubAxisInfo()));
         }
+        return new_axes;
+      };
+      SmallVector<sdy::DimensionShardingAttr> dim_shardings;
+      for (auto dim_sharding : sharding.getDimShardings()) {
         dim_shardings.push_back(sdy::DimensionShardingAttr::get(
-            module_op.getContext(), axes, dim_sharding.getIsClosed(),
-            dim_sharding.getPriority()));
+            module_op.getContext(), rename_axes(dim_sharding.getAxes()),
+            dim_sharding.getIsClosed(), dim_sharding.getPriority()));
       }
       SDY_CHECK(!llvm::is_contained(old_meshes, mesh_name))
           << "Invalid mesh name: " << mesh_name.str();
-      // TODO(b/440336690): Add support for replicated axes and unreduced axes.
       return sdy::TensorShardingAttr::get(
           sharding.getContext(), mesh_name, dim_shardings,
-          sharding.getReplicatedAxes(), sharding.getUnreducedAxes());
+          rename_axes(sharding.getReplicatedAxes()),
+          rename_axes(sharding.getUnreducedAxes()));
     });
 
     for (StringRef mesh_name : old_meshes) {

shardy/dialect/mpmd/transforms/import/test/generate_sdy_meshes_from_topology.mlir

@@ -108,3 +108,63 @@ module @fully_replicated_tensor {
     return %0 : tensor<16xf32>
   }
 }
+
+// -----
+
+// CHECK-LABEL: module @replicated_axes
+module @replicated_axes {
+  // CHECK-DAG: sdy.mesh @tpu = <["x"=2, "y"=4]>
+  // CHECK-DAG: sdy.mesh @cpu = <["z"=8]>
+  // CHECK-NOT: sdy.mesh @mesh
+  // CHECK-NOT: sdy.mesh @mesh_0
+  sdy.mesh @mesh = <["tpu_x"=8, "tpu_y"=8]>
+  sdy.mesh @mesh_0 = <["cpu_z"=8]>
+
+  // CHECK: %arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@tpu, [{"x"}], replicated={"y"}>}
+  // CHECK: %arg1: tensor<16xf32> {sdy.sharding = #sdy.sharding<@cpu, [{"z":(1)2}]>}
+  func.func @main(
+    %arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tpu_x"}], replicated={"tpu_y"}>},
+    %arg1: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"cpu_z":(1)2}]>})
+      -> (tensor<16xf32>) attributes {
+      topology = #mpmd.topology<<"tpu" : <["x"=2, "y"=4]>>, <"cpu" : <["z"=8]>>>} {
+    %0 = mpmd.named_computation<"stage1"> (%arg0, %arg0) (%arg3: tensor<16xf32>, %arg4: tensor<16xf32>) {
+      %2 = stablehlo.add %arg4, %arg3 : tensor<16xf32>
+      mpmd.return %2 : tensor<16xf32>
+    } : (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>
+    %1 = mpmd.named_computation<"stage2"> (%arg1, %0) (%arg3: tensor<16xf32>, %arg4: tensor<16xf32>) {
+      %2 = stablehlo.add %arg4, %arg3 : tensor<16xf32>
+      mpmd.return %2 : tensor<16xf32>
+    } : (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>
+    return %1 : tensor<16xf32>
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @unreduced_axes
+module @unreduced_axes {
+  // CHECK-DAG: sdy.mesh @tpu = <["x"=2, "y"=4]>
+  // CHECK-DAG: sdy.mesh @cpu = <["z"=8]>
+  // CHECK-NOT: sdy.mesh @mesh
+  // CHECK-NOT: sdy.mesh @mesh_0
+  sdy.mesh @mesh = <["tpu_x"=8, "tpu_y"=8]>
+  sdy.mesh @mesh_0 = <["cpu_z"=8]>
+
+  // CHECK: %arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@tpu, [{"x"}], unreduced={"y"}>}
+  // CHECK: %arg1: tensor<16xf32> {sdy.sharding = #sdy.sharding<@cpu, [{"z":(1)2}]>}
+  func.func @main(
+    %arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tpu_x"}], unreduced={"tpu_y"}>},
+    %arg1: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"cpu_z":(1)2}]>})
+      -> (tensor<16xf32>) attributes {
+      topology = #mpmd.topology<<"tpu" : <["x"=2, "y"=4]>>, <"cpu" : <["z"=8]>>>} {
+    %0 = mpmd.named_computation<"stage1"> (%arg0, %arg0) (%arg3: tensor<16xf32>, %arg4: tensor<16xf32>) {
+      %2 = stablehlo.add %arg4, %arg3 : tensor<16xf32>
+      mpmd.return %2 : tensor<16xf32>
+    } : (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>
+    %1 = mpmd.named_computation<"stage2"> (%arg1, %0) (%arg3: tensor<16xf32>, %arg4: tensor<16xf32>) {
+      %2 = stablehlo.add %arg4, %arg3 : tensor<16xf32>
+      mpmd.return %2 : tensor<16xf32>
+    } : (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>
+    return %1 : tensor<16xf32>
+  }
+}