feat: scatter op batch interface (#2072)

Author: avik-pal

src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp

@@ -467,10 +467,10 @@ class AutoDiffWhileFwd
         auto idx = arg.getArgNumber();
         if (resultPositionsToShadow.count(idx)) {
           if (gutils->isConstantValue(arg)) {
-            nb->insertArgument(
-                curidx,
-                cast<AutoDiffTypeInterface>(arg.getType()).getShadowType(gutils->width),
-                op.getLoc());
+            nb->insertArgument(curidx,
+                               cast<AutoDiffTypeInterface>(arg.getType())
+                                   .getShadowType(gutils->width),
+                               op.getLoc());
           }
           curidx++;
         }
@@ -4124,6 +4124,81 @@ struct SHLOConvolutionOpBatchInterface
   }
 };
 
+struct SHLOScatterOpBatchInterface
+    : public BatchOpInterface::ExternalModel<SHLOScatterOpBatchInterface,
+                                             stablehlo::ScatterOp> {
+  mlir::LogicalResult createBatch(Operation *src, OpBuilder &builder,
+                                  IRMapping &mapper,
+                                  ArrayRef<int64_t> batchSizes) const {
+    auto op = cast<stablehlo::ScatterOp>(src);
+
+    SmallVector<Value> newInputs;
+    newInputs.reserve(op.getInputs().size());
+    for (auto input : op.getInputs()) {
+      newInputs.push_back(mapper.lookup(input));
+    }
+
+    auto newScatterIndices = mapper.lookup(op.getScatterIndices());
+
+    SmallVector<Value> newUpdates;
+    newUpdates.reserve(op.getUpdates().size());
+    for (auto update : op.getUpdates()) {
+      newUpdates.push_back(mapper.lookup(update));
+    }
+
+    auto dimNumbers = op.getScatterDimensionNumbers();
+    int64_t nBatch = batchSizes.size();
+
+    SmallVector<int64_t> newUpdateWindowDims;
+    for (auto dim : dimNumbers.getUpdateWindowDims()) {
+      newUpdateWindowDims.push_back(dim + nBatch);
+    }
+
+    SmallVector<int64_t> newInsertedWindowDims;
+    for (auto dim : dimNumbers.getInsertedWindowDims()) {
+      newInsertedWindowDims.push_back(dim + nBatch);
+    }
+
+    SmallVector<int64_t> newInputBatchingDims, newScatterIndicesBatchingDims;
+    for (int64_t i = 0; i < nBatch; ++i) {
+      newInputBatchingDims.push_back(i);
+      newScatterIndicesBatchingDims.push_back(i);
+    }
+    for (auto dim : dimNumbers.getInputBatchingDims()) {
+      newInputBatchingDims.push_back(dim + nBatch);
+    }
+    for (auto dim : dimNumbers.getScatterIndicesBatchingDims()) {
+      newScatterIndicesBatchingDims.push_back(dim + nBatch);
+    }
+
+    SmallVector<int64_t> newScatterDimsToOperandDims;
+    for (auto dim : dimNumbers.getScatterDimsToOperandDims()) {
+      newScatterDimsToOperandDims.push_back(dim + nBatch);
+    }
+
+    auto newIndexVectorDim = dimNumbers.getIndexVectorDim() + nBatch;
+
+    auto newDimNumbers = stablehlo::ScatterDimensionNumbersAttr::get(
+        builder.getContext(), newUpdateWindowDims, newInsertedWindowDims,
+        newInputBatchingDims, newScatterIndicesBatchingDims,
+        newScatterDimsToOperandDims, newIndexVectorDim);
+
+    auto newScatterOp = stablehlo::ScatterOp::create(
+        builder, op.getLoc(), newInputs, newScatterIndices, newUpdates,
+        newDimNumbers, op.getIndicesAreSortedAttr(), op.getUniqueIndicesAttr());
+
+    IRMapping regionMapper;
+    op.getUpdateComputation().cloneInto(&newScatterOp.getUpdateComputation(),
+                                        regionMapper);
+
+    for (int i = 0; i < op.getNumResults(); ++i) {
+      mapper.map(op.getResult(i), newScatterOp.getResult(i));
+    }
+
+    return success();
+  }
+};
+
 struct StablehloAddSimplifyMathInterface
     : public MathSimplifyInterface::ExternalModel<
           StablehloAddSimplifyMathInterface, stablehlo::AddOp> {
@@ -4239,8 +4314,7 @@ void mlir::enzyme::registerStableHLODialectAutoDiffInterface(
     ConvolutionOp::attachInterface<SHLOConvolutionOpBatchInterface>(*context);
     PadOp::attachInterface<SHLOPadOpBatchInterface>(*context);
 
-    ScatterOp::attachInterface<SHLOGenericBatchOpInterface<ScatterOp>>(
-        *context); // TODO: simpler version with newly named dims
+    ScatterOp::attachInterface<SHLOScatterOpBatchInterface>(*context);
 
     AddOp::attachInterface<StablehloAddSimplifyMathInterface>(*context);
     SubtractOp::attachInterface<StablehloSubSimplifyMathInterface>(*context);

src/enzyme_ad/jax/Passes/AutoBatching.cpp

@@ -938,9 +938,8 @@ LogicalResult GreedyWhileLoopBatchFission::matchAndRewriteImpl(
         .Case<stablehlo::ReshapeOp, stablehlo::SliceOp, stablehlo::ReturnOp,
               // avoid ops that use SHLOGenericBatchOpInterface since that
               // lowers to loop
-              stablehlo::ScatterOp, stablehlo::IfOp, stablehlo::CaseOp,
-              stablehlo::WhileOp, stablehlo::CustomCallOp>(
-            [](auto op) { return true; })
+              stablehlo::IfOp, stablehlo::CaseOp, stablehlo::WhileOp,
+              stablehlo::CustomCallOp>([](auto op) { return true; })
         .Case<stablehlo::BroadcastInDimOp, stablehlo::TransposeOp>(
             [](auto op) { return stablehlo::OpIsReshapeLike(op); })
         .Default([](auto op) { return false; });
@@ -2327,6 +2326,7 @@ void populateAutoBatchingPassPatterns(RewritePatternSet &patterns,
              SliceToBatch<stablehlo::GetDimensionSizeOp>,
              SliceToBatch<stablehlo::ReverseOp>,
              SliceToBatch<stablehlo::ConvolutionOp>,
+             SliceToBatch<stablehlo::ScatterOp>,
              SliceToBatchWithReshapeLikeCheck<stablehlo::BroadcastInDimOp>,
              SliceToBatchWithReshapeLikeCheck<stablehlo::TransposeOp>,
              SliceToBatchElementwise>(ctx);
@@ -2338,8 +2338,7 @@ void populateAutoBatchingPassPatterns(RewritePatternSet &patterns,
                  ConcatInsertDimToBatch<stablehlo::IotaOp>,
                  ConcatInsertDimToBatchReduceLike<stablehlo::ReduceOp>,
                  ConcatInsertDimToBatchReduceLike<stablehlo::ReduceWindowOp>,
-                 // ConcatInsertDimToBatch<stablehlo::ScatterOp>, after batch
-                 // op interface is implemented
+                 ConcatInsertDimToBatch<stablehlo::ScatterOp>,
                  ConcatInsertDimToBatch<stablehlo::SortOp>,
                  ConcatInsertDimToBatch<stablehlo::ConcatenateOp>,
                  ConcatInsertDimToBatch<stablehlo::GetDimensionSizeOp>,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -2744,6 +2744,11 @@ def ApplyConcatInsertDimConvolutionPatterns : EnzymeHLOPatternOp<
   let patterns = ["ConcatInsertDimToBatch<stablehlo::ConvolutionOp>"];
 }
 
+def ApplyConcatInsertDimScatterPatterns : EnzymeHLOPatternOp<
+    "concat_insert_dim_scatter"> {
+  let patterns = ["ConcatInsertDimToBatch<stablehlo::ScatterOp>"];
+}
+
 def ApplyConcatInsertDimElementwisePatterns : EnzymeHLOPatternOp<
     "concat_insert_dim_elementwise"> {
   let patterns = ["ConcatInsertDimElementwiseToBatch"];
@@ -2769,6 +2774,10 @@ def ApplySortSliceToBatchPatterns : EnzymeHLOPatternOp<
     "sort_slice_to_batch"> {
   let patterns = ["SliceToBatch<stablehlo::SortOp>"];
 }
+def ApplyScatterSliceToBatchPatterns : EnzymeHLOPatternOp<
+    "scatter_slice_to_batch"> {
+  let patterns = ["SliceToBatch<stablehlo::ScatterOp>"];
+}
 def ApplyTransposeSliceToBatchPatterns : EnzymeHLOPatternOp<
     "transpose_slice_to_batch"> {
   let patterns = ["SliceToBatchWithReshapeLikeCheck<stablehlo::TransposeOp>"];

src/enzyme_ad/jax/primitives.py

@@ -395,7 +395,7 @@ def optimization_passes(
         "sub_const_prop",
         "xor_const_prop",
         # other constant propagations
-        "const_prop_through_barrier<16>",
+        # "const_prop_through_barrier<16>",
         f"concat_const_prop<1>({max_constant_threshold})",
         f"dynamic_update_slice_const_prop({max_constant_threshold})",
         "clamp_const_prop",
@@ -418,6 +418,7 @@ def optimization_passes(
             "reducewindow_slice_to_batch",
             "elementwise_slice_to_batch",
             "convolution_slice_to_batch",
+            "scatter_slice_to_batch",
         ]
 
     if enable_concat_to_batch_passes:
@@ -430,6 +431,7 @@ def optimization_passes(
             "concat_insert_dim_reduce_window",
             "concat_insert_dim_elementwise",
             "concat_insert_dim_convolution",
+            "concat_insert_dim_scatter",
         ]
 
     if enable_reduce_slice_fusion_passes:

test/lit_tests/batchtests/scatter.mlir

@@ -0,0 +1,45 @@
+// RUN: enzymexlamlir-opt %s --enzyme-batch | FileCheck %s
+
+module {
+  func.func private @unbatched_scatter(%arg0: tensor<8xf32>, %arg1: tensor<2x1xi32>, %arg2: tensor<2xf32>) -> tensor<8xf32> {
+    %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) <{
+      indices_are_sorted = false,
+      scatter_dimension_numbers = #stablehlo.scatter<
+        update_window_dims = [],
+        inserted_window_dims = [0],
+        scatter_dims_to_operand_dims = [0],
+        index_vector_dim = 1
+      >,
+      unique_indices = false
+    }> ({
+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
+      %1 = stablehlo.add %arg3, %arg4 : tensor<f32>
+      stablehlo.return %1 : tensor<f32>
+    }) : (tensor<8xf32>, tensor<2x1xi32>, tensor<2xf32>) -> tensor<8xf32>
+    return %0 : tensor<8xf32>
+  }
+
+  func.func @main(%arg0: tensor<4x8xf32>, %arg1: tensor<4x2x1xi32>, %arg2: tensor<4x2xf32>) -> tensor<4x8xf32> {
+    %0 = enzyme.batch @unbatched_scatter(%arg0, %arg1, %arg2) {batch_shape = array<i64: 4>} : (tensor<4x8xf32>, tensor<4x2x1xi32>, tensor<4x2xf32>) -> tensor<4x8xf32>
+    return %0 : tensor<4x8xf32>
+  }
+}
+
+// CHECK: func.func private @batched_unbatched_scatter(%arg0: tensor<4x8xf32>, %arg1: tensor<4x2x1xi32>, %arg2: tensor<4x2xf32>) -> tensor<4x8xf32> {
+// CHECK-NEXT:   %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) <{
+// CHECK-SAME:     indices_are_sorted = false,
+// CHECK-SAME:     scatter_dimension_numbers = #stablehlo.scatter<
+// CHECK-SAME:       inserted_window_dims = [1],
+// CHECK-SAME:       input_batching_dims = [0],
+// CHECK-SAME:       scatter_indices_batching_dims = [0],
+// CHECK-SAME:       scatter_dims_to_operand_dims = [1],
+// CHECK-SAME:       index_vector_dim = 2
+// CHECK-SAME:     >,
+// CHECK-SAME:     unique_indices = false
+// CHECK-SAME:   }> ({
+// CHECK-NEXT:   ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
+// CHECK-NEXT:     %1 = stablehlo.add %arg3, %arg4 : tensor<f32>
+// CHECK-NEXT:     stablehlo.return %1 : tensor<f32>
+// CHECK-NEXT:   }) : (tensor<4x8xf32>, tensor<4x2x1xi32>, tensor<4x2xf32>) -> tensor<4x8xf32>
+// CHECK-NEXT:   return %0 : tensor<4x8xf32>
+// CHECK-NEXT: }