Shardy: don't replicate update if constant input indices to DUS if enzyme comm opts on

PiperOrigin-RevId: 900259511

Author: wsmoses

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

@@ -834,6 +834,12 @@ OpShardingRuleAttr createOpShardingRule(Operation* op,
             SmallVector<int64_t> operandDims(
                 dynamicUpdateSlice->getNumOperands(), kNullDim);
             OpShardingRuleBuilder builder(dynamicUpdateSlice);
+
+            bool allIndicesConstant = llvm::all_of(dynamicUpdateSlice.getStartIndices(), [](Value v) {
+              Operation* defOp = v.getDefiningOp();
+              return defOp && (isa<sdy::ConstantOp>(defOp) || defOp->hasTrait<OpTrait::ConstantLike>());
+            });
+
             for (auto [dim, dimSizes] :
                  llvm::enumerate(llvm::zip_equal(operandShape, updateShape))) {
               auto [operandDimSize, updateDimSize] = dimSizes;
@@ -846,15 +852,20 @@ OpShardingRuleAttr createOpShardingRule(Operation* op,
                 //
                 // Thus, we add a factor for the operand/result slicing
                 // dimension with kPassThrough type. We also add a unique factor
-                // for the update with kNeedReplication type.
+                // for the update with kNeedReplication type (unless all indices
+                // are constant and we can rely on Enzyme comms opt).
                 operandDims[0] = dim;
                 operandDims[1] = kNullDim;
                 builder.addFactor(operandDims, dim, operandDimSize);
 
                 operandDims[0] = kNullDim;
                 operandDims[1] = dim;
-                builder.addFactor(operandDims, kNullDim, updateDimSize,
-                                  FactorType::kNeedReplication);
+                if (!allIndicesConstant) {
+                  builder.addFactor(operandDims, kNullDim, updateDimSize,
+                                    FactorType::kNeedReplication);
+                } else {
+                  builder.addFactor(operandDims, kNullDim, updateDimSize);
+                }
               }
             }
             return builder.build();

shardy/dialect/sdy/transforms/propagation/test/op_sharding_rule_registry.mlir

@@ -585,6 +585,15 @@ func.func @dynamic_update_slice(%arg0: tensor<32x4x8xf32>, %arg1: tensor<32x1x2x
   return %0 : tensor<32x4x8xf32>
 }
 
+// CHECK-LABEL: func @dynamic_update_slice_constant_indices
+func.func @dynamic_update_slice_constant_indices(%arg0: tensor<32x4x8xf32>, %arg1: tensor<32x1x2xf32>) -> tensor<32x4x8xf32> {
+  %0 = stablehlo.constant dense<0> : tensor<i32>
+  %1 = sdy.constant dense<0> : tensor<i32>
+  // CHECK: sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, l], [i, k, m], [], [], [])->([i, j, l]) {i=32, j=4, k=1, l=8, m=2}>
+  %2 = stablehlo.dynamic_update_slice %arg0, %arg1, %0, %0, %1 : (tensor<32x4x8xf32>, tensor<32x1x2xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<32x4x8xf32>
+  return %2 : tensor<32x4x8xf32>
+}
+
 // CHECK-LABEL: func @fft
 func.func @fft(%arg0: tensor<8x32x64xcomplex<f32>>) -> tensor<8x32x64xcomplex<f32>> {
   // CHECK: sdy.sharding_rule = #sdy.op_sharding_rule<([i, j, k])->([i, j, k]) {i=8, j=32, k=64} need_replication={j, k}>