#sdy. Add two special factor types (reduction, and need_replication) in OpShardingRule.

The sharding rule with the special factor types looks like
```
{sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} reduction={j} need_replication={i, l}, custom>}
```

The indices of special factors are sorted, unique and cannot overlap. The index should be vaild `0 <= i < num_factors`.

PiperOrigin-RevId: 707234719

Author: ZixuanJiang

shardy/dialect/sdy/ir/attrs.td

@@ -733,6 +733,11 @@ def Sdy_OpShardingRule : AttrDef<Sdy_Dialect, "OpShardingRule"> {
     this is mainly for completeness as many ops such as pointwise ops have size
     one dimensions that correspond across operands and results.
 
+    `reduction_factors` contains the indices of factors requiring reduction,
+    such as the contracting dimensions in a dot operation.
+    `need_replication_factors` contains the indices of factors requiring full
+    replication, such as the sorted dimension in a sort operation.
+
     `is_custom_rule` describes whether this is a rule defined by a user for a
     `stablehlo.custom_call` op. The partitioner doesn't know how to partition
     these ops, so a user must tell it how. When it is a custom rule, then the
@@ -744,6 +749,8 @@ def Sdy_OpShardingRule : AttrDef<Sdy_Dialect, "OpShardingRule"> {
       OptionalArrayRefParameter<"int64_t">:$factor_sizes,
       OptionalArrayRefParameter<"TensorMappingAttr">:$operand_mappings,
       OptionalArrayRefParameter<"TensorMappingAttr">:$result_mappings,
+      OptionalArrayRefParameter<"int64_t">:$reduction_factors,
+      OptionalArrayRefParameter<"int64_t">:$need_replication_factors,
       DefaultValuedParameter<"bool", "false">:$is_custom_rule
   );
 
@@ -752,16 +759,21 @@ def Sdy_OpShardingRule : AttrDef<Sdy_Dialect, "OpShardingRule"> {
     `(`$operand_mappings`)`
     `` `->` ``
     `(`$result_mappings`)` ``
-    custom<FactorSizes>($factor_sizes)
-    ``custom<IsCustomRule>($is_custom_rule)
+    custom<FactorSizes>($factor_sizes) ``
+    custom<ReductionFactors>($reduction_factors) ``
+    custom<NeedReplicationFactors>($need_replication_factors) ``
+    custom<IsCustomRule>($is_custom_rule)
     `>`
   }];
 
   let builders = [
     AttrBuilder<(ins "ArrayRef<int64_t>":$factor_sizes,
                      "ArrayRef<TensorMappingAttr>":$operand_mappings,
-                     "ArrayRef<TensorMappingAttr>":$result_mappings), [{
+                     "ArrayRef<TensorMappingAttr>":$result_mappings,
+                     "ArrayRef<int64_t>":$reduction_factors,
+                     "ArrayRef<int64_t>":$need_replication_factors), [{
       return $_get($_ctxt, factor_sizes, operand_mappings, result_mappings,
+                   reduction_factors, need_replication_factors,
                    /*is_custom_rule=*/false);
     }]>
   ];

shardy/dialect/sdy/ir/parsers.cc

@@ -289,6 +289,50 @@ ParseResult parseFactorSizes(AsmParser& parser,
   return success();
 }
 
+namespace {
+
+// Parses factor sizes. In a OpShardingRule, you could have `, type={k, i}`.
+// `k` is index 2, while `i` is index 0. Thus factors would be set to [2, 0].
+ParseResult parseFactorsWithType(AsmParser& parser,
+                                 SmallVector<int64_t>& factors,
+                                 StringRef type) {
+  auto parseElementFn = [&]() -> ParseResult {
+    StringRef factorSymbol;
+    if (parser.parseKeyword(&factorSymbol)) {
+      return failure();
+    }
+    FailureOr<int64_t> factorIndex =
+        parseFactorSymbolIndex(parser, factorSymbol);
+    if (failed(factorIndex)) {
+      return failure();
+    }
+    factors.push_back(*factorIndex);
+    return success();
+  };
+
+  if (!parser.parseOptionalKeyword(type)) {
+    if (parser.parseEqual()) {
+      return failure();
+    }
+    return parser.parseCommaSeparatedList(AsmParser::Delimiter::OptionalBraces,
+                                          parseElementFn);
+  }
+  return success();
+}
+
+}  // namespace
+
+ParseResult parseReductionFactors(AsmParser& parser,
+                                  SmallVector<int64_t>& reductionFactors) {
+  return parseFactorsWithType(parser, reductionFactors, "reduction");
+}
+
+ParseResult parseNeedReplicationFactors(
+    AsmParser& parser, SmallVector<int64_t>& needReplicationFactors) {
+  return parseFactorsWithType(parser, needReplicationFactors,
+                              "need_replication");
+}
+
 ParseResult parseIsCustomRule(AsmParser& parser, bool& isCustomRule) {
   isCustomRule = false;
   if (!parser.parseOptionalComma()) {

shardy/dialect/sdy/ir/parsers.h

@@ -38,6 +38,16 @@ ParseResult parseMeshOrRef(AsmParser& parser, Attribute& meshOrRef);
 ParseResult parseFactorSizes(AsmParser& parser,
                              SmallVector<int64_t>& factorSizes);
 
+// Parses the reduction factors of an OpShardingRule. We expect to parse
+// `reduction={i, k}` into a vector [0, 2].
+ParseResult parseReductionFactors(AsmParser& parser,
+                                  SmallVector<int64_t>& reductionFactors);
+
+// Parses the factors needing replication of an OpShardingRule. We expect to
+// parse `need_replication={i, k}` into a vector [0, 2].
+ParseResult parseNeedReplicationFactors(
+    AsmParser& parser, SmallVector<int64_t>& needReplicationFactors);
+
 ParseResult parseIsCustomRule(AsmParser& parser, bool& isCustomRule);
 
 // Parses a single block region without the block id. This is an example of what

shardy/dialect/sdy/ir/printers.cc

@@ -74,6 +74,33 @@ void printFactorSizes(AsmPrinter& printer, ArrayRef<int64_t> factorSizes) {
   printer << "}";
 }
 
+namespace {
+
+void printFactorsWithType(AsmPrinter& printer, ArrayRef<int64_t> factors,
+                          StringRef type) {
+  if (factors.empty()) {
+    return;
+  }
+  printer << " " << type << "={";
+  llvm::interleaveComma(factors, printer, [&](int64_t factor) {
+    printer << factorSymbolString(factor);
+  });
+  printer << "}";
+}
+
+}  // namespace
+
+void printReductionFactors(AsmPrinter& printer,
+                           ArrayRef<int64_t> reductionFactors) {
+  return printFactorsWithType(printer, reductionFactors, "reduction");
+}
+
+void printNeedReplicationFactors(AsmPrinter& printer,
+                                 ArrayRef<int64_t> needReplicationFactors) {
+  return printFactorsWithType(printer, needReplicationFactors,
+                              "need_replication");
+}
+
 void printIsCustomRule(AsmPrinter& printer, bool isCustomRule) {
   if (isCustomRule) {
     printer << ", custom";

shardy/dialect/sdy/ir/printers.h

@@ -38,6 +38,16 @@ void printMeshOrRef(AsmPrinter& printer, Attribute meshOrRef);
 // printed as `{i=6, j=2, k=4}`.
 void printFactorSizes(AsmPrinter& printer, ArrayRef<int64_t> factorSizes);
 
+// Prints the reduction factors of an OpShardingRule. Given a vector [0, 2], we
+// print `reduction={i, k}`.
+void printReductionFactors(AsmPrinter& printer,
+                           ArrayRef<int64_t> reductionFactors);
+
+// Prints the factors needing replication of an OpShardingRule. Given a vector
+// [0, 2], we print `need_replication={i, k}`.
+void printNeedReplicationFactors(AsmPrinter& printer,
+                                 ArrayRef<int64_t> needReplicationFactors);
+
 void printIsCustomRule(AsmPrinter& printer, bool isCustomRule);
 
 // Prints a single block region without the block id, for example:

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

@@ -46,3 +46,10 @@ func.func @custom_call_custom_rule(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32
   %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=32}, custom>} : (tensor<16x32xf32>) -> tensor<16x32xf32>
   func.return %0: tensor<16x32xf32>
 }
+
+// CHECK-LABEL: func @reduction_and_need_replication_factors
+func.func @reduction_and_need_replication_factors(%arg0: tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32> {
+  // CHECK: {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} reduction={j} need_replication={i, l}, custom>}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} reduction={j} need_replication={i, l}, custom>} : (tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32>
+  func.return %0: tensor<2x5x7xf32>
+}

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

@@ -181,3 +181,35 @@ func.func @no_results(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> {
   stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->() {i=2, j=8}, custom>} : (tensor<2x8xf32>) -> ()
   func.return %arg0 : tensor<2x8xf32>
 }
+
+// -----
+
+func.func @equality_sign_after_reduction(%arg0: tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32> {
+  // expected-error@+1 {{expected '='}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} reduction: {j}>} : (tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32>
+  func.return %0: tensor<2x5x7xf32>
+}
+
+// -----
+
+func.func @reduce_is_an_unknown_keyword(%arg0: tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32> {
+  // expected-error@+1 {{expected '>'}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} reduce={j}>} : (tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32>
+  func.return %0: tensor<2x5x7xf32>
+}
+
+// -----
+
+func.func @reduction_should_be_before_need_replication(%arg0: tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32> {
+  // expected-error@+1 {{expected '>'}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} need_replication={i} reduction={j}>} : (tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32>
+  func.return %0: tensor<2x5x7xf32>
+}
+
+// -----
+
+func.func @invalid_dimension_symbol(%arg0: tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32> {
+  // expected-error@+1 {{expecting symbol from 'i' to 'z'. Received: 'a'}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k, l])->([i, k, l]) {i=2, j=3, k=5, l=7} reduction={a}>} : (tensor<2x3x5x7xf32>) -> tensor<2x5x7xf32>
+  func.return %0: tensor<2x5x7xf32>
+}

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

@@ -109,3 +109,35 @@ func.func @duplicate_factor_different_dim(%arg0: tensor<2x2xf32>) -> tensor<4xf3
   %0 = stablehlo.reshape %arg0 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, i])->([ij]) {i=2, j=2}>} : (tensor<2x2xf32>) -> tensor<4xf32>
   return %0 : tensor<4xf32>
 }
+
+// -----
+
+func.func @unsorted_special_factors(%arg0: tensor<2x4x8xf32>) -> tensor<2x8xf32> {
+  // expected-error@+1 {{indices of special factors must be sorted}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k])->([i, k]) {i=2, j=4, k=8} need_replication={k, i}>} : (tensor<2x4x8xf32>) -> tensor<2x8xf32>
+  func.return %0: tensor<2x8xf32>
+}
+
+// -----
+
+func.func @repeated_special_factors(%arg0: tensor<2x4x8xf32>) -> tensor<2x8xf32> {
+  // expected-error@+1 {{indices of special factors must be unique}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k])->([i, k]) {i=2, j=4, k=8} need_replication={i, i}>} : (tensor<2x4x8xf32>) -> tensor<2x8xf32>
+  func.return %0: tensor<2x8xf32>
+}
+
+// -----
+
+func.func @invalid_special_factor_index(%arg0: tensor<2x4x8xf32>) -> tensor<2x8xf32> {
+  // expected-error@+1 {{index must be less than 3, got: 17}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k])->([i, k]) {i=2, j=4, k=8} need_replication={z}>} : (tensor<2x4x8xf32>) -> tensor<2x8xf32>
+  func.return %0: tensor<2x8xf32>
+}
+
+// -----
+
+func.func @invalid_special_factor_index(%arg0: tensor<2x4x8xf32>) -> tensor<2x8xf32> {
+  // expected-error@+1 {{reduction and need_replication factors must be disjoint}}
+  %0 = stablehlo.custom_call @foo(%arg0) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k])->([i, k]) {i=2, j=4, k=8} reduction={j} need_replication={j}>} : (tensor<2x4x8xf32>) -> tensor<2x8xf32>
+  func.return %0: tensor<2x8xf32>
+}

shardy/dialect/sdy/ir/verifiers.cc

@@ -17,6 +17,7 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <functional>
+#include <iterator>
 #include <numeric>
 #include <optional>
 #include <utility>
@@ -542,6 +543,30 @@ LogicalResult verifyShardingRuleMapping(Operation* op, TypeRange types,
   return success();
 }
 
+LogicalResult verifyIndicesOfSpecialFactors(Operation* op, int64_t numFactors,
+                                            ArrayRef<int64_t> indices) {
+  if (indices.empty()) {
+    return success();
+  }
+
+  if (!llvm::is_sorted(indices)) {
+    return op->emitOpError("indices of special factors must be sorted");
+  }
+  if (std::adjacent_find(indices.begin(), indices.end()) != indices.end()) {
+    return op->emitOpError("indices of special factors must be unique");
+  }
+
+  if (indices.front() < 0) {
+    return op->emitOpError("index must be non-negative");
+  }
+  if (indices.back() >= numFactors) {
+    return op->emitOpError("index must be less than ")
+           << numFactors << ", got: " << indices.back();
+  }
+
+  return success();
+}
+
 // Verifies the following for an `OpShardingRuleAttr`:
 //
 // - If the rule is custom, the operation the rule is attached to is a
@@ -575,6 +600,29 @@ LogicalResult verifyOpShardingRuleAttr(OpShardingRuleAttr shardingRule,
            << " that isn't used in operand and result mappings";
   }
 
+  ArrayRef<int64_t> reductionFactors = shardingRule.getReductionFactors();
+  ArrayRef<int64_t> needReplicationFactors =
+      shardingRule.getNeedReplicationFactors();
+
+  if (failed(verifyIndicesOfSpecialFactors(op, shardingRule.getNumFactors(),
+                                           reductionFactors))) {
+    return failure();
+  }
+  if (failed(verifyIndicesOfSpecialFactors(op, shardingRule.getNumFactors(),
+                                           needReplicationFactors))) {
+    return failure();
+  }
+
+  SmallVector<int64_t> intersection;
+  std::set_intersection(reductionFactors.begin(), reductionFactors.end(),
+                        needReplicationFactors.begin(),
+                        needReplicationFactors.end(),
+                        std::back_inserter(intersection));
+  if (!intersection.empty()) {
+    return op->emitOpError(
+        "reduction and need_replication factors must be disjoint");
+  }
+
   return success();
 }
 

shardy/dialect/sdy/transforms/propagation/op_sharding_rule_builder.cc

@@ -115,7 +115,8 @@ OpShardingRuleAttr OpShardingRuleBuilder::build() {
       buildTensorMappingAttrList(resultMappings, factorSizes, context);
 
   auto result = OpShardingRuleAttr::get(
-      context, factorSizes, operandMappingAttrs, resultMappingAttrs);
+      context, factorSizes, operandMappingAttrs, resultMappingAttrs,
+      reductionFactors, needReplicationFactors);
 
   // Erase all added factors, to return the builder to its original state before
   // calling this method.