[SDY] Add equi-sharding op ShardingGroupOp parsing and printing.

go/sdy-equi-sharding

PiperOrigin-RevId: 651172814

Author: Google-ML-Automation

shardy/dialect/sdy/ir/enums.td

@@ -38,4 +38,20 @@ def Sdy_PropagationDirection :
   let cppNamespace = Sdy_Dialect.cppNamespace;
 }
 
+// A sharding group type specifies the type of equi-sharding behavior applied
+// to a sharding group. This can be either
+
+// *   `AS`: all tensors in the group are strictly enforced to have the same
+//       sharding
+// *   `LIKE`: tensors in the group are encouraged to share shardings post
+//       propagation but the sameness of their shardings is not gauranteed.
+def Sdy_ShardingGroupType :
+    I32EnumAttr<"ShardingGroupType",
+        "sharding group type enum", [
+        I32EnumAttrCase<"AS", 0>,
+        I32EnumAttrCase<"LIKE", 1>]> {
+  let genSpecializedAttr = 1;
+  let cppNamespace = Sdy_Dialect.cppNamespace;
+}
+
 #endif  // SDY_ENUMS

shardy/dialect/sdy/ir/ops.td

@@ -187,6 +187,34 @@ def Sdy_ManualComputationOp : Sdy_Op<"manual_computation",
   }];
 }
 
+def Sdy_ShardingGroupOp : Sdy_Op<"sharding_group",
+  // The side-effects of this op are idempotent but since the op doesn't return
+  // the input tensor we don't mark it with the Idempotent op trait.
+  [NativeOpTrait<"ZeroResults">]>{
+  let summary = "Sharding group operation";
+  let description = [{
+    This op provides an interface to assign tensors to sharding groups
+    allowing users to have greater control over the similarity of shardings in
+    such groups. This operation takes arguments like group ID and sharding type
+    (LIKE or AS) and returns no result, but instead modifies the internal shard
+    group representation to guide the propagation process to meet the
+    "shard type" constraint specified.
+
+    Reference: go/sdy-equi-shard
+  }];
+
+  let arguments = (ins
+    AnyRankedTensor:$input,
+    I64Attr:$group_id,
+    Sdy_ShardingGroupType:$group_type,
+    OptionalAttr<Sdy_TensorSharding>:$sharding);
+
+  // Dangling op has no results.
+  let results = (outs);
+
+  let assemblyFormat = "$input `group_id````=```$group_id `type````=```$group_type (`sharding````=``` $sharding^)? attr-dict `:` type($input)";
+}
+
 def Sdy_ConstantOp : Sdy_Op<"constant",
   [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Constant operation";

shardy/dialect/sdy/ir/test/sharding_group_parse_print.mlir

@@ -0,0 +1,24 @@
+// RUN: sdy_opt %s 2>&1 | FileCheck %s
+
+sdy.mesh @mesh = <"a"=3>
+
+// CHECK-LABEL: func @add_to_shard_like
+func.func @add_to_shard_like(%arg0: tensor<8xf32>) {
+  // sdy.sharding_group %arg0 group_id=12 type=LIKE : tensor<8xf32>
+  sdy.sharding_group %arg0 group_id=12 type=LIKE : tensor<8xf32>
+  func.return
+}
+
+// CHECK-LABEL: func @add_to_shard_as
+func.func @add_to_shard_as(%arg0: tensor<8xf32>) {
+  // CHECK sdy.sharding_group %arg0 group_id=21 type=AS  : tensor<8xf32>
+  sdy.sharding_group %arg0 group_id=21 type=AS  : tensor<8xf32>
+  func.return
+}
+
+// CHECK-LABEL: func @add_to_shard_as_with_sharding
+func.func @add_to_shard_as_with_sharding(%arg0: tensor<8xf32>) {
+  // CHECK sdy.sharding_group %arg0 group_id=21 type=AS sharding=<@mesh, [{"a"}]> : tensor<8xf32>
+  sdy.sharding_group %arg0 group_id=21 type=AS sharding=<@mesh, [{"a"}]> : tensor<8xf32>
+  func.return
+}