[Preprocessing] Handle batching dims and dynamic-update-slice scatter ops (#24720)
See bug report https://github.com/iree-org/iree/issues/24719
JAX-emitted scatters use two forms that IREE's StableHLO scatter
preprocessing did not support, causing malformed shape errors during
input conversion:
1. Scatter batching dims were only passed through, so the downstream
canonicalizers collapsed batching dims together with the real scatter
loop dim and miscompiled. Added `ScatterBatchingDimsExpander` (benefit
2). This is ported with just a few style tweaks directly from upstream
StableHLO's `ScatterWithBatchingDimsExpander`. I verified that the
StableHLO pattern wasn't accessible to be called directly. The only way
to access it is to run the whole set of compatibility expanders which
seemed brittle since new expanders can be added over time, so just
copied it and its helper functions. I didn't spend time to fully
understand the code though since it's an already reviewed pattern used
to handle this exact case for backwards compatibility in StableHLO.
Basically, it lowers each batching dim to an explicit iota index column
so no batching dims remain for the existing patterns. Here's an example
```mlir
module @m {
func.func @s(%operand: tensor<8x33xi32>, %indices: tensor<8x32x1xi32>, %updates: tensor<8x32xi32>) -> tensor<8x33xi32> {
%0 = "stablehlo.scatter"(%operand, %indices, %updates) <{scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [1], input_batching_dims = [0], scatter_indices_batching_dims = [0], scatter_dims_to_operand_dims = [1], index_vector_dim = 2>}> ({
^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>):
%1 = stablehlo.minimum %arg2, %arg3 : tensor<i32>
stablehlo.return %1 : tensor<i32>
}) : (tensor<8x33xi32>, tensor<8x32x1xi32>, tensor<8x32xi32>) -> tensor<8x33xi32>
return %0 : tensor<8x33xi32>
}
}
```
gets preprocessed into
```mlir
module @m {
func.func @s(%arg0: tensor<8x33xi32>, %arg1: tensor<8x32x1xi32>, %arg2: tensor<8x32xi32>) -> tensor<8x33xi32> {
%0 = stablehlo.iota dim = 0 : tensor<8x32x1xi32>
%1 = stablehlo.concatenate %0, %arg1, dim = 2 : (tensor<8x32x1xi32>, tensor<8x32x1xi32>) -> tensor<8x32x2xi32>
%collapsed = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<8x32x2xi32> into tensor<256x2xi32>
%collapsed_0 = tensor.collapse_shape %arg2 [[0, 1]] : tensor<8x32xi32> into tensor<256xi32>
%2 = "stablehlo.scatter"(%arg0, %collapsed, %collapsed_0) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 1>, unique_indices = false}> ({
^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>):
%3 = stablehlo.minimum %arg3, %arg4 : tensor<i32>
stablehlo.return %3 : tensor<i32>
}) : (tensor<8x33xi32>, tensor<256x2xi32>, tensor<256xi32>) -> tensor<8x33xi32>
return %2 : tensor<8x33xi32>
}
}
```
2. Single-index, full-rank overwrite scatters are dynamic-update-slice
semantics and cannot be represented as `iree_linalg_ext.scatter`. Added
`ScatterToDynamicUpdateSlice` (benefit 3) to rewrite them into
`stablehlo.dynamic_update_slice`.
Here's an example:
```mlir
module @m {
func.func @s(%operand: tensor<8x7x24x32xf32>, %indices: tensor<1xi32>, %updates: tensor<8x4x24x32xf32>) -> tensor<8x7x24x32xf32> {
%0 = "stablehlo.scatter"(%operand, %indices, %updates) <{indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0, 1, 2, 3], scatter_dims_to_operand_dims = [1]>, unique_indices = true}> ({
^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
stablehlo.return %arg3 : tensor<f32>
}) : (tensor<8x7x24x32xf32>, tensor<1xi32>, tensor<8x4x24x32xf32>) -> tensor<8x7x24x32xf32>
return %0 : tensor<8x7x24x32xf32>
}
}
```
gets preprocessed into
```mlir
module @m {
func.func @s(%arg0: tensor<8x7x24x32xf32>, %arg1: tensor<1xi32>, %arg2: tensor<8x4x24x32xf32>) -> tensor<8x7x24x32xf32> {
%c = stablehlo.constant dense<0> : tensor<i32>
%0 = stablehlo.slice %arg1 [0:1] : (tensor<1xi32>) -> tensor<1xi32>
%1 = stablehlo.reshape %0 : (tensor<1xi32>) -> tensor<i32>
%2 = stablehlo.dynamic_update_slice %arg0, %arg2, %c, %1, %c, %c : (tensor<8x7x24x32xf32>, tensor<8x4x24x32xf32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<8x7x24x32xf32>
return %2 : tensor<8x7x24x32xf32>
}
}
```
Both patterns run before the existing scatter canonicalizers, because
the first expands out functionality not supported by later patterns, and
the second one removes the scatter entirely. Lit tests added for both.
---------
Signed-off-by: Paul Stark <paul.stark@cdprojektred.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>diff --git a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
index 93f5c4b..e152d9c 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
+++ b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
@@ -454,12 +454,293 @@
}
};
+// The following helpers and the ScatterBatchingDimsExpander pattern below are
+// ported from upstream StableHLO's `ScatterWithBatchingDimsExpander` and its
+// supporting helpers in
+// third_party/stablehlo/stablehlo/transforms/StablehloCompatibilityExpander.cpp
+// (mergeSortedDims / fitsInIntegralType / promoteTypeForSize /
+// getUpdatedIndicesAreSorted / createConcatIndices). They are copied rather
+// than reused because the upstream pattern lives in an anonymous namespace, and
+// its only public entry point (populateStablehloCompatibilityExpanderPatterns)
+// is a version-gated bundle that would also pull in unrelated compatibility
+// expanders.
+//
+// Merges two sorted lists of dimensions into a single sorted list.
+SmallVector<int64_t> mergeSortedDims(ArrayRef<int64_t> dims1,
+ ArrayRef<int64_t> dims2) {
+ SmallVector<int64_t> result;
+ result.reserve(dims1.size() + dims2.size());
+ std::merge(dims1.begin(), dims1.end(), dims2.begin(), dims2.end(),
+ std::back_inserter(result));
+ return result;
+}
+
+bool fitsInIntegralType(int64_t size, IntegerType type) {
+ if (type.isUnsigned()) {
+ return llvm::isUIntN(type.getWidth(), size);
+ }
+ return llvm::isIntN(type.getWidth(), size);
+}
+
+// If `type` is an integer type in which `size` doesn't fit, promote it to i32
+// or i64 (depending on `size`).
+Type promoteTypeForSize(Type type, int64_t size, OpBuilder &builder) {
+ // Gather/Scatter should have an integer type, but we check just in case.
+ auto intType = dyn_cast<IntegerType>(type);
+ if (!intType || fitsInIntegralType(size, intType)) {
+ return type;
+ }
+ if (fitsInIntegralType(size, builder.getI32Type())) {
+ return builder.getI32Type();
+ }
+ return builder.getI64Type();
+}
+
+// If `indicesBatchingDims` and `updatedIndexMap` are both sorted, then the
+// `indices_are_sorted` property is preserved: each concatenated iota is
+// monotonically increasing.
+bool getUpdatedIndicesAreSorted(bool indicesAreSorted,
+ ArrayRef<int64_t> indicesBatchingDims,
+ ArrayRef<int64_t> updatedIndexMap) {
+ return indicesAreSorted && llvm::is_sorted(indicesBatchingDims) &&
+ llvm::is_sorted(updatedIndexMap);
+}
+
+// Returns an updated indices tensor such that an `IotaOp` is prepended for each
+// dim in `indicesBatchingDims` with a `ConcatenateOp`.
+//
+// If `indexVectorDim` is equal to the rank of `indices`, it is reshaped to have
+// a trailing dimension of size 1 so it can be concatenated with the `IotaOp`s.
+Value createConcatIndices(Value indices, int64_t indexVectorDim,
+ ArrayRef<int64_t> indicesBatchingDims,
+ PatternRewriter &rewriter) {
+ Location loc = indices.getLoc();
+ auto indicesType = cast<RankedTensorType>(indices.getType());
+ Type elementType = indicesType.getElementType();
+
+ // The batching dim sizes might not fit in the existing element type, in which
+ // case we need to promote it.
+ for (int64_t batchingDim : indicesBatchingDims) {
+ elementType = promoteTypeForSize(
+ elementType, indicesType.getDimSize(batchingDim), rewriter);
+ }
+ if (elementType != indicesType.getElementType()) {
+ indicesType = RankedTensorType::get(indicesType.getShape(), elementType);
+ indices =
+ mlir::stablehlo::ConvertOp::create(rewriter, loc, indicesType, indices);
+ }
+
+ bool indexVectorDimOnLastDim = indexVectorDim == indicesType.getRank();
+ SmallVector<int64_t> iotaShape(indicesType.getShape());
+ if (indexVectorDimOnLastDim) {
+ iotaShape.push_back(1);
+ } else {
+ iotaShape[indexVectorDim] = 1;
+ }
+ auto iotaType = RankedTensorType::get(iotaShape, elementType);
+
+ if (indexVectorDimOnLastDim) {
+ indices =
+ mlir::stablehlo::ReshapeOp::create(rewriter, loc, iotaType, indices);
+ }
+
+ SmallVector<Value> indicesToConcat;
+ indicesToConcat.reserve(indicesBatchingDims.size() + 1);
+ for (int64_t batchingDim : indicesBatchingDims) {
+ indicesToConcat.push_back(
+ mlir::stablehlo::IotaOp::create(rewriter, loc, iotaType, batchingDim));
+ }
+ indicesToConcat.push_back(indices);
+ return mlir::stablehlo::ConcatenateOp::create(rewriter, loc, indicesToConcat,
+ indexVectorDim);
+}
+
+// Converts a `stablehlo.scatter` with batching dims to one without batching
+// dims, such that each batching dim becomes an inserted window dim with a
+// corresponding `IotaOp` concatenated to the scatter indices. This mirrors
+// upstream StableHLO's `ScatterWithBatchingDimsExpander`. The other scatter
+// canonicalization patterns do not understand batching dims and bail while any
+// remain (see `failIfScatterHasBatchingDims`).
+struct ScatterBatchingDimsExpander final
+ : OpRewritePattern<mlir::stablehlo::ScatterOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
+ PatternRewriter &rewriter) const override {
+ auto dimNumbers = op.getScatterDimensionNumbers();
+ ArrayRef<int64_t> inputBatchingDims = dimNumbers.getInputBatchingDims();
+ ArrayRef<int64_t> scatterIndicesBatchingDims =
+ dimNumbers.getScatterIndicesBatchingDims();
+ if (inputBatchingDims.empty()) {
+ return rewriter.notifyMatchFailure(op, "scatter op has no batching dims");
+ }
+
+ if (!cast<ShapedType>(op.getScatterIndices().getType()).hasStaticShape()) {
+ return rewriter.notifyMatchFailure(
+ op, "scatter indices have dynamic shape, can't expand");
+ }
+
+ SmallVector<int64_t> newInsertedWindowDims =
+ mergeSortedDims(inputBatchingDims, dimNumbers.getInsertedWindowDims());
+ SmallVector<int64_t> newScatterDimsToOperandDims =
+ llvm::to_vector(llvm::concat<const int64_t>(
+ inputBatchingDims, dimNumbers.getScatterDimsToOperandDims()));
+ Value newIndices = createConcatIndices(
+ op.getScatterIndices(), dimNumbers.getIndexVectorDim(),
+ scatterIndicesBatchingDims, rewriter);
+
+ auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
+ op.getContext(), dimNumbers.getUpdateWindowDims(),
+ newInsertedWindowDims,
+ /*inputBatchingDims=*/{}, /*scatterIndicesBatchingDims=*/{},
+ newScatterDimsToOperandDims, dimNumbers.getIndexVectorDim());
+
+ auto newScatter = mlir::stablehlo::ScatterOp::create(
+ rewriter, op.getLoc(), op->getResultTypes(), op.getInputs(), newIndices,
+ op.getUpdates(), newDimNumbers,
+ getUpdatedIndicesAreSorted(op.getIndicesAreSorted(),
+ scatterIndicesBatchingDims,
+ newScatterDimsToOperandDims),
+ op.getUniqueIndices());
+ newScatter.getUpdateComputation().takeBody(op.getUpdateComputation());
+ rewriter.replaceOp(op, newScatter.getResults());
+ return success();
+ }
+};
+
+// The scatter canonicalization patterns below assume a scatter with no batching
+// dims. This helper checks if there are any batching dims remaining so the
+// patterns can bail and let them be expanded by `ScatterBatchingDimsExpander`.
+static LogicalResult failIfScatterHasBatchingDims(mlir::stablehlo::ScatterOp op,
+ PatternRewriter &rewriter) {
+ auto dimNumbers = op.getScatterDimensionNumbers();
+ if (!dimNumbers.getInputBatchingDims().empty() ||
+ !dimNumbers.getScatterIndicesBatchingDims().empty()) {
+ return rewriter.notifyMatchFailure(
+ op,
+ "scatter has batching dims; ScatterBatchingDimsExpander runs first");
+ }
+ return success();
+}
+
+// Converts a `stablehlo.scatter` that writes a single, full-rank contiguous
+// slice with an overwrite computation into a `stablehlo.dynamic_update_slice`.
+//
+// Such scatters (a scalar index selecting an offset into an operand dim that is
+// *not* collapsed, i.e. a partial slice) are dynamic-update-slice semantics and
+// cannot be represented as an `iree_linalg_ext.scatter`, which requires the
+// index depth to map to fully-inserted leading operand dims.
+struct ScatterToDynamicUpdateSlice final
+ : OpRewritePattern<mlir::stablehlo::ScatterOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.getInputs().size() != 1 || op.getUpdates().size() != 1) {
+ return rewriter.notifyMatchFailure(op, "variadic scatter");
+ }
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
+ auto dimNumbers = op.getScatterDimensionNumbers();
+ if (!dimNumbers.getInsertedWindowDims().empty()) {
+ return rewriter.notifyMatchFailure(op, "has inserted dims");
+ }
+
+ Value operand = op.getInputs().front();
+ Value update = op.getUpdates().front();
+ Value indices = op.getScatterIndices();
+ auto operandTy = dyn_cast<RankedTensorType>(operand.getType());
+ auto updateTy = dyn_cast<RankedTensorType>(update.getType());
+ auto indicesTy = dyn_cast<RankedTensorType>(indices.getType());
+ if (!operandTy || !updateTy || !indicesTy || !operandTy.hasStaticShape() ||
+ !updateTy.hasStaticShape() || !indicesTy.hasStaticShape()) {
+ return rewriter.notifyMatchFailure(op, "dynamic/unranked shapes");
+ }
+
+ int64_t rank = operandTy.getRank();
+
+ // The update must be a single full-rank contiguous slice: no batch dims and
+ // an identity window mapping over all operand dims.
+ ArrayRef<int64_t> updateWindowDims = dimNumbers.getUpdateWindowDims();
+ if (updateTy.getRank() != rank ||
+ static_cast<int64_t>(updateWindowDims.size()) != rank) {
+ return rewriter.notifyMatchFailure(op, "update is not a full-rank slice");
+ }
+ for (auto [idx, dim] : llvm::enumerate(updateWindowDims)) {
+ if (static_cast<int64_t>(idx) != dim) {
+ return rewriter.notifyMatchFailure(op, "non-identity window dims");
+ }
+ }
+
+ // There must be exactly one scatter index vector.
+ int64_t indexVectorDim = dimNumbers.getIndexVectorDim();
+ int64_t numUpdates = 1;
+ for (int64_t d = 0, s = indicesTy.getRank(); d < s; ++d) {
+ if (d != indexVectorDim) {
+ numUpdates *= indicesTy.getDimSize(d);
+ }
+ }
+ if (numUpdates != 1) {
+ return rewriter.notifyMatchFailure(op, "more than one scatter index");
+ }
+
+ // The update computation must be a plain overwrite: return the update value
+ // (block argument #1) unchanged.
+ Region ®ion = op.getUpdateComputation();
+ if (!region.hasOneBlock() || region.front().getNumArguments() != 2) {
+ return rewriter.notifyMatchFailure(op, "unexpected update computation");
+ }
+ Block &block = region.front();
+ auto retOp = dyn_cast<mlir::stablehlo::ReturnOp>(block.getTerminator());
+ if (!retOp || retOp.getNumOperands() != 1 ||
+ retOp.getOperand(0) != block.getArgument(1)) {
+ return rewriter.notifyMatchFailure(op, "not an overwrite computation");
+ }
+
+ ArrayRef<int64_t> scatterDims = dimNumbers.getScatterDimsToOperandDims();
+ int64_t indexDepth = scatterDims.size();
+
+ ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+ Type indexElemTy = indicesTy.getElementType();
+ auto scalarTy = RankedTensorType::get({}, indexElemTy);
+
+ // Flatten indices to `indexDepth` scalar components (numUpdates == 1).
+ auto flatTy = RankedTensorType::get({indexDepth}, indexElemTy);
+ Value flatIndices = mlir::stablehlo::ReshapeOp::create(b, flatTy, indices);
+
+ Value zero =
+ mlir::stablehlo::ConstantOp::create(b, b.getZeroAttr(indexElemTy));
+
+ // Map each operand dim to a start offset: the scatter index component if
+ // the dim is indexed, otherwise zero. dynamic_update_slice clamps offsets.
+ auto sliceTy = RankedTensorType::get({1}, indexElemTy);
+ llvm::SmallVector<Value> startIndices(rank, zero);
+ for (auto [j, operandDim] : llvm::enumerate(scatterDims)) {
+ Value comp = mlir::stablehlo::SliceOp::create(
+ b, sliceTy, flatIndices,
+ b.getDenseI64ArrayAttr({static_cast<int64_t>(j)}),
+ b.getDenseI64ArrayAttr({static_cast<int64_t>(j) + 1}),
+ b.getDenseI64ArrayAttr({1}));
+ startIndices[operandDim] =
+ mlir::stablehlo::ReshapeOp::create(b, scalarTy, comp);
+ }
+
+ rewriter.replaceOpWithNewOp<mlir::stablehlo::DynamicUpdateSliceOp>(
+ op, operandTy, operand, update, startIndices);
+ return success();
+ }
+};
+
struct ScatterInt64Indices final
: OpRewritePattern<mlir::stablehlo::ScatterOp> {
using Base::Base;
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
auto indices = op.getScatterIndices();
auto indicesTy = indices.getType();
auto indicesETy = indicesTy.getElementType();
@@ -504,6 +785,9 @@
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
auto dimNumbers = op.getScatterDimensionNumbers();
auto indexVectorDim = dimNumbers.getIndexVectorDim();
Value indices = op.getScatterIndices();
@@ -572,6 +856,9 @@
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
auto dimNumbers = op.getScatterDimensionNumbers();
auto indexVectorDim = dimNumbers.getIndexVectorDim();
auto indices = cast<Value>(op.getScatterIndices());
@@ -662,6 +949,9 @@
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
auto dimNumbers = op.getScatterDimensionNumbers();
auto indexVectorDim = dimNumbers.getIndexVectorDim();
auto indices = cast<Value>(op.getScatterIndices());
@@ -736,6 +1026,9 @@
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
auto dimNumbers = op.getScatterDimensionNumbers();
@@ -851,6 +1144,9 @@
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
Location loc = op.getLoc();
auto dimNumbers = op.getScatterDimensionNumbers();
ArrayRef<int64_t> scatterDimsToOperandDims =
@@ -1074,6 +1370,9 @@
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
+ if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
+ return failure();
+ }
auto indices = op.getScatterIndices();
Value operand = op.getInputs().front();
auto indicesTy = cast<ShapedType>(indices.getType());
@@ -2176,10 +2475,11 @@
patterns.insert<RngBitcastFloat>(context);
// scatter canonicalization patterns
- patterns
- .insert<ScatterInt64Indices, ScatterImplicitIndex, ScatterImplicitBatch,
- ScatterMaterializeInsertedDim, ScatterCollapseBatch,
- ScatterBatchFirst, ScatterIndexedDimsFirst>(context);
+ patterns.insert<ScatterToDynamicUpdateSlice>(context, /*benefit=*/2);
+ patterns.insert<ScatterBatchingDimsExpander, ScatterInt64Indices,
+ ScatterImplicitIndex, ScatterImplicitBatch,
+ ScatterMaterializeInsertedDim, ScatterCollapseBatch,
+ ScatterBatchFirst, ScatterIndexedDimsFirst>(context);
// dot_general canonicalization patterns.
populatePreprocessingDotGeneralToDotPatterns(context, &patterns);
diff --git a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir
index 97bcedb..3d06187 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir
+++ b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir
@@ -557,6 +557,25 @@
// -----
+// Scatter batching dims are lowered to an explicit iota index column so that no
+// batching dims remain for the downstream canonicalizers.
+// CHECK-LABEL: @scatter_batching_dims
+func.func @scatter_batching_dims(%operand: tensor<2x5xi32>, %indices: tensor<2x3x1xi32>, %updates: tensor<2x3xi32>) -> tensor<2x5xi32> {
+ // CHECK: %[[IOTA:.+]] = stablehlo.iota dim = 0 : tensor<2x3x1xi32>
+ // CHECK: %[[CAT:.+]] = stablehlo.concatenate %[[IOTA]], %{{.+}}, dim = 2
+ // CHECK: stablehlo.scatter
+ // CHECK-NOT: input_batching_dims
+ // CHECK-SAME: scatter_dims_to_operand_dims = [0, 1]
+ %0 = "stablehlo.scatter"(%operand, %indices, %updates) <{scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [1], input_batching_dims = [0], scatter_indices_batching_dims = [0], scatter_dims_to_operand_dims = [1], index_vector_dim = 2>}> ({
+ ^bb0(%a: tensor<i32>, %b: tensor<i32>):
+ %m = stablehlo.minimum %a, %b : tensor<i32>
+ stablehlo.return %m : tensor<i32>
+ }) : (tensor<2x5xi32>, tensor<2x3x1xi32>, tensor<2x3xi32>) -> tensor<2x5xi32>
+ return %0 : tensor<2x5xi32>
+}
+
+// -----
+
// ScatterIndexedDimsFirst transposes the indexed operand dim to the front, so the rewritten
// scatter indexes operand dim 0 and the idx == dim check passes.
@@ -609,6 +628,23 @@
// -----
+// A single-index, full-rank, overwrite scatter (dynamic-update-slice semantics)
+// becomes a stablehlo.dynamic_update_slice.
+// CHECK-LABEL: @scatter_to_dynamic_update_slice
+func.func @scatter_to_dynamic_update_slice(%operand: tensor<2x7x3x3xf32>, %indices: tensor<1xi32>, %updates: tensor<2x4x3x3xf32>) -> tensor<2x7x3x3xf32> {
+ // CHECK: %[[C0:.+]] = stablehlo.constant dense<0> : tensor<i32>
+ // CHECK: %[[IDX:.+]] = stablehlo.reshape %{{.+}} : (tensor<1xi32>) -> tensor<i32>
+ // CHECK: stablehlo.dynamic_update_slice %{{.+}}, %{{.+}}, %[[C0]], %[[IDX]], %[[C0]], %[[C0]]
+ // CHECK-NOT: stablehlo.scatter
+ %0 = "stablehlo.scatter"(%operand, %indices, %updates) <{indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0, 1, 2, 3], scatter_dims_to_operand_dims = [1]>, unique_indices = true}> ({
+ ^bb0(%a: tensor<f32>, %b: tensor<f32>):
+ stablehlo.return %b : tensor<f32>
+ }) : (tensor<2x7x3x3xf32>, tensor<1xi32>, tensor<2x4x3x3xf32>) -> tensor<2x7x3x3xf32>
+ return %0 : tensor<2x7x3x3xf32>
+}
+
+// -----
+
// CHECK-LABEL: @scatter_implicit_index_dim
// CHECK: %[[INDICES:.+]] = tensor.expand_shape %arg1 {{.*}} : tensor<1xi32> into tensor<1x1xi32>
// CHECK: %[[SCATTER:.+]] = "stablehlo.scatter"(%arg0, %[[INDICES]], %arg2)
