[Codegen][VectorExt] Add TransferScatterOp unrolling support (#23704)
Add scatter unrolling to `LowerTransferGatherScatterOps` (renamed from
`LowerTransferGatherOps`). Templatized helpers with gather:
`removeDim0FromMap`, `computeUnrollDim0Maps`,
`extractSlicesForIteration`, `computeNewOffsets`.
Added subsequent lit tests.
Co-Authored-By: Claude Opus 4.6
[noreply@anthropic.com](mailto:noreply@anthropic.com)
Part 4/4 from https://github.com/iree-org/iree/pull/23610
---------
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BUILD.bazel
index 8827f14..a31b928 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BUILD.bazel
@@ -37,7 +37,7 @@
srcs = [
"BufferizationInterfaces.cpp",
"DistributionPatterns.cpp",
- "LowerTransferGatherOps.cpp",
+ "LowerTransferGatherScatterOps.cpp",
"Passes.cpp",
"VectorExtFoldUnitExtentDims.cpp",
],
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/CMakeLists.txt
index 57219b5..1e53aa4 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/CMakeLists.txt
@@ -31,7 +31,7 @@
SRCS
"BufferizationInterfaces.cpp"
"DistributionPatterns.cpp"
- "LowerTransferGatherOps.cpp"
+ "LowerTransferGatherScatterOps.cpp"
"Passes.cpp"
"VectorExtFoldUnitExtentDims.cpp"
DEPS
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/LowerTransferGatherOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/LowerTransferGatherOps.cpp
deleted file mode 100644
index 74db010..0000000
--- a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/LowerTransferGatherOps.cpp
+++ /dev/null
@@ -1,235 +0,0 @@
-// Copyright 2025 The IREE Authors
-//
-// Licensed under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
-#include "iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
-#include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/PatternMatch.h"
-
-using namespace mlir;
-using namespace mlir::iree_compiler::IREE::VectorExt;
-
-namespace {
-
-/// Remove dim 0 from an AffineMap by:
-/// 1. Replacing AffineDimExpr(0) with AffineConstantExpr(0)
-/// 2. Renumbering AffineDimExpr(k) where k > 0 to AffineDimExpr(k-1)
-/// 3. Reducing numDims by 1
-static AffineMap removeDim0FromMap(AffineMap map) {
- MLIRContext *ctx = map.getContext();
- SmallVector<AffineExpr> newResults;
- for (AffineExpr expr : map.getResults()) {
- if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
- unsigned pos = dimExpr.getPosition();
- if (pos == 0) {
- newResults.push_back(getAffineConstantExpr(0, ctx));
- } else {
- newResults.push_back(getAffineDimExpr(pos - 1, ctx));
- }
- } else {
- newResults.push_back(expr);
- }
- }
- return AffineMap::get(map.getNumDims() - 1, map.getNumSymbols(), newResults,
- ctx);
-}
-
-/// Remove dim 0 references from an index vec map. Returns the new map with
-/// results that referenced dim 0 dropped, and the axis positions in the index
-/// vec that need to be sliced.
-static AffineMap removeDim0FromIndexVecMap(AffineMap map,
- SmallVectorImpl<int64_t> &axes) {
- MLIRContext *ctx = map.getContext();
- SmallVector<AffineExpr> newResults;
- for (auto [resultIdx, expr] : llvm::enumerate(map.getResults())) {
- if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
- unsigned pos = dimExpr.getPosition();
- if (pos == 0) {
- axes.push_back(resultIdx);
- continue; // Drop this result.
- }
- newResults.push_back(getAffineDimExpr(pos - 1, ctx));
- } else {
- newResults.push_back(expr);
- }
- }
- return AffineMap::get(map.getNumDims() - 1, map.getNumSymbols(), newResults,
- ctx);
-}
-
-/// Extract a slice from a vector at position `idx` along the given `axis`.
-/// For a vector<4x8xindex>, extracting axis=0, idx=2 gives vector<8xindex>.
-static Value extractVecSlice(OpBuilder &b, Location loc, Value vec,
- int64_t axis, int64_t idx) {
- auto vecType = cast<VectorType>(vec.getType());
- int64_t rank = vecType.getRank();
-
- if (axis == 0) {
- // Extracting from rank-1 along axis 0 gives a scalar.
- return vector::ExtractOp::create(b, loc, vec, SmallVector<int64_t>{idx});
- }
-
- // General case: use extract_strided_slice.
- SmallVector<int64_t> offsets(rank, 0);
- SmallVector<int64_t> sizes(vecType.getShape());
- SmallVector<int64_t> strides(rank, 1);
- offsets[axis] = idx;
- sizes[axis] = 1;
- Value slice = vector::ExtractStridedSliceOp::create(b, loc, vec, offsets,
- sizes, strides);
- // Drop the unit dim.
- SmallVector<int64_t> newShape;
- for (int64_t i = 0; i < rank; ++i) {
- if (i != axis) {
- newShape.push_back(vecType.getShape()[i]);
- }
- }
- auto newType = VectorType::get(newShape, vecType.getElementType());
- return vector::ShapeCastOp::create(b, loc, newType, slice);
-}
-
-//===----------------------------------------------------------------------===//
-// UnrollTransferGatherDim
-//===----------------------------------------------------------------------===//
-
-/// Unrolls dim 0 of a transfer_gather, reducing vector rank by 1 each
-/// application. Stops at rank 1.
-struct UnrollTransferGatherDim : OpRewritePattern<TransferGatherOp> {
- using Base::Base;
-
- LogicalResult matchAndRewrite(TransferGatherOp op,
- PatternRewriter &rewriter) const override {
- VectorType resultType = op.getVector().getType();
- int64_t rank = resultType.getRank();
- if (rank <= 1) {
- return rewriter.notifyMatchFailure(op, "already rank <= 1");
- }
-
- Location loc = op.getLoc();
- int64_t dim0Size = resultType.getShape()[0];
- SmallVector<AffineMap> indexingMaps = op.getIndexingMapsArray();
- AffineMap sourceMap = indexingMaps[0];
- OperandRange indexVecs = op.getIndexVecs();
- int64_t numIndexVecs = indexVecs.size();
- Value mask = op.getMask();
-
- // Compute the new source map (dim 0 removed).
- AffineMap newSourceMap = removeDim0FromMap(sourceMap);
-
- // For each index vec, compute how dim 0 removal affects it.
- SmallVector<AffineMap> newIndexVecMaps;
- SmallVector<SmallVector<int64_t>> indexVecAxes; // axes to slice per vec
- for (int64_t i = 0; i < numIndexVecs; ++i) {
- SmallVector<int64_t> axes;
- AffineMap newMap = removeDim0FromIndexVecMap(indexingMaps[1 + i], axes);
- newIndexVecMaps.push_back(newMap);
- indexVecAxes.push_back(std::move(axes));
- }
-
- // Handle mask map.
- AffineMap newMaskMap;
- SmallVector<int64_t> maskAxes;
- if (mask) {
- newMaskMap = removeDim0FromIndexVecMap(indexingMaps.back(), maskAxes);
- }
-
- // Find which source dims use AffineDimExpr(0) — these need offset updates.
- SmallVector<int64_t> sourceDimsUsingDim0;
- for (auto [j, expr] : llvm::enumerate(sourceMap.getResults())) {
- if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
- if (dimExpr.getPosition() == 0) {
- sourceDimsUsingDim0.push_back(j);
- }
- }
- }
-
- // Build the new result vector type (dim 0 removed).
- SmallVector<int64_t> newShape(resultType.getShape().drop_front());
- auto newResultType = VectorType::get(newShape, resultType.getElementType());
-
- // Build new indexing_maps array.
- SmallVector<AffineMap> newAllMaps;
- newAllMaps.push_back(newSourceMap);
- for (AffineMap &m : newIndexVecMaps) {
- newAllMaps.push_back(m);
- }
- if (mask) {
- newAllMaps.push_back(newMaskMap);
- }
-
- // Initialize accumulator.
- Value acc = ub::PoisonOp::create(rewriter, loc, resultType);
-
- for (int64_t i = 0; i < dim0Size; ++i) {
- // Compute new base offsets.
- SmallVector<Value> newOffsets(op.getOffsets());
- for (int64_t srcDim : sourceDimsUsingDim0) {
- Value offset = newOffsets[srcDim];
- Value iVal = arith::ConstantIndexOp::create(rewriter, loc, i);
- newOffsets[srcDim] = arith::AddIOp::create(rewriter, loc, offset, iVal);
- }
-
- // Extract index vec slices.
- SmallVector<Value> newIndexVecs;
- for (int64_t k = 0; k < numIndexVecs; ++k) {
- Value idxVec = indexVecs[k];
- if (indexVecAxes[k].empty()) {
- // This index vec doesn't reference dim 0 — use as-is.
- newIndexVecs.push_back(idxVec);
- } else {
- // Extract along each axis that referenced dim 0.
- // Since maps only have simple dim exprs, there should be at most
- // one axis referencing dim 0.
- for (int64_t axis : indexVecAxes[k]) {
- idxVec = extractVecSlice(rewriter, loc, idxVec, axis, i);
- }
- newIndexVecs.push_back(idxVec);
- }
- }
-
- // Extract mask slice.
- Value newMask;
- if (mask) {
- if (maskAxes.empty()) {
- newMask = mask;
- } else {
- Value m = mask;
- for (int64_t axis : maskAxes) {
- m = extractVecSlice(rewriter, loc, m, axis, i);
- }
- newMask = m;
- }
- }
-
- auto subGather = TransferGatherOp::create(
- rewriter, loc, newResultType, op.getBase(), newOffsets, newIndexVecs,
- rewriter.getAffineMapArrayAttr(newAllMaps), op.getPadding(), newMask);
-
- // Insert into accumulator.
- SmallVector<int64_t> offsets(rank, 0);
- offsets[0] = i;
- SmallVector<int64_t> strides(newShape.size(), 1);
- acc = vector::InsertStridedSliceOp::create(
- rewriter, loc, subGather.getResult(), acc, offsets, strides);
- }
-
- rewriter.replaceOp(op, acc);
- return success();
- }
-};
-
-} // namespace
-
-namespace mlir::iree_compiler::IREE::VectorExt {
-
-void populateVectorTransferGatherLoweringPatterns(RewritePatternSet &patterns) {
- patterns.add<UnrollTransferGatherDim>(patterns.getContext());
-}
-
-} // namespace mlir::iree_compiler::IREE::VectorExt
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/LowerTransferGatherScatterOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/LowerTransferGatherScatterOps.cpp
new file mode 100644
index 0000000..c74bc88
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/LowerTransferGatherScatterOps.cpp
@@ -0,0 +1,314 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::VectorExt;
+
+namespace {
+
+/// Remove dim 0 from an AffineMap by:
+/// 1. Replacing AffineDimExpr(0) with AffineConstantExpr(0)
+/// 2. Renumbering AffineDimExpr(k) where k > 0 to AffineDimExpr(k-1)
+/// 3. Reducing numDims by 1
+static AffineMap removeDim0FromMap(AffineMap map) {
+ MLIRContext *ctx = map.getContext();
+ SmallVector<AffineExpr> newResults;
+ for (AffineExpr expr : map.getResults()) {
+ if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+ unsigned pos = dimExpr.getPosition();
+ if (pos == 0) {
+ newResults.push_back(getAffineConstantExpr(0, ctx));
+ } else {
+ newResults.push_back(getAffineDimExpr(pos - 1, ctx));
+ }
+ } else {
+ newResults.push_back(expr);
+ }
+ }
+ return AffineMap::get(map.getNumDims() - 1, map.getNumSymbols(), newResults,
+ ctx);
+}
+
+/// Remove dim 0 references from an index vec map. Returns the new map with
+/// results that referenced dim 0 dropped, and the axis positions in the index
+/// vec that need to be sliced.
+static AffineMap removeDim0FromIndexVecMap(AffineMap map,
+ SmallVectorImpl<int64_t> &axes) {
+ MLIRContext *ctx = map.getContext();
+ SmallVector<AffineExpr> newResults;
+ for (auto [resultIdx, expr] : llvm::enumerate(map.getResults())) {
+ if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+ unsigned pos = dimExpr.getPosition();
+ if (pos == 0) {
+ axes.push_back(resultIdx);
+ continue;
+ }
+ newResults.push_back(getAffineDimExpr(pos - 1, ctx));
+ } else {
+ newResults.push_back(expr);
+ }
+ }
+ return AffineMap::get(map.getNumDims() - 1, map.getNumSymbols(), newResults,
+ ctx);
+}
+
+/// Extract a slice from a vector at position `idx` along the given `axis`.
+/// For a vector<4x8xindex>, extracting axis=0, idx=2 gives vector<8xindex>.
+static Value extractVecSlice(OpBuilder &b, Location loc, Value vec,
+ int64_t axis, int64_t idx) {
+ auto vecType = cast<VectorType>(vec.getType());
+ int64_t rank = vecType.getRank();
+
+ if (axis == 0) {
+ // Extracting from rank-1 along axis 0 gives a scalar.
+ return vector::ExtractOp::create(b, loc, vec, int64_t{idx});
+ }
+
+ // General case: use extract_strided_slice.
+ SmallVector<int64_t> offsets(rank, 0);
+ SmallVector<int64_t> sizes(vecType.getShape());
+ SmallVector<int64_t> strides(rank, 1);
+ offsets[axis] = idx;
+ sizes[axis] = 1;
+ Value slice = vector::ExtractStridedSliceOp::create(b, loc, vec, offsets,
+ sizes, strides);
+ // Drop the unit dim.
+ SmallVector<int64_t> newShape;
+ for (int64_t i = 0; i < rank; ++i) {
+ if (i != axis) {
+ newShape.push_back(vecType.getShape()[i]);
+ }
+ }
+ auto newType = VectorType::get(newShape, vecType.getElementType());
+ return vector::ShapeCastOp::create(b, loc, newType, slice);
+}
+
+//===----------------------------------------------------------------------===//
+// Shared unroll helpers
+//===----------------------------------------------------------------------===//
+
+/// Compute dim-0-removed indexing maps for unrolling. Populates the new base
+/// map, per-index-vec maps and axes, mask map and axes, base dims using dim 0,
+/// and the combined new indexing maps array.
+static void
+computeUnrollDim0Maps(ArrayRef<AffineMap> indexingMaps, int64_t numIndexVecs,
+ bool hasMask, AffineMap baseMap,
+ SmallVectorImpl<AffineMap> &newAllMaps,
+ SmallVectorImpl<SmallVector<int64_t>> &indexVecAxes,
+ SmallVectorImpl<int64_t> &maskAxes,
+ SmallVectorImpl<int64_t> &baseDimsUsingDim0) {
+ AffineMap newBaseMap = removeDim0FromMap(baseMap);
+ newAllMaps.push_back(newBaseMap);
+
+ for (int64_t i = 0; i < numIndexVecs; ++i) {
+ SmallVector<int64_t> axes;
+ AffineMap newMap = removeDim0FromIndexVecMap(indexingMaps[1 + i], axes);
+ newAllMaps.push_back(newMap);
+ indexVecAxes.push_back(std::move(axes));
+ }
+
+ if (hasMask) {
+ AffineMap newMaskMap =
+ removeDim0FromIndexVecMap(indexingMaps.back(), maskAxes);
+ newAllMaps.push_back(newMaskMap);
+ }
+
+ for (auto [j, expr] : llvm::enumerate(baseMap.getResults())) {
+ if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+ if (dimExpr.getPosition() == 0) {
+ baseDimsUsingDim0.push_back(j);
+ }
+ }
+ }
+}
+
+/// Extract sliced index vecs and mask for iteration `i` of dim-0 unrolling.
+/// Returns the sliced mask (or nullptr if no mask).
+static Value
+extractSlicesForIteration(OpBuilder &rewriter, Location loc, int64_t i,
+ OperandRange indexVecs, int64_t numIndexVecs,
+ ArrayRef<SmallVector<int64_t>> indexVecAxes,
+ Value mask, ArrayRef<int64_t> maskAxes,
+ SmallVectorImpl<Value> &newIndexVecs) {
+ for (int64_t k = 0; k < numIndexVecs; ++k) {
+ Value idxVec = indexVecs[k];
+ if (!indexVecAxes[k].empty()) {
+ for (int64_t axis : indexVecAxes[k]) {
+ idxVec = extractVecSlice(rewriter, loc, idxVec, axis, i);
+ }
+ }
+ newIndexVecs.push_back(idxVec);
+ }
+
+ if (!mask) {
+ return nullptr;
+ }
+ Value m = mask;
+ for (int64_t axis : maskAxes) {
+ m = extractVecSlice(rewriter, loc, m, axis, i);
+ }
+ return m;
+}
+
+/// Update base offsets for dim-0 iteration `i`.
+static SmallVector<Value>
+computeNewOffsets(OpBuilder &rewriter, Location loc, ValueRange offsets,
+ int64_t i, ArrayRef<int64_t> baseDimsUsingDim0) {
+ SmallVector<Value> newOffsets(offsets);
+ for (int64_t baseDim : baseDimsUsingDim0) {
+ Value offset = newOffsets[baseDim];
+ Value iVal = arith::ConstantIndexOp::create(rewriter, loc, i);
+ newOffsets[baseDim] = arith::AddIOp::create(rewriter, loc, offset, iVal);
+ }
+ return newOffsets;
+}
+
+//===----------------------------------------------------------------------===//
+// UnrollTransferGatherDim / UnrollTransferScatterDim
+//===----------------------------------------------------------------------===//
+
+/// Unrolls dim 0 of a transfer_gather, reducing vector rank by 1 each
+/// application. Sub-gathers are assembled into the result via
+/// insert_strided_slice. Stops at rank 1.
+struct UnrollTransferGatherDim final : OpRewritePattern<TransferGatherOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(TransferGatherOp op,
+ PatternRewriter &rewriter) const override {
+ VectorType vectorType = op.getVector().getType();
+ int64_t rank = vectorType.getRank();
+ if (rank <= 1) {
+ return rewriter.notifyMatchFailure(op, "already rank <= 1");
+ }
+
+ Location loc = op.getLoc();
+ int64_t dim0Size = vectorType.getShape()[0];
+ SmallVector<AffineMap> indexingMaps = op.getIndexingMapsArray();
+ OperandRange indexVecs = op.getIndexVecs();
+ int64_t numIndexVecs = indexVecs.size();
+ Value mask = op.getMask();
+
+ SmallVector<AffineMap> newAllMaps;
+ SmallVector<SmallVector<int64_t>> indexVecAxes;
+ SmallVector<int64_t> maskAxes, baseDimsUsingDim0;
+ computeUnrollDim0Maps(indexingMaps, numIndexVecs, !!mask, indexingMaps[0],
+ newAllMaps, indexVecAxes, maskAxes,
+ baseDimsUsingDim0);
+
+ SmallVector<int64_t> newShape(vectorType.getShape().drop_front());
+ auto newVectorType = VectorType::get(newShape, vectorType.getElementType());
+
+ Value acc = ub::PoisonOp::create(rewriter, loc, vectorType);
+
+ for (int64_t i = 0; i < dim0Size; ++i) {
+ SmallVector<Value> newOffsets = computeNewOffsets(
+ rewriter, loc, op.getOffsets(), i, baseDimsUsingDim0);
+
+ SmallVector<Value> newIndexVecs;
+ Value newMask =
+ extractSlicesForIteration(rewriter, loc, i, indexVecs, numIndexVecs,
+ indexVecAxes, mask, maskAxes, newIndexVecs);
+
+ auto subGather = TransferGatherOp::create(
+ rewriter, loc, newVectorType, op.getBase(), newOffsets, newIndexVecs,
+ rewriter.getAffineMapArrayAttr(newAllMaps), op.getPadding(), newMask);
+
+ SmallVector<int64_t> offsets(rank, 0);
+ offsets[0] = i;
+ SmallVector<int64_t> strides(newShape.size(), 1);
+ acc = vector::InsertStridedSliceOp::create(
+ rewriter, loc, subGather.getResult(), acc, offsets, strides);
+ }
+
+ rewriter.replaceOp(op, acc);
+ return success();
+ }
+};
+
+/// Unrolls dim 0 of a transfer_scatter, reducing vector rank by 1 each
+/// application. For tensor semantics, sub-scatters are chained via SSA
+/// results. For memref semantics, sub-scatters write in-place. Stops at
+/// rank 1.
+struct UnrollTransferScatterDim final : OpRewritePattern<TransferScatterOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(TransferScatterOp op,
+ PatternRewriter &rewriter) const override {
+ VectorType vectorType = op.getVectorType();
+ int64_t rank = vectorType.getRank();
+ if (rank <= 1) {
+ return rewriter.notifyMatchFailure(op, "already rank <= 1");
+ }
+
+ Location loc = op.getLoc();
+ int64_t dim0Size = vectorType.getShape()[0];
+ SmallVector<AffineMap> indexingMaps = op.getIndexingMapsArray();
+ OperandRange indexVecs = op.getIndexVecs();
+ int64_t numIndexVecs = indexVecs.size();
+ Value mask = op.getMask();
+
+ SmallVector<AffineMap> newAllMaps;
+ SmallVector<SmallVector<int64_t>> indexVecAxes;
+ SmallVector<int64_t> maskAxes, baseDimsUsingDim0;
+ computeUnrollDim0Maps(indexingMaps, numIndexVecs, !!mask, indexingMaps[0],
+ newAllMaps, indexVecAxes, maskAxes,
+ baseDimsUsingDim0);
+
+ Value dest = op.getBase();
+
+ for (int64_t i = 0; i < dim0Size; ++i) {
+ SmallVector<Value> newOffsets = computeNewOffsets(
+ rewriter, loc, op.getOffsets(), i, baseDimsUsingDim0);
+
+ SmallVector<Value> newIndexVecs;
+ Value newMask =
+ extractSlicesForIteration(rewriter, loc, i, indexVecs, numIndexVecs,
+ indexVecAxes, mask, maskAxes, newIndexVecs);
+
+ Value vecSlice =
+ vector::ExtractOp::create(rewriter, loc, op.getVector(), int64_t{i});
+
+ if (op.hasTensorSemantics()) {
+ auto subScatter = TransferScatterOp::create(
+ rewriter, loc, dest.getType(), dest, vecSlice, newOffsets,
+ newIndexVecs, rewriter.getAffineMapArrayAttr(newAllMaps), newMask);
+ dest = subScatter.getResult();
+ continue;
+ }
+ TransferScatterOp::create(rewriter, loc, /*resultTypes=*/TypeRange{},
+ dest, vecSlice, newOffsets, newIndexVecs,
+ rewriter.getAffineMapArrayAttr(newAllMaps),
+ newMask);
+ }
+
+ if (op.hasTensorSemantics()) {
+ rewriter.replaceOp(op, dest);
+ } else {
+ rewriter.eraseOp(op);
+ }
+ return success();
+ }
+};
+
+} // namespace
+
+namespace mlir::iree_compiler::IREE::VectorExt {
+
+void populateVectorTransferGatherScatterLoweringPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<UnrollTransferGatherDim, UnrollTransferScatterDim>(
+ patterns.getContext());
+}
+
+} // namespace mlir::iree_compiler::IREE::VectorExt
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h
index d533e42..c15f2e7 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h
@@ -11,7 +11,8 @@
namespace mlir::iree_compiler::IREE::VectorExt {
-void populateVectorTransferGatherLoweringPatterns(RewritePatternSet &patterns);
+void populateVectorTransferGatherScatterLoweringPatterns(
+ RewritePatternSet &patterns);
}; // namespace mlir::iree_compiler::IREE::VectorExt
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUVectorLowering.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUVectorLowering.cpp
index ebdb9b4..b8e7c53 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUVectorLowering.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUVectorLowering.cpp
@@ -617,12 +617,14 @@
contractLoweringPatterns,
vector::VectorMultiReductionLowering::InnerReduction);
contractLoweringPatterns.add<UnrollElementwiseOps>(funcOp->getContext());
- // Unroll transfer_gather ops to rank 1 and lower contiguous ones to
- // vector.transfer_read.
- IREE::VectorExt::populateVectorTransferGatherLoweringPatterns(
+ // Unroll transfer_gather/scatter ops to rank 1 and lower contiguous ones
+ // to vector.transfer_read/write.
+ IREE::VectorExt::populateVectorTransferGatherScatterLoweringPatterns(
contractLoweringPatterns);
IREE::VectorExt::TransferGatherOp::getCanonicalizationPatterns(
contractLoweringPatterns, ctx);
+ IREE::VectorExt::TransferScatterOp::getCanonicalizationPatterns(
+ contractLoweringPatterns, ctx);
if (failed(applyPatternsGreedily(funcOp,
std::move(contractLoweringPatterns)))) {
return signalPassFailure();
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/vector_lowering.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/vector_lowering.mlir
index 4519a73..f8727cc 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/vector_lowering.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/vector_lowering.mlir
@@ -300,3 +300,95 @@
// CHECK-NOT: transfer_gather
// CHECK-COUNT-32: vector.load
// CHECK-NOT: transfer_gather
+
+// -----
+
+// Test unrolling of a 2D transfer_scatter: outer dim is scattered (indices),
+// inner dim is contiguous.
+
+func.func @transfer_scatter_unroll_embedding_write(
+ %source: memref<4096x64xf16>,
+ %vector: vector<4x64xf16>,
+ %indices: vector<4xindex>) {
+ %c0 = arith.constant 0 : index
+ iree_vector_ext.transfer_scatter %vector into %source[%c0, %c0]
+ [%indices : vector<4xindex>] {
+ indexing_maps = [affine_map<(d0, d1)[s0] -> (s0, d1)>,
+ affine_map<(d0, d1)[s0] -> (d0)>]
+ } : vector<4x64xf16>, memref<4096x64xf16>
+ return
+}
+
+// After unrolling, the 2D scatter becomes 4 rank-1 sub-scatters.
+// CHECK-LABEL: func.func @transfer_scatter_unroll_embedding_write
+// CHECK-COUNT-4: transfer_scatter {{.+}} : vector<64xf16>
+
+// -----
+
+// Test unrolling of a masked 2D transfer_scatter.
+
+func.func @transfer_scatter_unroll_masked(
+ %source: memref<4096x64xf16>,
+ %vector: vector<4x64xf16>,
+ %indices: vector<4xindex>,
+ %mask: vector<4x64xi1>) {
+ %c0 = arith.constant 0 : index
+ iree_vector_ext.transfer_scatter %vector into %source[%c0, %c0]
+ [%indices : vector<4xindex>], %mask {
+ indexing_maps = [affine_map<(d0, d1)[s0] -> (s0, d1)>,
+ affine_map<(d0, d1)[s0] -> (d0)>,
+ affine_map<(d0, d1)[s0] -> (d0, d1)>]
+ } : vector<4x64xf16>, memref<4096x64xf16>, vector<4x64xi1>
+ return
+}
+
+// After unrolling, mask slices are passed to each sub-scatter.
+// CHECK-LABEL: func.func @transfer_scatter_unroll_masked
+// CHECK-COUNT-4: transfer_scatter {{.+}} : vector<64xf16>
+
+// -----
+
+// Test unrolling of a 2D transfer_scatter with tensor semantics.
+
+func.func @transfer_scatter_unroll_tensor(
+ %dest: tensor<4096x64xf16>,
+ %vector: vector<4x64xf16>,
+ %indices: vector<4xindex>) -> tensor<4096x64xf16> {
+ %c0 = arith.constant 0 : index
+ %out = iree_vector_ext.transfer_scatter %vector into %dest[%c0, %c0]
+ [%indices : vector<4xindex>] {
+ indexing_maps = [affine_map<(d0, d1)[s0] -> (s0, d1)>,
+ affine_map<(d0, d1)[s0] -> (d0)>]
+ } : vector<4x64xf16>, tensor<4096x64xf16> -> tensor<4096x64xf16>
+ return %out : tensor<4096x64xf16>
+}
+
+// After unrolling, the 2D scatter becomes 4 rank-1 sub-scatters chained
+// via tensor SSA results.
+// CHECK-LABEL: func.func @transfer_scatter_unroll_tensor
+// CHECK-COUNT-4: transfer_scatter {{.+}} : vector<64xf16>
+
+// -----
+
+// Test unrolling of a 3D transfer_scatter with a transposed 2D index vector.
+// The first two output dims (d0=4, d1=8) are both scattered via a single
+// index vec of shape 8x4 (note: d1 before d0, i.e. "transposed").
+// The inner dim (d2=64) is contiguous.
+
+func.func @transfer_scatter_unroll_transposed_index(
+ %dest: memref<4096x64xf16>,
+ %vector: vector<4x8x64xf16>,
+ %indices: vector<8x4xindex>) {
+ %c0 = arith.constant 0 : index
+ iree_vector_ext.transfer_scatter %vector into %dest[%c0, %c0]
+ [%indices : vector<8x4xindex>] {
+ indexing_maps = [affine_map<(d0, d1, d2)[s0] -> (s0, d2)>,
+ affine_map<(d0, d1, d2)[s0] -> (d1, d0)>]
+ } : vector<4x8x64xf16>, memref<4096x64xf16>
+ return
+}
+
+// After two rounds of unrolling (d0=4 then d1=8), the 3D scatter
+// becomes 4*8=32 rank-1 sub-scatters.
+// CHECK-LABEL: func.func @transfer_scatter_unroll_transposed_index
+// CHECK-COUNT-32: transfer_scatter {{.+}} : vector<64xf16>
