[Mlir-commits] [mlir] [mlir][vector] Make gather/scatter index dimensions separately (PR #194395)
Krzysztof Drewniak
llvmlistbot at llvm.org
Mon Apr 27 07:57:21 PDT 2026
https://github.com/krzysz00 created https://github.com/llvm/llvm-project/pull/194395
This commit updates the semantics of vector.gather and vector.scatter into something that makes multi-dimensional memrefs/tensors easier to work with and resolves the inconsistencies highlighted by #187215 by allowing both possible semantics.
Previously, vector.gather / vector.scatter had one "index_vec" argument, which specified the offsets to gather at once you'd reached the start location given by base[offsets]. While it was agreed that these indices would be interpreted in terms of the memref's layout in that you'd apply the innermost memref stride, but it wasn't agreed if having indices that went past the dimension was meant to index into the underlying memory or into the next dimension of the subview/strided memref.
To resolve this ambiguity and get rid of a bunch of complex delinearization logic, vector.gather and vector.scatter now take 1 <= r <= [the nank of the base] index vectors which specify how we're gathering/scattering along the last r memref dimensions.
This means that you can manipulate the indexing into each memref dimension separately without needing to break up the indices with possibly dynamic values - for example, folding an expand_shape into a gather will now can look like just merging together the relevant gather indices.
This commit also clarifies the required relationship between the index_vecs and the strides of the memref of the index vecs use a short integer type - you must be able to truncate the relevant memref strides to match the index vec types.
This commit also updates the pattern that folds certain subviews into gathers to use the new semantics (and therefore support more strides), updates the lowerings to vector.loads along with the lowerings to LLVM to account for the new semantics, fixes some XeGPU patterns, and updates the one upstream creator of vector.gather to use the new form.
It also adds a canonicalization that moves splat indices from the index_vecs into the offsets.
>From f96c9ca1ea83a4bfd0ee71765ddf4fe2e5b943ad Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 6 Apr 2026 22:33:40 +0000
Subject: [PATCH] [mlir][vector] Make gather/scatter index dimensions
separately
This commit updates the semantics of vector.gather and vector.scatter
into something that makes multi-dimensional memrefs/tensors easier to
work with and resolves the inconsistencies highlighted by #187215 by
allowing both possible semantics.
Previously, vector.gather / vector.scatter had one "index_vec"
argument, which specified the offsets to gather at once you'd reached
the start location given by base[offsets]. While it was agreed that
these indices would be interpreted in terms of the memref's layout in
that you'd apply the innermost memref stride, but it wasn't agreed if
having indices that went past the dimension was meant to index into
the underlying memory or into the next dimension of the
subview/strided memref.
To resolve this ambiguity and get rid of a bunch of complex
delinearization logic, vector.gather and vector.scatter now take
1 <= r <= [the nank of the base] index vectors which specify how we're
gathering/scattering along the last r memref dimensions.
This means that you can manipulate the indexing into each memref
dimension separately without needing to break up the indices with
possibly dynamic values - for example, folding an expand_shape into
a gather will now can look like just merging together the relevant
gather indices.
This commit also clarifies the required relationship between the
index_vecs and the strides of the memref of the index vecs use a short
integer type - you must be able to truncate the relevant memref
strides to match the index vec types.
This commit also updates the pattern that folds certain subviews into
gathers to use the new semantics (and therefore support more strides),
updates the lowerings to vector.loads along with the lowerings to LLVM
to account for the new semantics, fixes some XeGPU patterns, and
updates the one upstream creator of vector.gather to use the new form.
It also adds a canonicalization that moves splat indices from the
index_vecs into the offsets.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 133 ++++++++----
.../VectorToLLVM/ConvertVectorToLLVM.cpp | 61 +++++-
.../VectorToXeGPU/VectorToXeGPU.cpp | 36 ++--
.../Linalg/Transforms/Vectorization.cpp | 55 ++---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 162 ++++++++++++++-
.../Vector/Transforms/LowerVectorGather.cpp | 189 ++++++++----------
.../Vector/Transforms/VectorUnroll.cpp | 10 +-
.../vector-to-llvm-interface.mlir | 2 +
.../vectorization/extract-with-patterns.mlir | 107 +++++-----
.../Dialect/Linalg/vectorization/extract.mlir | 55 ++---
mlir/test/Dialect/Vector/canonicalize.mlir | 39 ++++
mlir/test/Dialect/Vector/invalid.mlir | 72 +++++++
mlir/test/Dialect/Vector/ops.mlir | 20 ++
.../Vector/vector-gather-lowering.mlir | 184 ++++++++++++++---
.../Dialect/XeGPU/xegpu-vector-linearize.mlir | 11 +-
.../Dialect/Vector/CPU/gather.mlir | 31 ++-
16 files changed, 818 insertions(+), 349 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index fdde3995f6333..794ae192bb5a4 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2057,24 +2057,29 @@ def Vector_MaskedStoreOp :
];
}
+def Vector_AtLeastOneIntegerVec : Variadic<VectorOfNonZeroRankOf<[AnyInteger, Index]>> {
+ let minSize = 1;
+}
+
def Vector_GatherOp :
Vector_Op<"gather", [
DeclareOpInterfaceMethods<MaskableOpInterface>,
DeclareOpInterfaceMethods<MemorySpaceCastConsumerOpInterface>,
DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
- DeclareOpInterfaceMethods<AlignmentAttrOpInterface>
+ DeclareOpInterfaceMethods<AlignmentAttrOpInterface>,
+ AttrSizedOperandSegments
]>,
Arguments<(ins Arg<TensorOrMemRef<[AnyType]>, "", [MemRead]>:$base,
Variadic<Index>:$offsets,
- VectorOfNonZeroRankOf<[AnyInteger, Index]>:$indices,
+ Vector_AtLeastOneIntegerVec:$indices,
VectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$pass_thru,
OptionalAttr<IntValidAlignment<I64Attr>>: $alignment)>,
Results<(outs AnyVectorOfNonZeroRank:$result)> {
let summary = [{
- Gathers elements from memory or ranked tensor into a vector as defined by an
- index vector and a mask vector.
+ Gathers elements from memory or ranked tensor into a vector as defined by
+ index vectors and a mask vector.
}];
let description = [{
@@ -2083,17 +2088,18 @@ def Vector_GatherOp :
on the values of an n-D mask vector.
If a mask bit is set, the corresponding result element is taken from `base`
- at an index defined by k indices and n-D `index_vec`. Otherwise, the element
- is taken from the pass-through vector. As an example, suppose that `base` is
- 3-D and the result is 2-D:
+ at an index defined by k indices and the values in r (for gather rank) n-D
+ `index_vec`s. Otherwise, the element is taken from the pass-through vector.
+ As an example, suppose that `base` is 3-D and the result is 2-D,
+ where we have r=2 index vectors:
```mlir
func.func @gather_3D_to_2D(
%base: memref<?x10x?xf32>, %ofs_0: index, %ofs_1: index, %ofs_2: index,
- %indices: vector<2x3xi32>, %mask: vector<2x3xi1>,
- %fall_thru: vector<2x3xf32>) -> vector<2x3xf32> {
+ %indices0: vector<2x3xi32>, %indices1: vector<2x3xi32>,
+ %mask: vector<2x3xi1>, %fall_thru: vector<2x3xf32>) -> vector<2x3xf32> {
%result = vector.gather %base[%ofs_0, %ofs_1, %ofs_2]
- [%indices], %mask, %fall_thru : [...]
+ [%indices0, %indices1], %mask, %fall_thru : [...]
return %result : vector<2x3xf32>
}
```
@@ -2101,10 +2107,10 @@ def Vector_GatherOp :
The indexing semantics are then,
```
- result[i,j] := if mask[i,j] then base[i0, i1, i2 + indices[i,j]]
+ result[i,j] := if mask[i,j] then base[i0, i1 + indices0[i, j], i2 + indices1[i, j]]
else pass_thru[i,j]
```
- The index into `base` only varies in the innermost ((k-1)-th) dimension.
+ The index into `base` only varies in the r innermost ((k-r)-th) dimensions.
If a mask bit is set and the corresponding index is out-of-bounds for the
given base, the behavior is undefined. If a mask bit is not set, the value
@@ -2115,6 +2121,16 @@ def Vector_GatherOp :
during progressively lowering to bring other memory operations closer to
hardware ISA support for a gather.
+ If the type of the memref offsets is potentially larger than the type of the
+ gather indices, it is assumed that the final `r` strides of the memref can be
+ truncated to the index type such that the linear gather index can be computed
+ without signed overflow (i.e. the truncated stride sign-extends back to the
+ original value). If an extension to the offset type is needed, the gather
+ indices are interpreted as signed integers.
+
+ Similarly, if the gather indices are wider than the memref offset type,
+ they must be losslessly truncatable to that type.
+
An optional `alignment` attribute allows to specify the byte alignment of the
gather operation. It must be a positive power of 2. The operation must access
memory at an address aligned to this boundary. Violating this requirement
@@ -2127,37 +2143,43 @@ def Vector_GatherOp :
%0 = vector.gather %base[%c0][%v], %mask, %pass_thru
: memref<?xf32>, vector<2x16xi32>, vector<2x16xi1>, vector<2x16xf32> into vector<2x16xf32>
- // 2-D memref gathered to 1-D vector.
- %1 = vector.gather %base[%i, %j][%v], %mask, %pass_thru
- : memref<16x16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+ // 2-D memref gathered to 1-D vector along two dimensions. The single
+ // index type is shared across all index vectors.
+ %1 = vector.gather %base[%i, %j][%v0, %v1], %mask, %pass_thru
+ : memref<16x16xf32>, vector<16xi32>, vector<16xi1>,
+ vector<16xf32> into vector<16xf32>
```
}];
let extraClassDeclaration = [{
ShapedType getBaseType() { return getBase().getType(); }
- VectorType getIndexVectorType() { return getIndices().getType(); }
+ VectorType getIndexVectorType() {
+ return cast<VectorType>(getIndices().front().getType());
+ }
VectorType getMaskVectorType() { return getMask().getType(); }
VectorType getPassThruVectorType() { return getPassThru().getType(); }
VectorType getVectorType() { return getResult().getType(); }
}];
- let assemblyFormat =
- "$base `[` $offsets `]` `[` $indices `]` `,` "
- "$mask `,` $pass_thru attr-dict `:` type($base) `,` "
- "type($indices) `,` type($mask) `,` type($pass_thru) "
- "`into` type($result)";
+ let assemblyFormat = [{
+ $base `[` $offsets `]` `[` $indices `]` `,`
+ $mask `,` $pass_thru attr-dict `:` type($base) `,`
+ custom<SameTypeVariadicOperands>(ref($indices), type($indices)) `,`
+ type($mask) `,` type($pass_thru) `into` type($result)
+ }];
let hasCanonicalizer = 1;
let hasVerifier = 1;
let builders = [
OpBuilder<(ins "VectorType":$resultType,
"Value":$base,
- "ValueRange":$indices,
- "Value":$index_vec,
+ "ValueRange":$offsets,
+ "ValueRange":$index_vecs,
"Value":$mask,
"Value":$passthrough,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment), [{
- return build($_builder, $_state, resultType, base, indices, index_vec, mask, passthrough,
+ return build($_builder, $_state, resultType, base, offsets, index_vecs,
+ mask, passthrough,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>
@@ -2167,10 +2189,11 @@ def Vector_GatherOp :
def Vector_ScatterOp
: Vector_Op<"scatter",
[DeclareOpInterfaceMethods<MemorySpaceCastConsumerOpInterface>,
- DeclareOpInterfaceMethods<AlignmentAttrOpInterface>]>,
+ DeclareOpInterfaceMethods<AlignmentAttrOpInterface>,
+ AttrSizedOperandSegments]>,
Arguments<(ins Arg<TensorOrMemRef<[AnyType]>, "", [MemWrite]>:$base,
Variadic<Index>:$offsets,
- VectorOfNonZeroRankOf<[AnyInteger, Index]>:$indices,
+ Vector_AtLeastOneIntegerVec:$indices,
VectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$valueToStore,
OptionalAttr<IntValidAlignment<I64Attr>>:$alignment)>,
@@ -2183,12 +2206,20 @@ def Vector_ScatterOp
let description = [{
The scatter operation stores elements from a n-D vector into memory or ranked tensor as
- defined by a base with indices and an additional n-D index vector, but
- only if the corresponding bit in a n-D mask vector is set. Otherwise, no
- action is taken for that element. Informally the semantics are:
- ```
- if (mask[0]) base[index[0]] = value[0]
- if (mask[1]) base[index[1]] = value[1]
+ defined by a base with indices and an additional r (the scatter rank) n-D index
+ vectors, but only if the corresponding bit in a n-D mask vector is set.
+ Otherwise, no action is taken for that element. Informally the semantics are:
+ ```
+ if (mask[0, ..., 0]) {
+ base[offsets[0], ...,
+ offsets[k - r] + indices[0][0, 0, ... 0],
+ ...
+ offsets[k - 1] + indices[r - 1][0, 0, ... 0]] = value[0, 0 ..., 0]
+ }
+ if (mask[0, ..., 1]) {
+ base[offsets[0],
+ ...,
+ offsets[k - 1] + indices[r - 1][0, 0, ... 1]] = value[9, 0, ..., 1]
etc.
```
@@ -2208,6 +2239,16 @@ def Vector_ScatterOp
correspond to those of the `llvm.masked.scatter`
[intrinsic](https://llvm.org/docs/LangRef.html#llvm-masked-scatter-intrinsics).
+ If the type of the memref offsets is potentially larger than the type of the
+ scatter indices, it is assumed that the final `r` strides of the memref can be
+ truncated to the index type such that the linear scatter index can be computed
+ without signed overflow (i.e. the truncated stride sign-extends back to the
+ original value). If an extension to the offset type is needed, the scatter
+ indices are interpreted as signed integers.
+
+ Similarly, if the scatter indices are wider than the memref offset type,
+ they must be losslessly truncatable to that type.
+
An optional `alignment` attribute allows to specify the byte alignment of the
scatter operation. It must be a positive power of 2. The operation must access
memory at an address aligned to this boundary. Violating this requirement
@@ -2219,31 +2260,37 @@ def Vector_ScatterOp
vector.scatter %base[%c0][%v], %mask, %value
: memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
- vector.scatter %base[%i, %j][%v], %mask, %value
+ // The single index type is shared across all index vectors.
+ vector.scatter %base[%i, %j][%v0, %v1], %mask, %value
: memref<16x16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
```
}];
let extraClassDeclaration = [{
ShapedType getBaseType() { return getBase().getType(); }
- VectorType getIndexVectorType() { return getIndices().getType(); }
+ VectorType getIndexVectorType() {
+ return cast<VectorType>(getIndices().front().getType());
+ }
VectorType getMaskVectorType() { return getMask().getType(); }
VectorType getVectorType() { return getValueToStore().getType(); }
}];
- let assemblyFormat = "$base `[` $offsets `]` `[` $indices `]` `,` "
- "$mask `,` $valueToStore attr-dict `:` type($base) `,` "
- "type($indices) `,` type($mask) `,` "
- "type($valueToStore) (`->` type($result)^)?";
+ let assemblyFormat = [{
+ $base `[` $offsets `]` `[` $indices `]` `,`
+ $mask `,` $valueToStore attr-dict `:` type($base) `,`
+ custom<SameTypeVariadicOperands>(ref($indices), type($indices)) `,`
+ type($mask) `,` type($valueToStore) (`->` type($result)^)?
+ }];
let hasCanonicalizer = 1;
let hasVerifier = 1;
let builders = [OpBuilder<
- (ins "Type":$resultType, "Value":$base, "ValueRange":$indices,
- "Value":$index_vec, "Value":$mask, "Value":$valueToStore,
+ (ins "Type":$resultType, "Value":$base, "ValueRange":$offsets,
+ "ValueRange":$index_vecs, "Value":$mask, "Value":$valueToStore,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment),
[{
- return build($_builder, $_state, resultType, base, indices, index_vec, mask, valueToStore,
+ return build($_builder, $_state, resultType, base, offsets, index_vecs,
+ mask, valueToStore,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>];
@@ -2551,7 +2598,7 @@ def Vector_TypeCastOp :
}
def Vector_ConstantMaskOp :
- Vector_Op<"constant_mask", [Pure,
+ Vector_Op<"constant_mask", [Pure,
DeclareOpInterfaceMethods<VectorUnrollOpInterface>
]>,
Arguments<(ins DenseI64ArrayAttr:$mask_dim_sizes)>,
@@ -2611,7 +2658,7 @@ def Vector_ConstantMaskOp :
}
def Vector_CreateMaskOp :
- Vector_Op<"create_mask", [Pure,
+ Vector_Op<"create_mask", [Pure,
DeclareOpInterfaceMethods<VectorUnrollOpInterface>
]>,
Arguments<(ins Variadic<Index>:$mask_dim_sizes)>,
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index 43e0824fef6cd..87a2d79f08677 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -131,30 +131,75 @@ static LogicalResult isMemRefTypeSupported(MemRefType memRefType,
return success();
}
-// Add an index vector component to a base pointer.
+// Add a sequence of index vectors componentwise to a base pointer, using the
+// strides from the given memref. The index vectors are linearized:
+// index = sum(strides.take_back(len(indices)) * indices)
static Value getIndexedPtrs(ConversionPatternRewriter &rewriter, Location loc,
const LLVMTypeConverter &typeConverter,
MemRefType memRefType, Value llvmMemref, Value base,
- Value index, VectorType vectorType) {
+ ValueRange indices, VectorType vectorType) {
assert(succeeded(isMemRefTypeSupported(memRefType, typeConverter)) &&
"unsupported memref type");
assert(vectorType.getRank() == 1 && "expected a 1-d vector type");
- auto pType = MemRefDescriptor(llvmMemref).getElementPtrType();
+
+ unsigned rank = memRefType.getRank();
+ MemRefDescriptor desc(llvmMemref);
+
+ int64_t numElems = vectorType.getDimSize(0);
+ bool isScalable = vectorType.getScalableDims()[0];
+ SmallVector<int32_t> zeroMask(numElems, 0);
+ Value i32Zero = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
+ rewriter.getI32IntegerAttr(0));
+
+ Value linearized;
+ auto idxElemType =
+ cast<VectorType>(indices.front().getType()).getElementType();
+ auto idxVecType = LLVM::getVectorType(idxElemType, numElems, isScalable);
+ unsigned r = indices.size();
+ for (auto [idx, stride] : llvm::zip_equal(
+ indices, llvm::map_range(llvm::seq(rank - r, rank), [&](unsigned d) {
+ return desc.stride(rewriter, loc, d);
+ }))) {
+
+ Value castStride = stride;
+ if (stride.getType() != idxElemType) {
+ unsigned strideBits = cast<IntegerType>(stride.getType()).getWidth();
+ unsigned idxBits = cast<IntegerType>(idxElemType).getWidth();
+ if (strideBits > idxBits) {
+ castStride = LLVM::TruncOp::create(rewriter, loc, idxElemType, stride,
+ LLVM::IntegerOverflowFlags::nsw);
+ } else {
+ castStride = LLVM::SExtOp::create(rewriter, loc, idxElemType, stride);
+ }
+ }
+ Value strideVec = [&]() {
+ Value poison = LLVM::PoisonOp::create(rewriter, loc, idxVecType);
+ Value inserted = LLVM::InsertElementOp::create(rewriter, loc, poison,
+ castStride, i32Zero);
+ return LLVM::ShuffleVectorOp::create(rewriter, loc, inserted, poison,
+ zeroMask);
+ }();
+ Value contribution = LLVM::MulOp::create(rewriter, loc, idx, strideVec);
+ linearized = linearized ? LLVM::AddOp::create(rewriter, loc, linearized,
+ contribution)
+ : contribution;
+ }
+
+ auto pType = desc.getElementPtrType();
auto ptrsType =
LLVM::getVectorType(pType, vectorType.getDimSize(0),
/*isScalable=*/vectorType.getScalableDims()[0]);
return LLVM::GEPOp::create(
rewriter, loc, ptrsType,
- typeConverter.convertType(memRefType.getElementType()), base, index);
+ typeConverter.convertType(memRefType.getElementType()), base, linearized);
}
-/// Convert `foldResult` into a Value. Integer attribute is converted to
-/// an LLVM constant op.
+/// Convert `foldResult` into a Value, using `llvm.mlir.constant` if needed.
static Value getAsLLVMValue(OpBuilder &builder, Location loc,
OpFoldResult foldResult) {
if (auto attr = dyn_cast<Attribute>(foldResult)) {
- auto intAttr = cast<IntegerAttr>(attr);
- return LLVM::ConstantOp::create(builder, loc, intAttr).getResult();
+ return LLVM::ConstantOp::create(builder, loc, cast<TypedAttr>(attr))
+ .getResult();
}
return cast<Value>(foldResult);
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 0aae5cb5bb6ad..afcd4caa7a323 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -379,22 +379,34 @@ static Value computeOffsets(PatternRewriter &rewriter, OpType gatScatOp,
baseOffset =
arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
}
- Value indices = gatScatOp.getIndices();
- VectorType vecType = cast<VectorType>(indices.getType());
- Value strideVector =
- vector::BroadcastOp::create(rewriter, loc, vecType, strides.back())
- .getResult();
- Value stridedIndices =
- arith::MulIOp::create(rewriter, loc, strideVector, indices).getResult();
+ // The op's index vectors may use a narrower element type (e.g. i32), but
+ // strides and the base vector are always index — cast everything to the
+ // index-typed vector before combining.
+ OperandRange indexVecs = gatScatOp.getIndices();
+ VectorType vecType = gatScatOp.getIndexVectorType();
+ auto indexVecType =
+ VectorType::get(vecType.getShape(), rewriter.getIndexType());
+ Value combinedIndices = arith::ConstantOp::create(
+ rewriter, loc,
+ DenseElementsAttr::get(indexVecType, rewriter.getIndexAttr(0)));
+
+ auto tailStrides = ArrayRef<Value>(strides).take_back(indexVecs.size());
+ for (auto [idx, stride] : llvm::zip_equal(indexVecs, tailStrides)) {
+ Value castIdx = idx;
+ if (vecType != indexVecType)
+ castIdx = arith::IndexCastOp::create(rewriter, loc, indexVecType, idx);
+ Value strideVec =
+ vector::BroadcastOp::create(rewriter, loc, indexVecType, stride);
+ Value stridedIdx = arith::MulIOp::create(rewriter, loc, strideVec, castIdx);
+ combinedIndices =
+ arith::AddIOp::create(rewriter, loc, combinedIndices, stridedIdx);
+ }
Value baseVector =
- vector::BroadcastOp::create(
- rewriter, loc,
- VectorType::get(vecType.getShape(), rewriter.getIndexType()),
- baseOffset)
+ vector::BroadcastOp::create(rewriter, loc, indexVecType, baseOffset)
.getResult();
- return arith::AddIOp::create(rewriter, loc, baseVector, stridedIndices)
+ return arith::AddIOp::create(rewriter, loc, baseVector, combinedIndices)
.getResult();
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index b57e66a1c3580..a44ae49647f06 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -894,45 +894,21 @@ tensorExtractVectorizationPrecondition(Operation *op, bool vectorizeNDExtract) {
return success();
}
-/// Calculates the offsets (`$index_vec`) for `vector.gather` operations
-/// generated from `tensor.extract`. The offset is calculated as follows
-/// (example using scalar values):
-///
-/// offset = extractOp.indices[0]
-/// for (i = 1; i < numIndices; i++)
-/// offset = extractOp.dimSize[i] * offset + extractOp.indices[i];
-///
-/// For tensor<45 x 80 x 15 x f32> and index [1, 2, 3], this leads to:
-/// offset = ( ( 1 ) * 80 + 2 ) * 15 + 3
-static Value calculateGatherOffset(RewriterBase &rewriter,
- VectorizationState &state,
- tensor::ExtractOp extractOp,
- const IRMapping &bvm) {
- // The vector of indices for GatherOp should be shaped as the output vector.
+/// Calculates the per-dimension index vectors for `vector.gather` operations
+/// generated from `tensor.extract`. Returns one index vector per tensor
+/// dimension, each broadcast to the canonical vector shape.
+static SmallVector<Value> calculateGatherIndices(RewriterBase &rewriter,
+ VectorizationState &state,
+ tensor::ExtractOp extractOp,
+ const IRMapping &bvm) {
auto indexVecType = state.getCanonicalVecType(rewriter.getIndexType());
- auto loc = extractOp.getLoc();
-
- Value offset = broadcastIfNeeded(
- rewriter, bvm.lookup(extractOp.getIndices()[0]), indexVecType);
-
- const size_t numIndices = extractOp.getIndices().size();
- for (size_t i = 1; i < numIndices; i++) {
- Value dimIdx = arith::ConstantIndexOp::create(rewriter, loc, i);
- auto dimSize = broadcastIfNeeded(
- rewriter,
- tensor::DimOp::create(rewriter, loc, extractOp.getTensor(), dimIdx),
- indexVecType);
-
- offset = arith::MulIOp::create(rewriter, loc, offset, dimSize);
-
- auto extractOpIndex = broadcastIfNeeded(
- rewriter, bvm.lookup(extractOp.getIndices()[i]), indexVecType);
-
- offset = arith::AddIOp::create(rewriter, loc, extractOpIndex, offset);
- }
+ SmallVector<Value> indexVecs;
+ for (Value idx : extractOp.getIndices())
+ indexVecs.push_back(
+ broadcastIfNeeded(rewriter, bvm.lookup(idx), indexVecType));
- return offset;
+ return indexVecs;
}
enum VectorMemoryAccessKind { ScalarBroadcast, Contiguous, Gather };
@@ -1197,12 +1173,13 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
// 1. Handle gather access
if (memAccessKind == VectorMemoryAccessKind::Gather) {
- Value offset = calculateGatherOffset(rewriter, state, extractOp, bvm);
+ SmallVector<Value> indexVecs =
+ calculateGatherIndices(rewriter, state, extractOp, bvm);
// Generate the gather load
Operation *gatherOp = vector::GatherOp::create(
- rewriter, loc, resultType, extractOp.getTensor(), baseIndices, offset,
- maskConstantOp, passThruConstantOp);
+ rewriter, loc, resultType, extractOp.getTensor(), baseIndices,
+ indexVecs, maskConstantOp, passThruConstantOp);
gatherOp = state.maskOperation(rewriter, gatherOp, linalgOp);
LDBG() << "Vectorised as gather load: " << extractOp;
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 0ab6050ecc99b..3cc3733bb9500 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6361,20 +6361,56 @@ MaskedStoreOp::bubbleDownCasts(OpBuilder &builder) {
// GatherOp
//===----------------------------------------------------------------------===//
+/// Custom printer for a variadic operand list where all operands share the
+/// same type. Prints the type once regardless of operand count.
+static void printSameTypeVariadicOperands(OpAsmPrinter &p, Operation *op,
+ OperandRange operands,
+ TypeRange types) {
+ assert(!types.empty() && "expected at least one operand type");
+ p << types.front();
+}
+
+/// Custom parser for a variadic operand list where all operands share the
+/// same type. Parses a single type and replicates it for each operand.
+static ParseResult parseSameTypeVariadicOperands(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
+ SmallVectorImpl<Type> &types) {
+ Type type;
+ if (parser.parseType(type))
+ return failure();
+ types.assign(operands.size(), type);
+ return success();
+}
+
LogicalResult GatherOp::verify() {
- VectorType indVType = getIndexVectorType();
- VectorType maskVType = getMaskVectorType();
- VectorType resVType = getVectorType();
ShapedType baseType = getBaseType();
if (!llvm::isa<MemRefType, RankedTensorType>(baseType))
return emitOpError("requires base to be a memref or ranked tensor type");
+ if (getIndices().empty())
+ return emitOpError("requires at least one index vector");
+
+ VectorType indVType = getIndexVectorType();
+ VectorType maskVType = getMaskVectorType();
+ VectorType resVType = getVectorType();
+
if (failed(
verifyElementTypesMatch(*this, baseType, resVType, "base", "result")))
return failure();
if (llvm::size(getOffsets()) != baseType.getRank())
return emitOpError("requires ") << baseType.getRank() << " indices";
+
+ for (Value idx : getIndices()) {
+ if (cast<VectorType>(idx.getType()) != indVType)
+ return emitOpError("all index vectors must have the same type");
+ }
+
+ if (static_cast<int64_t>(getIndices().size()) > baseType.getRank())
+ return emitOpError("number of index vectors (")
+ << getIndices().size() << ") exceeds base rank ("
+ << baseType.getRank() << ")";
if (resVType.getShape() != indVType.getShape())
return emitOpError("expected result dim to match indices dim");
if (resVType.getShape() != maskVType.getShape())
@@ -6449,7 +6485,10 @@ class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
if (!isa<MemRefType>(op.getBase().getType()))
return rewriter.notifyMatchFailure(op, "base must be of memref type");
- if (failed(isZeroBasedContiguousSeq(op.getIndices())))
+ if (op.getIndices().size() != 1)
+ return rewriter.notifyMatchFailure(op, "expected single index vector");
+
+ if (failed(isZeroBasedContiguousSeq(op.getIndices().front())))
return failure();
rewriter.replaceOpWithNewOp<MaskedLoadOp>(op, op.getType(), op.getBase(),
@@ -6458,11 +6497,54 @@ class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
return success();
}
};
+
+/// Fold `gather %m[%o0, %o1, ..., %oK-R, ... %oK][splat(%i0), ..., %iR]`
+/// into `gather %m[%o0, %o1, ..., %oK-R + %i0, ..., %oK][%i1, ..., %iR]`.
+/// The 1-D case (one index vector) is intentionally skipped: dropping the
+/// only index vector would leave the gather without one, which is invalid.
+struct FoldBroadcastIndexIntoGatherOffset final : OpRewritePattern<GatherOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(GatherOp op,
+ PatternRewriter &rewriter) const override {
+ ValueRange indices = op.getIndices();
+ if (indices.size() < 2)
+ return rewriter.notifyMatchFailure(
+ op, "can't fold broadcast into 1-D gather");
+
+ auto bcast = indices.front().getDefiningOp<BroadcastOp>();
+ if (!bcast || !bcast.getSource().getType().isIntOrIndex())
+ return rewriter.notifyMatchFailure(op,
+ "first index isn't a splat scalar");
+
+ Value scalar = bcast.getSource();
+ Location loc = op.getLoc();
+
+ if (!scalar.getType().isIndex())
+ scalar = arith::IndexCastOp::create(rewriter, loc,
+ rewriter.getIndexType(), scalar);
+
+ unsigned k = op.getBaseType().getRank();
+ unsigned r = indices.size();
+ unsigned offsetIdx = k - r;
+
+ SmallVector<Value> newOffsets(op.getOffsets());
+ newOffsets[offsetIdx] =
+ arith::AddIOp::create(rewriter, loc, newOffsets[offsetIdx], scalar);
+
+ ValueRange newIndices = indices.drop_front();
+
+ rewriter.replaceOpWithNewOp<GatherOp>(
+ op, op.getVectorType(), op.getBase(), newOffsets, newIndices,
+ op.getMask(), op.getPassThru(), op.getAlignmentAttr());
+ return success();
+ }
+};
} // namespace
void GatherOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<GatherFolder, FoldContiguousGather>(context);
+ results.add<GatherFolder, FoldContiguousGather,
+ FoldBroadcastIndexIntoGatherOffset>(context);
}
FailureOr<std::optional<SmallVector<Value>>>
@@ -6476,19 +6558,33 @@ GatherOp::bubbleDownCasts(OpBuilder &builder) {
//===----------------------------------------------------------------------===//
LogicalResult ScatterOp::verify() {
- VectorType indVType = getIndexVectorType();
- VectorType maskVType = getMaskVectorType();
- VectorType valueVType = getVectorType();
ShapedType baseType = getBaseType();
if (!llvm::isa<MemRefType, RankedTensorType>(baseType))
return emitOpError("requires base to be a memref or ranked tensor type");
+ if (getIndices().empty())
+ return emitOpError("requires at least one index vector");
+
+ VectorType indVType = getIndexVectorType();
+ VectorType maskVType = getMaskVectorType();
+ VectorType valueVType = getVectorType();
+
if (failed(verifyElementTypesMatch(*this, baseType, valueVType, "base",
"valueToStore")))
return failure();
if (llvm::size(getOffsets()) != baseType.getRank())
return emitOpError("requires ") << baseType.getRank() << " indices";
+
+ for (Value idx : getIndices()) {
+ if (cast<VectorType>(idx.getType()) != indVType)
+ return emitOpError("all index vectors must have the same type");
+ }
+
+ if (static_cast<int64_t>(getIndices().size()) > baseType.getRank())
+ return emitOpError("number of index vectors (")
+ << getIndices().size() << ") exceeds base rank ("
+ << baseType.getRank() << ")";
if (valueVType.getShape() != indVType.getShape())
return emitOpError("expected valueToStore dim to match indices dim");
if (valueVType.getShape() != maskVType.getShape())
@@ -6541,7 +6637,10 @@ class FoldContiguousScatter final : public OpRewritePattern<ScatterOp> {
if (!isa<MemRefType>(op.getBase().getType()))
return failure();
- if (failed(isZeroBasedContiguousSeq(op.getIndices())))
+ if (op.getIndices().size() != 1)
+ return rewriter.notifyMatchFailure(op, "expected single index vector");
+
+ if (failed(isZeroBasedContiguousSeq(op.getIndices().front())))
return failure();
rewriter.replaceOpWithNewOp<MaskedStoreOp>(
@@ -6549,11 +6648,54 @@ class FoldContiguousScatter final : public OpRewritePattern<ScatterOp> {
return success();
}
};
+
+/// Fold `scatter %v, %m[%o0, ..., %oK-R, ... %oK][splat(%i0), ..., %iR]`
+/// into `scatter %v, %m[%o0, %o1, ..., %oK-R + %i0, ..., %oK][%i1, ..., %iR]`.
+/// The 1-D case (one index vector) is intentionally skipped: dropping the
+/// only index vector would leave the scatter without one, which is invalid.
+struct FoldBroadcastIndexIntoScatterOffset final : OpRewritePattern<ScatterOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(ScatterOp op,
+ PatternRewriter &rewriter) const override {
+ ValueRange indices = op.getIndices();
+ if (indices.size() < 2)
+ return rewriter.notifyMatchFailure(
+ op, "can't fold broadcast into 1-D scatter");
+
+ auto bcast = indices.front().getDefiningOp<BroadcastOp>();
+ if (!bcast || !bcast.getSource().getType().isIntOrIndex())
+ return rewriter.notifyMatchFailure(op,
+ "first index isn't a splat scalar");
+
+ Value scalar = bcast.getSource();
+ Location loc = op.getLoc();
+
+ if (!scalar.getType().isIndex())
+ scalar = arith::IndexCastOp::create(rewriter, loc,
+ rewriter.getIndexType(), scalar);
+
+ unsigned k = op.getBaseType().getRank();
+ unsigned r = indices.size();
+ unsigned offsetIdx = k - r;
+
+ SmallVector<Value> newOffsets(op.getOffsets());
+ newOffsets[offsetIdx] =
+ arith::AddIOp::create(rewriter, loc, newOffsets[offsetIdx], scalar);
+
+ ValueRange newIndices = indices.drop_front();
+
+ rewriter.replaceOpWithNewOp<ScatterOp>(
+ op, op->getResultTypes(), op.getBase(), newOffsets, newIndices,
+ op.getMask(), op.getValueToStore(), op.getAlignmentAttr());
+ return success();
+ }
+};
} // namespace
void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<ScatterFolder, FoldContiguousScatter>(context);
+ results.add<ScatterFolder, FoldContiguousScatter,
+ FoldBroadcastIndexIntoScatterOffset>(context);
}
FailureOr<std::optional<SmallVector<Value>>>
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
index 7194d41d60df7..b384a3fec8153 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
@@ -11,7 +11,6 @@
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -55,7 +54,7 @@ struct UnrollGather : OpRewritePattern<vector::GatherOp> {
LogicalResult matchAndRewrite(vector::GatherOp op,
PatternRewriter &rewriter) const override {
- Value indexVec = op.getIndices();
+ OperandRange indexVecs = op.getIndices();
Value maskVec = op.getMask();
Value passThruVec = op.getPassThru();
@@ -63,14 +62,16 @@ struct UnrollGather : OpRewritePattern<vector::GatherOp> {
VectorType subTy, int64_t index) {
int64_t thisIdx[1] = {index};
- Value indexSubVec =
- vector::ExtractOp::create(rewriter, loc, indexVec, thisIdx);
+ SmallVector<Value> indexSubVecs;
+ for (Value iv : indexVecs)
+ indexSubVecs.push_back(
+ vector::ExtractOp::create(rewriter, loc, iv, thisIdx));
Value maskSubVec =
vector::ExtractOp::create(rewriter, loc, maskVec, thisIdx);
Value passThruSubVec =
vector::ExtractOp::create(rewriter, loc, passThruVec, thisIdx);
return vector::GatherOp::create(rewriter, loc, subTy, op.getBase(),
- op.getOffsets(), indexSubVec, maskSubVec,
+ op.getOffsets(), indexSubVecs, maskSubVec,
passThruSubVec, op.getAlignmentAttr());
};
@@ -78,22 +79,23 @@ struct UnrollGather : OpRewritePattern<vector::GatherOp> {
}
};
-/// Rewrites a vector.gather of a strided MemRef as a gather of a non-strided
-/// MemRef with updated indices that model the strided access.
+/// Rewrites a vector.gather of a strided MemRef (from a subview) as a gather
+/// of the original non-strided source MemRef with multi-dimensional indices.
+///
+/// The original single index vector is used for the outer dimension, and a
+/// zero vector is added for the trailing dimension.
///
/// ```mlir
/// %subview = memref.subview %M (...)
/// : memref<100x3xf32> to memref<100xf32, strided<[3]>>
-/// %gather = vector.gather %subview[%idxs] (...)
+/// %gather = vector.gather %subview[%c0][%idxs] (...)
/// : memref<100xf32, strided<[3]>>
/// ```
/// ==>
/// ```mlir
-/// %collapse_shape = memref.collapse_shape %M (...)
-/// : memref<100x3xf32> into memref<300xf32>
-/// %new_idxs = arith.muli %idxs, %c3 : vector<4xindex>
-/// %gather = vector.gather %collapse_shape[%new_idxs] (...)
-/// : memref<300xf32> (...)
+/// %zeros = arith.constant dense<0> : vector<4xindex>
+/// %gather = vector.gather %M[%c0, %c0][%idxs, %zeros] (...)
+/// : memref<100x3xf32> (...)
/// ```
///
/// ATM this is effectively limited to reading a 1D Vector from a 2D MemRef,
@@ -103,58 +105,59 @@ struct RemoveStrideFromGatherSource : OpRewritePattern<vector::GatherOp> {
LogicalResult matchAndRewrite(vector::GatherOp op,
PatternRewriter &rewriter) const override {
- Value base = op.getBase();
+ if (op.getIndices().size() != 1)
+ return rewriter.notifyMatchFailure(op, "expected single index vector");
- // TODO: Strided accesses might be coming from other ops as well
- auto subview = base.getDefiningOp<memref::SubViewOp>();
+ auto subview = op.getBase().getDefiningOp<memref::SubViewOp>();
if (!subview)
- return failure();
-
- auto sourceType = subview.getSource().getType();
-
- // TODO: Allow ranks > 2.
- if (sourceType.getRank() != 2)
- return failure();
-
- // Get strides
- auto layout = subview.getResult().getType().getLayout();
- auto stridedLayoutAttr = llvm::dyn_cast<StridedLayoutAttr>(layout);
- if (!stridedLayoutAttr)
- return failure();
-
- // TODO: Allow the access to be strided in multiple dimensions.
- if (stridedLayoutAttr.getStrides().size() != 1)
- return failure();
-
- int64_t srcTrailingDim = sourceType.getShape().back();
+ return rewriter.notifyMatchFailure(op, "base is not a memref.subview");
+
+ // Restricted to a 1-D subview of a 2-D source where the trailing source
+ // dim was rank-reduced (so the subview iterates rows of the source).
+ // TODO: Generalize.
+ if (subview.getSource().getType().getRank() != 2 ||
+ subview.getResult().getType().getRank() != 1 ||
+ !subview.getDroppedDims().test(1))
+ return rewriter.notifyMatchFailure(
+ op, "expected 1-D subview rank-reducing the trailing dim of a 2-D "
+ "source");
- // Assume that the stride matches the trailing dimension of the source
- // memref.
- // TODO: Relax this assumption.
- if (stridedLayoutAttr.getStrides()[0] != srcTrailingDim)
- return failure();
-
- // 1. Collapse the input memref so that it's "flat".
- SmallVector<ReassociationIndices> reassoc = {{0, 1}};
- Value collapsed = memref::CollapseShapeOp::create(
- rewriter, op.getLoc(), subview.getSource(), reassoc);
-
- // 2. Generate new gather indices that will model the
- // strided access.
- IntegerAttr stride = rewriter.getIndexAttr(srcTrailingDim);
- VectorType vType = op.getIndices().getType();
- Value mulCst = arith::ConstantOp::create(
- rewriter, op.getLoc(), vType, DenseElementsAttr::get(vType, stride));
+ Location loc = op.getLoc();
- Value newIdxs =
- arith::MulIOp::create(rewriter, op.getLoc(), op.getIndices(), mulCst);
+ Value outerStride = getValueOrCreateConstantIndexOp(
+ rewriter, loc, subview.getMixedStrides()[0]);
+
+ SmallVector<OpFoldResult> subviewOffsets = subview.getMixedOffsets();
+ Value subviewOuterOff =
+ getValueOrCreateConstantIndexOp(rewriter, loc, subviewOffsets[0]);
+ Value subviewInnerOff =
+ getValueOrCreateConstantIndexOp(rewriter, loc, subviewOffsets[1]);
+ Value scaledGatherOff = arith::MulIOp::create(
+ rewriter, loc, op.getOffsets().front(), outerStride);
+ Value newOuterOff =
+ arith::AddIOp::create(rewriter, loc, subviewOuterOff, scaledGatherOff);
+
+ // Scale the gather indices by the outer stride for the same reason. Cast
+ // the scalar stride to the index vector's element type before broadcasting.
+ VectorType vType = op.getIndexVectorType();
+ Type eltTy = vType.getElementType();
+ Value strideForVec = outerStride;
+ if (eltTy != rewriter.getIndexType())
+ strideForVec =
+ arith::IndexCastOp::create(rewriter, loc, eltTy, outerStride);
+ Value strideVec =
+ vector::BroadcastOp::create(rewriter, loc, vType, strideForVec);
+ Value newOuterIndices = arith::MulIOp::create(rewriter, loc, strideVec,
+ op.getIndices().front());
+
+ Value zeroVec = arith::ConstantOp::create(rewriter, loc, vType,
+ rewriter.getZeroAttr(vType));
- // 3. Create an updated gather op with the collapsed input memref and the
- // updated indices.
Value newGather = vector::GatherOp::create(
- rewriter, op.getLoc(), op.getResult().getType(), collapsed,
- op.getOffsets(), newIdxs, op.getMask(), op.getPassThru(),
- op.getAlignmentAttr());
+ rewriter, loc, op.getResult().getType(), subview.getSource(),
+ SmallVector<Value>{newOuterOff, subviewInnerOff},
+ SmallVector<Value>{newOuterIndices, zeroVec}, op.getMask(),
+ op.getPassThru(), op.getAlignmentAttr());
rewriter.replaceOp(op, newGather);
return success();
@@ -165,12 +168,9 @@ struct RemoveStrideFromGatherSource : OpRewritePattern<vector::GatherOp> {
/// `tensor.extract`s. To avoid out-of-bounds memory accesses, these
/// loads/extracts are made conditional using `scf.if` ops.
///
-/// For multi-dimensional memrefs (rank > 1), the gather index is combined
-/// with the offsets via linearize-then-delinearize to produce correct
-/// N-D load indices:
-/// idx = indices[i]
-/// flatIdx = linearize(offsets, memrefShape) + idx
-/// loadIndices = delinearize(flatIdx, memrefShape)
+/// With r index vectors, each one directly offsets one of the r innermost
+/// base dimensions. The load indices are:
+/// loadOffsets[k-r+j] = offsets[k-r+j] + indexCast(indices[j][i])
struct Gather1DToConditionalLoads : OpRewritePattern<vector::GatherOp> {
using Base::Base;
@@ -191,9 +191,6 @@ struct Gather1DToConditionalLoads : OpRewritePattern<vector::GatherOp> {
Value condMask = op.getMask();
Value base = op.getBase();
- // For multi-dimensional memrefs, use linearize+delinearize to compute
- // correct N-D load indices from the 1-D gather index.
- bool useDelinearization = false;
if (auto memType = dyn_cast<MemRefType>(base.getType())) {
// vector.load requires the most minor memref dim to have unit stride
// (unless reading exactly 1 element).
@@ -204,26 +201,27 @@ struct Gather1DToConditionalLoads : OpRewritePattern<vector::GatherOp> {
return rewriter.notifyMatchFailure(
op, "most minor memref dim must have unit stride");
}
+ }
- if (memType.getRank() > 1)
- useDelinearization = true;
+ unsigned r = op.getIndices().size();
+ VectorType indexVecType =
+ op.getIndexVectorType().clone(rewriter.getIndexType());
+ SmallVector<Value> indexVecs;
+ for (Value iv : op.getIndices()) {
+ if (iv.getType() == indexVecType)
+ indexVecs.push_back(iv);
+ else
+ indexVecs.push_back(
+ rewriter.createOrFold<arith::IndexCastOp>(loc, indexVecType, iv));
}
- Value indexVec = rewriter.createOrFold<arith::IndexCastOp>(
- loc, op.getIndexVectorType().clone(rewriter.getIndexType()),
- op.getIndices());
+ // Snapshot the offsets so per-element rewrites of the trailing r entries
+ // (loadOffsets[rank-r+j] += indices[j][i]) start from the originals each
+ // iteration, not from the previous iteration's sum.
auto loadOffsets = llvm::to_vector(op.getOffsets());
- Value lastLoadOffset = loadOffsets.back();
-
- // Compute the memref shape and linearized offsets once, outside the
- // per-element loop.
- SmallVector<OpFoldResult> baseShape;
- Value linearizedOffsets;
- if (useDelinearization) {
- baseShape = memref::getMixedSizes(rewriter, loc, base);
- linearizedOffsets = affine::AffineLinearizeIndexOp::create(
- rewriter, loc, loadOffsets, baseShape, /*disjoint=*/false);
- }
+ unsigned rank = loadOffsets.size();
+ SmallVector<Value> savedOffsets =
+ llvm::to_vector(ArrayRef<Value>(loadOffsets).take_back(r));
Value result = op.getPassThru();
BoolAttr nontemporalAttr = nullptr;
@@ -234,23 +232,12 @@ struct Gather1DToConditionalLoads : OpRewritePattern<vector::GatherOp> {
int64_t thisIdx[1] = {i};
Value condition =
vector::ExtractOp::create(rewriter, loc, condMask, thisIdx);
- Value index = vector::ExtractOp::create(rewriter, loc, indexVec, thisIdx);
-
- if (useDelinearization) {
- // The gather index offsets the innermost dimension. Combine with
- // the offsets by linearizing, adding the gather index, then
- // delinearizing back to N-D indices:
- // flatIdx = linearize(offsets, shape) + idx
- // loadIndices = delinearize(flatIdx, shape)
- Value flatIdx =
- rewriter.createOrFold<arith::AddIOp>(loc, linearizedOffsets, index);
- auto delinOp = affine::AffineDelinearizeIndexOp::create(
- rewriter, loc, flatIdx, baseShape, /*hasOuterBound=*/true);
- for (int64_t d = 0, rank = loadOffsets.size(); d < rank; ++d)
- loadOffsets[d] = delinOp.getResult(d);
- } else {
- loadOffsets.back() =
- rewriter.createOrFold<arith::AddIOp>(loc, lastLoadOffset, index);
+
+ for (unsigned j = 0; j < r; ++j) {
+ Value index =
+ vector::ExtractOp::create(rewriter, loc, indexVecs[j], thisIdx);
+ loadOffsets[rank - r + j] =
+ rewriter.createOrFold<arith::AddIOp>(loc, savedOffsets[j], index);
}
auto loadBuilder = [&](OpBuilder &b, Location loc) {
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index ec08f01d2a4b9..095e4594b337a 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -683,8 +683,12 @@ struct UnrollGatherPattern : public OpRewritePattern<vector::GatherOp> {
// To get the unrolled gather, extract the same slice based on the
// decomposed shape from each of the index, mask, and pass-through
// vectors.
- Value indexSubVec = rewriter.createOrFold<vector::ExtractStridedSliceOp>(
- loc, gatherOp.getIndices(), elementOffsets, *targetShape, strides);
+ SmallVector<Value> indexSubVecs;
+ for (Value indexVec : gatherOp.getIndices()) {
+ indexSubVecs.push_back(
+ rewriter.createOrFold<vector::ExtractStridedSliceOp>(
+ loc, indexVec, elementOffsets, *targetShape, strides));
+ }
Value maskSubVec = rewriter.createOrFold<vector::ExtractStridedSliceOp>(
loc, gatherOp.getMask(), elementOffsets, *targetShape, strides);
Value passThruSubVec =
@@ -693,7 +697,7 @@ struct UnrollGatherPattern : public OpRewritePattern<vector::GatherOp> {
strides);
auto slicedGather = vector::GatherOp::create(
rewriter, loc, targetType, gatherOp.getBase(), gatherOp.getOffsets(),
- indexSubVec, maskSubVec, passThruSubVec);
+ indexSubVecs, maskSubVec, passThruSubVec);
result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
loc, slicedGather, result, elementOffsets, strides);
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
index d570d46e11b4a..1222d13afdda0 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
@@ -2022,6 +2022,8 @@ func.func @gather_1d_from_2d(%arg0: memref<4x4xf32>, %arg1: vector<4xi32>, %arg2
// CHECK-LABEL: func @gather_1d_from_2d
// CHECK: %[[B:.*]] = llvm.getelementptr %{{.*}}[%{{.*}}] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+// Stride is truncated from i64 to i32 (index vector elem type) with nsw.
+// CHECK: llvm.trunc %{{.*}} overflow<nsw> : i64 to i32
// CHECK: %[[P:.*]] = llvm.getelementptr %[[B]][%{{.*}}] : (!llvm.ptr, vector<4xi32>) -> vector<4x!llvm.ptr>, f32
// CHECK: %[[G:.*]] = llvm.intr.masked.gather %[[P]], %{{.*}}, %{{.*}} {alignment = 4 : i32} : (vector<4x!llvm.ptr>, vector<4xi1>, vector<4xf32>) -> vector<4xf32>
// CHECK: return %[[G]] : vector<4xf32>
diff --git a/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
index e04a3f1a83d35..1e5d0844272ad 100644
--- a/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
@@ -230,18 +230,17 @@ func.func @vectorize_nd_tensor_extract_index_from_tensor(%arg0: tensor<3x3xf32>,
// CHECK-SAME: %[[ARG4:.*]]: tensor<4x7x3x2xf32>
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[PV:.*]] = ub.poison : i32
-// CHECK-DAG: %[[CST:.*]] = arith.constant dense<3> : vector<4x3xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<4x7x3x2xi1>
// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<4x7x3x2xf32>
// CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], %[[PV]] {in_bounds = [true, true]} : tensor<4x3xi32>, vector<4x3xi32>
// CHECK: %[[V1:.*]] = vector.transfer_read %[[ARG2]][%[[C0]], %[[C0]]], %[[PV]] {in_bounds = [true, true]} : tensor<4x3xi32>, vector<4x3xi32>
// CHECK: %[[CAST:.*]] = arith.index_cast %[[V0]] : vector<4x3xi32> to vector<4x3xindex>
+// CHECK: %[[B0:.*]] = vector.broadcast %[[CAST]] : vector<4x3xindex> to vector<7x2x4x3xindex>
+// CHECK: %[[T0:.*]] = vector.transpose %[[B0]], [2, 0, 3, 1] : vector<7x2x4x3xindex> to vector<4x7x3x2xindex>
// CHECK: %[[CAST_1:.*]] = arith.index_cast %[[V1]] : vector<4x3xi32> to vector<4x3xindex>
-// CHECK: %[[MULI:.*]] = arith.muli %[[CAST]], %[[CST]] : vector<4x3xindex>
-// CHECK: %[[ADDI:.*]] = arith.addi %[[CAST_1]], %[[MULI]] : vector<4x3xindex>
-// CHECK: %[[B:.*]] = vector.broadcast %[[ADDI]] : vector<4x3xindex> to vector<7x2x4x3xindex>
-// CHECK: %[[T:.*]] = vector.transpose %[[B]], [2, 0, 3, 1] : vector<7x2x4x3xindex> to vector<4x7x3x2xindex>
-// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]]] [%[[T]]], %[[CST_1]], %[[PASSTHRU]] : tensor<3x3xf32>, vector<4x7x3x2xindex>, vector<4x7x3x2xi1>, vector<4x7x3x2xf32> into vector<4x7x3x2xf32>
+// CHECK: %[[B1:.*]] = vector.broadcast %[[CAST_1]] : vector<4x3xindex> to vector<7x2x4x3xindex>
+// CHECK: %[[T1:.*]] = vector.transpose %[[B1]], [2, 0, 3, 1] : vector<7x2x4x3xindex> to vector<4x7x3x2xindex>
+// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]]] [%[[T0]], %[[T1]]], %[[CST_1]], %[[PASSTHRU]] : tensor<3x3xf32>, vector<4x7x3x2xindex>, vector<4x7x3x2xi1>, vector<4x7x3x2xf32> into vector<4x7x3x2xf32>
// CHECK: vector.transfer_write %[[GATHER]], %[[ARG4]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true, true]} : vector<4x7x3x2xf32>, tensor<4x7x3x2xf32>
// -----
@@ -268,17 +267,17 @@ func.func @vectorize_nd_tensor_extract_load_1d_column_vector_using_gather_load(%
// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_load_1d_column_vector_using_gather_load
// CHECK-SAME: %[[ARG0:.*]]: tensor<8x128x768xf32>
// CHECK-SAME: %[[ARG1:.*]]: index
+// CHECK-DAG: %[[CST_ZEROS:.*]] = arith.constant dense<0> : vector<8x1xindex>
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<true> : vector<8x1xi1>
-// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<[0, 98304, 196608, 294912, 393216, 491520, 589824, 688128]> : vector<8xindex>
+// CHECK-DAG: %[[CST_SEQ:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<8x1xf32>
-// CHECK: %[[ADDI_ARG1:.*]] = arith.addi %[[ARG1]], %[[ARG1]] : index
-// CHECK: %[[B1:.*]] = vector.broadcast %[[CST_1]] : vector<8xindex> to vector<1x8xindex>
+// CHECK: %[[B1:.*]] = vector.broadcast %[[CST_SEQ]] : vector<8xindex> to vector<1x8xindex>
// CHECK: %[[T:.*]] = vector.transpose %[[B1]], [1, 0] : vector<1x8xindex> to vector<8x1xindex>
+// CHECK: %[[ADDI_ARG1:.*]] = arith.addi %[[ARG1]], %[[ARG1]] : index
// CHECK: %[[B2:.*]] = vector.broadcast %[[ADDI_ARG1]] : index to vector<8x1xindex>
-// CHECK: %[[ADDI:.*]] = arith.addi %[[B2]], %[[T]] : vector<8x1xindex>
-// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[ADDI]]], %[[CST_0]], %[[PASSTHRU]] : tensor<8x128x768xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>
+// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[T]], %[[CST_ZEROS]], %[[B2]]], %[[CST_0]], %[[PASSTHRU]] : tensor<8x128x768xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>
// CHECK: vector.transfer_write %[[GATHER]], %[[EMPTY]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>
// -----
@@ -304,14 +303,15 @@ func.func @index_from_output_column_vector_gather_load(%src: tensor<8x128xf32>)
// CHECK-LABEL: func.func @index_from_output_column_vector_gather_load(
// CHECK-SAME: %[[SRC:.*]]: tensor<8x128xf32>) -> tensor<8x1xf32> {
-// CHECK: %[[IDX_VEC:.*]] = arith.constant dense<[0, 128, 256, 384, 512, 640, 768, 896]> : vector<8xindex>
-// CHECK: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8x1xi1>
+// CHECK-DAG: %[[CST_ZEROS:.*]] = arith.constant dense<0> : vector<8x1xindex>
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>
+// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<8x1xi1>
+// CHECK-DAG: %[[IDX_VEC:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
// CHECK: %[[OUT:.*]] = tensor.empty() : tensor<8x1xf32>
// CHECK: %[[B:.*]] = vector.broadcast %[[IDX_VEC]] : vector<8xindex> to vector<1x8xindex>
// CHECK: %[[TR:.*]] = vector.transpose %[[B]], [1, 0] : vector<1x8xindex> to vector<8x1xindex>
-// CHECK: %[[GATHER:.*]] = vector.gather %[[SRC]]{{\[}}%[[C0]], %[[C0]]] {{\[}}%[[TR]]], %[[MASK]], %[[PASS_THRU]] : tensor<8x128xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>
+// CHECK: %[[GATHER:.*]] = vector.gather %[[SRC]]{{\[}}%[[C0]], %[[C0]]] {{\[}}%[[TR]], %[[CST_ZEROS]]], %[[MASK]], %[[PASS_THRU]] : tensor<8x128xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>
// CHECK: %[[RES:.*]] = vector.transfer_write %[[GATHER]], %[[OUT]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>
// CHECK: return %[[RES]] : tensor<8x1xf32>
@@ -340,14 +340,15 @@ func.func @index_from_output_column_vector_contiguous_load(%src: tensor<8x128xf3
// CHECK-LABEL: func.func @index_from_output_column_vector_contiguous_load(
// CHECK-SAME: %[[SRC:.*]]: tensor<8x128xf32>) -> tensor<8x1xf32> {
-// CHECK: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8x1xi1>
-// CHECK: %[[IDX_VEC:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
+// CHECK-DAG: %[[CST_ZEROS:.*]] = arith.constant dense<0> : vector<8x1xindex>
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>
+// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<8x1xi1>
+// CHECK-DAG: %[[IDX_VEC:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
// CHECK: %[[OUT:.*]] = tensor.empty() : tensor<8x1xf32>
// CHECK: %[[B:.*]] = vector.broadcast %[[IDX_VEC]] : vector<8xindex> to vector<1x8xindex>
// CHECK: %[[TR:.*]] = vector.transpose %[[B]], [1, 0] : vector<1x8xindex> to vector<8x1xindex>
-// CHECK: %[[GATHER:.*]] = vector.gather %[[SRC]]{{\[}}%[[C0]], %[[C0]]] {{\[}}%[[TR]]], %[[MASK]], %[[PASS_THRU]] : tensor<8x128xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>
+// CHECK: %[[GATHER:.*]] = vector.gather %[[SRC]]{{\[}}%[[C0]], %[[C0]]] {{\[}}%[[CST_ZEROS]], %[[TR]]], %[[MASK]], %[[PASS_THRU]] : tensor<8x128xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>
// CHECK: %[[RES:.*]] = vector.transfer_write %[[GATHER]], %[[OUT]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>
// CHECK: return %[[RES]] : tensor<8x1xf32>
@@ -409,19 +410,17 @@ func.func @vectorize_nd_tensor_extract_with_affine_apply_gather(%6: tensor<80x16
// CHECK-SAME: %[[VAL_0:.*]]: tensor<80x16xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: index,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {
+// CHECK-DAG: %[[CST_16:.*]] = arith.constant dense<16> : vector<1x4xindex>
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<1x4xi1>
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32>
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
-// CHECK-DAG: %[[VAL_7:.*]] = arith.constant dense<16> : vector<4xindex>
-// CHECK: %[[VAL_8:.*]] = vector.broadcast %[[VAL_1]] : index to vector<4xindex>
-// CHECK: %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_3]] : vector<4xindex>
-// CHECK: %[[VAL_10:.*]] = arith.muli %[[VAL_9]], %[[VAL_7]] : vector<4xindex>
-// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_7]] : vector<4xindex>
-// CHECK: %[[VAL_12:.*]] = vector.broadcast %[[VAL_11]] : vector<4xindex> to vector<1x4xindex>
-// CHECK: %[[VAL_13:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {{\[}}%[[VAL_12]]], %[[VAL_4]], %[[VAL_5]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
-// CHECK: %[[VAL_14:.*]] = vector.transfer_write %[[VAL_13]], %[[VAL_2]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
-// CHECK: return %[[VAL_14]] : tensor<1x4xf32>
+// CHECK: %[[VAL_7:.*]] = vector.broadcast %[[VAL_1]] : index to vector<4xindex>
+// CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_3]] : vector<4xindex>
+// CHECK: %[[VAL_9:.*]] = vector.broadcast %[[VAL_8]] : vector<4xindex> to vector<1x4xindex>
+// CHECK: %[[VAL_10:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {{\[}}%[[VAL_9]], %[[CST_16]]], %[[VAL_4]], %[[VAL_5]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
+// CHECK: %[[VAL_11:.*]] = vector.transfer_write %[[VAL_10]], %[[VAL_2]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
+// CHECK: return %[[VAL_11]] : tensor<1x4xf32>
// CHECK: }
// Make sure that non-linear arithmetic operations (e.g. arith.maxsi) are allowed when calculating indices for load operations. Gather load.
@@ -443,14 +442,15 @@ func.func @vectorize_nd_tensor_extract_with_maxsi_gather(%arg0: tensor<80x16xf32
// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_with_maxsi_gather(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<80x16xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {
-// CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[1264, 1265, 1266, 1267]> : vector<4xindex>
-// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<1x4xi1>
-// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32>
-// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_7:.*]] = vector.broadcast %[[VAL_2]] : vector<4xindex> to vector<1x4xindex>
-// CHECK: %[[VAL_9:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {{\[}}%[[VAL_7]]], %[[VAL_4]], %[[VAL_5]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
-// CHECK: %[[VAL_10:.*]] = vector.transfer_write %[[VAL_9]], %[[VAL_1]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
-// CHECK: return %[[VAL_10]] : tensor<1x4xf32>
+// CHECK-DAG: %[[STEP:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
+// CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<true> : vector<1x4xi1>
+// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32>
+// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[CST_79:.*]] = arith.constant dense<79> : vector<1x4xindex>
+// CHECK: %[[B_STEP:.*]] = vector.broadcast %[[STEP]] : vector<4xindex> to vector<1x4xindex>
+// CHECK: %[[VAL_5:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_4]], %[[VAL_4]]] {{\[}}%[[CST_79]], %[[B_STEP]]], %[[VAL_2]], %[[VAL_3]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
+// CHECK: %[[VAL_6:.*]] = vector.transfer_write %[[VAL_5]], %[[VAL_1]]{{\[}}%[[VAL_4]], %[[VAL_4]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
+// CHECK: return %[[VAL_6]] : tensor<1x4xf32>
// CHECK: }
// -----
@@ -471,19 +471,16 @@ func.func @vectorize_nd_tensor_extract_block_arg(%arg0: tensor<5x6xf32>, %arg1:
// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_block_arg(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<5x6xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<5xindex>) -> tensor<5xf32> {
+// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<5xf32>
+// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<5xi1>
+// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<[0, 1, 2, 3, 4]> : vector<5xindex>
// CHECK-DAG: %[[PAD:.*]] = ub.poison : index
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index
-// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<[0, 1, 2, 3, 4]> : vector<5xindex>
-// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<5xi1>
-// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<5xf32>
-// CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<6> : vector<5xindex>
-// CHECK: %[[VAL_7:.*]] = tensor.empty() : tensor<5xf32>
-// CHECK: %[[VAL_8:.*]] = vector.transfer_read %[[VAL_1]]{{\[}}%[[VAL_2]]], %[[PAD]] {in_bounds = [true]} : tensor<5xindex>, vector<5xindex>
-// CHECK: %[[VAL_9:.*]] = arith.muli %[[VAL_8]], %[[VAL_6]] : vector<5xindex>
-// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_3]] : vector<5xindex>
-// CHECK: %[[VAL_11:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_2]], %[[VAL_2]]] {{\[}}%[[VAL_10]]], %[[VAL_4]], %[[VAL_5]] : tensor<5x6xf32>, vector<5xindex>, vector<5xi1>, vector<5xf32> into vector<5xf32>
-// CHECK: %[[VAL_12:.*]] = vector.transfer_write %[[VAL_11]], %[[VAL_7]]{{\[}}%[[VAL_2]]] {in_bounds = [true]} : vector<5xf32>, tensor<5xf32>
-// CHECK: return %[[VAL_12]] : tensor<5xf32>
+// CHECK: %[[VAL_6:.*]] = tensor.empty() : tensor<5xf32>
+// CHECK: %[[VAL_7:.*]] = vector.transfer_read %[[VAL_1]]{{\[}}%[[VAL_2]]], %[[PAD]] {in_bounds = [true]} : tensor<5xindex>, vector<5xindex>
+// CHECK: %[[VAL_8:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_2]], %[[VAL_2]]] {{\[}}%[[VAL_7]], %[[VAL_3]]], %[[VAL_4]], %[[VAL_5]] : tensor<5x6xf32>, vector<5xindex>, vector<5xi1>, vector<5xf32> into vector<5xf32>
+// CHECK: %[[VAL_9:.*]] = vector.transfer_write %[[VAL_8]], %[[VAL_6]]{{\[}}%[[VAL_2]]] {in_bounds = [true]} : vector<5xf32>, tensor<5xf32>
+// CHECK: return %[[VAL_9]] : tensor<5xf32>
// CHECK: }
// -----
@@ -510,16 +507,14 @@ func.func @vectorize_reverse_like_tensor_extract(%arg0: tensor<1x2x3xf32>, %arg1
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]*]]
-// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[CST_ZEROS:.+]] = arith.constant dense<0> : vector<1x1x3xindex>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<1x1x3xi1>
// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<1x1x3xf32>
+// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<1x1x3xi1>
// CHECK-DAG: %[[INIT_IDX:.+]] = arith.constant dense<[2, 1, 0]> : vector<3xindex>
-// CHECK: %[[T0:.+]] = arith.muli %[[ARG2]], %[[C3]] : index
-// CHECK: %[[T1:.+]] = vector.broadcast %[[T0]] : index to vector<1x1x3xindex>
-// CHECK: %[[T2:.+]] = vector.broadcast %[[INIT_IDX]]
-// CHECK: %[[T3:.+]] = arith.addi %[[T2]], %[[T1]]
-// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[T3]]], %[[MASK]], %[[PASSTHRU]]
+// CHECK: %[[T0:.+]] = vector.broadcast %[[ARG2]] : index to vector<1x1x3xindex>
+// CHECK: %[[T1:.+]] = vector.broadcast %[[INIT_IDX]]
+// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[CST_ZEROS]], %[[T0]], %[[T1]]], %[[MASK]], %[[PASSTHRU]]
// CHECK: vector.transfer_write %[[GATHER]]
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/vectorization/extract.mlir b/mlir/test/Dialect/Linalg/vectorization/extract.mlir
index 76ac4b8398069..66b3b215b6adf 100644
--- a/mlir/test/Dialect/Linalg/vectorization/extract.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/extract.mlir
@@ -252,17 +252,14 @@ func.func @masked_vectorize_nd_tensor_extract_with_affine_apply_gather(%6: tenso
// CHECK-LABEL: func.func @masked_vectorize_nd_tensor_extract_with_affine_apply_gather
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 3 : index
-// CHECK: %[[VAL_8:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_5]] : vector<1x4xi1>
-// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<1x3xf32>, vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
-// CHECK: %[[VAL_11:.*]] = vector.broadcast {{.*}} : index to vector<4xindex>
-// CHECK: %[[VAL_12:.*]] = arith.addi {{.*}} : vector<4xindex>
-// CHECK: %[[VAL_16:.*]] = vector.broadcast {{.*}} : vector<4xindex> to vector<1x4xindex>
-// CHECK: %[[VAL_18:.*]] = tensor.dim {{.*}} : tensor<80x16xf32>
-// CHECK: %[[VAL_19:.*]] = vector.broadcast {{.*}} : index to vector<1x4xindex>
-// CHECK: %[[VAL_20:.*]] = arith.muli {{.*}} : vector<1x4xindex>
-// CHECK: %[[VAL_22:.*]] = arith.addi {{.*}} : vector<1x4xindex>
-// CHECK: %[[VAL_23:.*]] = vector.mask %[[VAL_8]] { vector.gather {{.*}} : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
-// CHECK: %[[VAL_25:.*]] = vector.mask %[[VAL_8]] { vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x3xf32> } : vector<1x4xi1> -> tensor<1x3xf32>
+// CHECK: %[[VAL_6:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_5]] : vector<1x4xi1>
+// CHECK: %[[VAL_7:.*]] = vector.mask %[[VAL_6]] { vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<1x3xf32>, vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
+// CHECK: %[[VAL_8:.*]] = vector.broadcast {{.*}} : index to vector<4xindex>
+// CHECK: %[[VAL_9:.*]] = arith.addi {{.*}} : vector<4xindex>
+// CHECK: %[[VAL_10:.*]] = vector.broadcast {{.*}} : vector<4xindex> to vector<1x4xindex>
+// CHECK: %[[CST_16:.*]] = arith.constant dense<16> : vector<1x4xindex>
+// CHECK: %[[VAL_11:.*]] = vector.mask %[[VAL_6]] { vector.gather {{.*}}[%{{.*}}, %{{.*}}] [%[[VAL_10]], %[[CST_16]]], {{.*}} : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
+// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_6]] { vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x3xf32> } : vector<1x4xi1> -> tensor<1x3xf32>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -293,7 +290,6 @@ func.func @masked_dynamic_vectorize_nd_tensor_extract_with_affine_apply_gather(%
// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: index,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
-// CHECK: %[[VAL_3:.*]] = arith.constant 16 : index
// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_5:.*]] = tensor.dim %[[VAL_2]], %[[VAL_4]] : tensor<?x?xf32>
// CHECK: %[[VAL_6:.*]] = arith.constant 1 : index
@@ -309,16 +305,11 @@ func.func @masked_dynamic_vectorize_nd_tensor_extract_with_affine_apply_gather(%
// CHECK: %[[VAL_16:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32>
// CHECK: %[[VAL_17:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_18:.*]] = vector.broadcast %[[VAL_14]] : vector<4xindex> to vector<1x4xindex>
-// CHECK: %[[VAL_19:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_20:.*]] = tensor.dim %[[VAL_0]], %[[VAL_19]] : tensor<?x?xf32>
-// CHECK: %[[VAL_21:.*]] = vector.broadcast %[[VAL_20]] : index to vector<1x4xindex>
-// CHECK: %[[VAL_22:.*]] = arith.muli %[[VAL_18]], %[[VAL_21]] : vector<1x4xindex>
-// CHECK: %[[VAL_23:.*]] = arith.constant dense<16> : vector<1x4xindex>
-// CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_22]] : vector<1x4xindex>
-// CHECK: %[[VAL_25:.*]] = vector.mask %[[VAL_10]] { vector.gather %[[VAL_0]]{{\[}}%[[VAL_17]], %[[VAL_17]]] {{\[}}%[[VAL_24]]], %[[VAL_15]], %[[VAL_16]] : tensor<?x?xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
-// CHECK: %[[VAL_26:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_27:.*]] = vector.mask %[[VAL_10]] { vector.transfer_write %[[VAL_25]], %[[VAL_2]]{{\[}}%[[VAL_26]], %[[VAL_26]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<?x?xf32> } : vector<1x4xi1> -> tensor<?x?xf32>
-// CHECK: return %[[VAL_27]] : tensor<?x?xf32>
+// CHECK: %[[CST_16:.*]] = arith.constant dense<16> : vector<1x4xindex>
+// CHECK: %[[VAL_19:.*]] = vector.mask %[[VAL_10]] { vector.gather %[[VAL_0]]{{\[}}%[[VAL_17]], %[[VAL_17]]] {{\[}}%[[VAL_18]], %[[CST_16]]], %[[VAL_15]], %[[VAL_16]] : tensor<?x?xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
+// CHECK: %[[VAL_20:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_21:.*]] = vector.mask %[[VAL_10]] { vector.transfer_write %[[VAL_19]], %[[VAL_2]]{{\[}}%[[VAL_20]], %[[VAL_20]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<?x?xf32> } : vector<1x4xi1> -> tensor<?x?xf32>
+// CHECK: return %[[VAL_21]] : tensor<?x?xf32>
// CHECK: }
module attributes {transform.with_named_sequence} {
@@ -349,8 +340,6 @@ func.func @extract_masked_vectorize(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf3
// CHECK-LABEL: func.func @extract_masked_vectorize(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
-// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_3:.*]] = arith.constant 2 : index
// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_5:.*]] = tensor.dim %[[VAL_1]], %[[VAL_4]] : tensor<?x?xf32>
// CHECK: %[[VAL_6:.*]] = arith.constant 1 : index
@@ -362,16 +351,11 @@ func.func @extract_masked_vectorize(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf3
// CHECK: %[[VAL_12:.*]] = arith.constant dense<true> : vector<3x3xi1>
// CHECK: %[[VAL_13:.*]] = arith.constant dense<0.000000e+00> : vector<3x3xf32>
// CHECK: %[[VAL_14:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_15:.*]] = arith.constant dense<1> : vector<3x3xindex>
-// CHECK: %[[VAL_16:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_17:.*]] = tensor.dim %[[VAL_0]], %[[VAL_16]] : tensor<?x?xf32>
-// CHECK: %[[VAL_18:.*]] = vector.broadcast %[[VAL_17]] : index to vector<3x3xindex>
-// CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_15]], %[[VAL_18]] : vector<3x3xindex>
-// CHECK: %[[VAL_20:.*]] = arith.constant dense<2> : vector<3x3xindex>
-// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_19]] : vector<3x3xindex>
-// CHECK: %[[VAL_22:.*]] = vector.mask %[[VAL_10]] { vector.gather %[[VAL_0]]{{\[}}%[[VAL_14]], %[[VAL_14]]] {{\[}}%[[VAL_21]]], %[[VAL_12]], %[[VAL_13]] : tensor<?x?xf32>, vector<3x3xindex>, vector<3x3xi1>, vector<3x3xf32> into vector<3x3xf32> } : vector<3x3xi1> -> vector<3x3xf32>
-// CHECK: %[[VAL_23:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_24:.*]] = vector.mask %[[VAL_10]] { vector.transfer_write %[[VAL_22]], %[[VAL_1]]{{\[}}%[[VAL_23]], %[[VAL_23]]] {in_bounds = [true, true]} : vector<3x3xf32>, tensor<?x?xf32> } : vector<3x3xi1> -> tensor<?x?xf32>
+// CHECK: %[[CST_1:.*]] = arith.constant dense<1> : vector<3x3xindex>
+// CHECK: %[[CST_2:.*]] = arith.constant dense<2> : vector<3x3xindex>
+// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_10]] { vector.gather %[[VAL_0]]{{\[}}%[[VAL_14]], %[[VAL_14]]] {{\[}}%[[CST_1]], %[[CST_2]]], %[[VAL_12]], %[[VAL_13]] : tensor<?x?xf32>, vector<3x3xindex>, vector<3x3xi1>, vector<3x3xf32> into vector<3x3xf32> } : vector<3x3xi1> -> vector<3x3xf32>
+// CHECK: %[[VAL_16:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_17:.*]] = vector.mask %[[VAL_10]] { vector.transfer_write %[[VAL_15]], %[[VAL_1]]{{\[}}%[[VAL_16]], %[[VAL_16]]] {in_bounds = [true, true]} : vector<3x3xf32>, tensor<?x?xf32> } : vector<3x3xi1> -> tensor<?x?xf32>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -416,7 +400,8 @@ func.func @tensor_extract_dynamic_shape(%arg1: tensor<123x321xf32>, %arg2: tenso
// CHECK: %[[MASK_2:.*]] = arith.constant dense<true> : vector<1x3x8xi1>
// CHECK: %[[FALLTHROUGH:.*]] = arith.constant dense<0.000000e+00> : vector<1x3x8xf32>
// CHECK: %[[C0_1:.*]] = arith.constant 0 : index
-// CHECK: vector.mask %[[MASK]] { vector.gather %[[ARG_1]][%[[C0_1]], %[[C0_1]]] [%{{.*}}], %[[MASK_2]], %[[FALLTHROUGH]] : tensor<123x321xf32>, vector<1x3x8xindex>, vector<1x3x8xi1>, vector<1x3x8xf32> into vector<1x3x8xf32> } : vector<1x3x8xi1> -> vector<1x3x8xf32>
+// CHECK: %[[CST_1:.*]] = arith.constant dense<1> : vector<1x3x8xindex>
+// CHECK: vector.mask %[[MASK]] { vector.gather %[[ARG_1]][%[[C0_1]], %[[C0_1]]] [%[[CST_1]], %{{.*}}], %[[MASK_2]], %[[FALLTHROUGH]] : tensor<123x321xf32>, vector<1x3x8xindex>, vector<1x3x8xi1>, vector<1x3x8xf32> into vector<1x3x8xf32> } : vector<1x3x8xi1> -> vector<1x3x8xf32>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6aa92ab79a0dd..f3fb2eec04816 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4280,6 +4280,45 @@ func.func @scatter_tensor_all_false(%base: tensor<16xf32>,
// -----
+// CHECK-LABEL: func @fold_broadcast_index_into_gather_offset(
+// CHECK-SAME: %[[BASE:[a-z0-9]+]]: memref<?x?xf32>, %[[OFF0:[a-z0-9]+]]: index,
+// CHECK-SAME: %[[OFF1:[a-z0-9]+]]: index, %[[SCALAR:[a-z0-9]+]]: index,
+// CHECK-SAME: %[[IDX1:[a-z0-9]+]]: vector<4xindex>, %[[MASK:[a-z0-9]+]]: vector<4xi1>,
+// CHECK-SAME: %[[PT:[a-z0-9]+]]: vector<4xf32>)
+// CHECK: %[[NEW_OFF:.*]] = arith.addi %[[OFF0]], %[[SCALAR]] : index
+// CHECK: %[[RES:.*]] = vector.gather %[[BASE]][%[[NEW_OFF]], %[[OFF1]]] [%[[IDX1]]], %[[MASK]], %[[PT]] : memref<?x?xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+// CHECK: return %[[RES]] : vector<4xf32>
+func.func @fold_broadcast_index_into_gather_offset(
+ %base: memref<?x?xf32>, %off0: index, %off1: index,
+ %scalar: index, %idx1: vector<4xindex>,
+ %mask: vector<4xi1>, %pass_thru: vector<4xf32>) -> vector<4xf32> {
+ %bcast = vector.broadcast %scalar : index to vector<4xindex>
+ %0 = vector.gather %base[%off0, %off1][%bcast, %idx1], %mask, %pass_thru
+ : memref<?x?xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_broadcast_index_into_scatter_offset(
+// CHECK-SAME: %[[BASE:[a-z0-9]+]]: memref<?x?xf32>, %[[OFF0:[a-z0-9]+]]: index,
+// CHECK-SAME: %[[OFF1:[a-z0-9]+]]: index, %[[SCALAR:[a-z0-9]+]]: index,
+// CHECK-SAME: %[[IDX1:[a-z0-9]+]]: vector<4xindex>, %[[MASK:[a-z0-9]+]]: vector<4xi1>,
+// CHECK-SAME: %[[VAL:[a-z0-9]+]]: vector<4xf32>)
+// CHECK: %[[NEW_OFF:.*]] = arith.addi %[[OFF0]], %[[SCALAR]] : index
+// CHECK: vector.scatter %[[BASE]][%[[NEW_OFF]], %[[OFF1]]] [%[[IDX1]]], %[[MASK]], %[[VAL]] : memref<?x?xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32>
+func.func @fold_broadcast_index_into_scatter_offset(
+ %base: memref<?x?xf32>, %off0: index, %off1: index,
+ %scalar: index, %idx1: vector<4xindex>,
+ %mask: vector<4xi1>, %value: vector<4xf32>) {
+ %bcast = vector.broadcast %scalar : index to vector<4xindex>
+ vector.scatter %base[%off0, %off1][%bcast, %idx1], %mask, %value
+ : memref<?x?xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @fold_extract_constant_indices
// CHECK-SAME: %[[ARG:.*]]: vector<32x1xi32>) -> i32 {
// CHECK: %[[RES:.*]] = vector.extract %[[ARG]][0, 0] : i32 from vector<32x1xi32>
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index f90312c915334..e0205e7c93a70 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1554,6 +1554,78 @@ func.func @gather_tensor_alignment(%base: tensor<16xf32>, %indices: vector<16xi3
// -----
+// minSize on variadics doesn't currently add a verifier entry, check that our
+// manual diagnostic works.
+func.func @gather_no_indices(%base: memref<?xf32>,
+ %mask: vector<4xi1>, %pass_thru: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{'vector.gather' op requires at least one index vector}}
+ %0 = "vector.gather"(%base, %c0, %mask, %pass_thru)
+ <{operandSegmentSizes = array<i32: 1, 1, 0, 1, 1>}>
+ : (memref<?xf32>, index, vector<4xi1>, vector<4xf32>) -> vector<4xf32>
+}
+
+// -----
+
+func.func @scatter_no_indices(%base: memref<?xf32>,
+ %mask: vector<4xi1>, %value: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{'vector.scatter' op requires at least one index vector}}
+ "vector.scatter"(%base, %c0, %mask, %value)
+ <{operandSegmentSizes = array<i32: 1, 1, 0, 1, 1>}>
+ : (memref<?xf32>, index, vector<4xi1>, vector<4xf32>) -> ()
+}
+
+// -----
+
+func.func @gather_mismatched_index_types(%base: memref<?x?xf32>,
+ %idx0: vector<4xi32>, %idx1: vector<4xi64>,
+ %mask: vector<4xi1>, %pass_thru: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{'vector.gather' op all index vectors must have the same type}}
+ %0 = "vector.gather"(%base, %c0, %c0, %idx0, %idx1, %mask, %pass_thru)
+ <{operandSegmentSizes = array<i32: 1, 2, 2, 1, 1>}>
+ : (memref<?x?xf32>, index, index, vector<4xi32>, vector<4xi64>,
+ vector<4xi1>, vector<4xf32>) -> vector<4xf32>
+}
+
+// -----
+
+func.func @scatter_mismatched_index_types(%base: memref<?x?xf32>,
+ %idx0: vector<4xi32>, %idx1: vector<4xi64>,
+ %mask: vector<4xi1>, %value: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{'vector.scatter' op all index vectors must have the same type}}
+ "vector.scatter"(%base, %c0, %c0, %idx0, %idx1, %mask, %value)
+ <{operandSegmentSizes = array<i32: 1, 2, 2, 1, 1>}>
+ : (memref<?x?xf32>, index, index, vector<4xi32>, vector<4xi64>,
+ vector<4xi1>, vector<4xf32>) -> ()
+}
+
+// -----
+
+func.func @gather_too_many_indices(%base: memref<?xf32>,
+ %idx0: vector<4xi32>, %idx1: vector<4xi32>,
+ %mask: vector<4xi1>, %pass_thru: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{'vector.gather' op number of index vectors (2) exceeds base rank (1)}}
+ %0 = vector.gather %base[%c0][%idx0, %idx1], %mask, %pass_thru
+ : memref<?xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+}
+
+// -----
+
+func.func @scatter_too_many_indices(%base: memref<?xf32>,
+ %idx0: vector<4xi32>, %idx1: vector<4xi32>,
+ %mask: vector<4xi1>, %value: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{'vector.scatter' op number of index vectors (2) exceeds base rank (1)}}
+ vector.scatter %base[%c0][%idx0, %idx1], %mask, %value
+ : memref<?xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32>
+}
+
+// -----
+
func.func @scatter_to_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
%c0 = arith.constant 0 : index
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index de620221944de..4ee79897ca14f 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -867,6 +867,26 @@ func.func @gather_and_scatter_multi_dims(%base: memref<?xf32>, %v: vector<2x16xi
return %0 : vector<2x16xf32>
}
+// CHECK-LABEL: @gather_multi_index
+func.func @gather_multi_index(%base: memref<?x?x?xf32>, %v0: vector<4xi32>,
+ %v1: vector<4xi32>, %mask: vector<4xi1>, %pass_thru: vector<4xf32>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ // CHECK: vector.gather %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] [%{{.*}}, %{{.*}}], %{{.*}}, %{{.*}} : memref<?x?x?xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ %0 = vector.gather %base[%c0, %c0, %c0][%v0, %v1], %mask, %pass_thru
+ : memref<?x?x?xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: @scatter_multi_index
+func.func @scatter_multi_index(%base: memref<?x?xf32>, %v0: vector<4xi32>,
+ %v1: vector<4xi32>, %mask: vector<4xi1>, %value: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ // CHECK: vector.scatter %{{.*}}[%{{.*}}, %{{.*}}] [%{{.*}}, %{{.*}}], %{{.*}}, %{{.*}} : memref<?x?xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32>
+ vector.scatter %base[%c0, %c0][%v0, %v1], %mask, %value
+ : memref<?x?xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32>
+ return
+}
+
// CHECK-LABEL: @gather_on_tensor
func.func @gather_on_tensor(%base: tensor<?xf32>, %v: vector<16xi32>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>) -> vector<16xf32> {
%c0 = arith.constant 0 : index
diff --git a/mlir/test/Dialect/Vector/vector-gather-lowering.mlir b/mlir/test/Dialect/Vector/vector-gather-lowering.mlir
index 59b13e300e5e5..6bf0b4fe34476 100644
--- a/mlir/test/Dialect/Vector/vector-gather-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-gather-lowering.mlir
@@ -54,13 +54,11 @@ func.func @gather_memref_1d_i32_index(%base: memref<?xf32>, %v: vector<2xi32>, %
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: [[PTV0:%.+]] = vector.extract [[PASS]][0] : vector<3xf32> from vector<2x3xf32>
-// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[C0]], %[[C1]]] by
// CHECK-DAG: [[M0:%.+]] = vector.extract [[MASK]][0, 0] : i1 from vector<2x3xi1>
// CHECK-DAG: [[IDX0:%.+]] = vector.extract [[IDXVEC]][0, 0] : index from vector<2x3xindex>
-// CHECK: %[[FLAT0:.+]] = arith.addi %[[LIN]], [[IDX0]] : index
-// CHECK: %[[DL0:.+]]:2 = affine.delinearize_index %[[FLAT0]] into
+// CHECK: %[[OFF0:.+]] = arith.addi [[IDX0]], %[[C1]] : index
// CHECK: [[RES0:%.+]] = scf.if [[M0]] -> (vector<3xf32>)
-// CHECK-NEXT: [[LD0:%.+]] = vector.load [[BASE]][%[[DL0]]#0, %[[DL0]]#1] : memref<?x?xf32>, vector<1xf32>
+// CHECK-NEXT: [[LD0:%.+]] = vector.load [[BASE]][%[[C0]], %[[OFF0]]] : memref<?x?xf32>, vector<1xf32>
// CHECK-NEXT: [[ELEM0:%.+]] = vector.extract [[LD0]][0] : f32 from vector<1xf32>
// CHECK-NEXT: [[INS0:%.+]] = vector.insert [[ELEM0]], [[PTV0]] [0] : f32 into vector<3xf32>
// CHECK-NEXT: scf.yield [[INS0]] : vector<3xf32>
@@ -256,32 +254,157 @@ func.func @strided_gather(%base : memref<100x3xf32>,
// CHECK-SAME: %[[IDXS:.*]]: vector<4xindex>,
// CHECK-SAME: %[[VAL_4:.*]]: index,
// CHECK-SAME: %[[VAL_5:.*]]: index) -> vector<4xf32> {
-// CHECK: %[[TRUE:.*]] = arith.constant true
-// CHECK: %[[CST_3:.*]] = arith.constant dense<3> : vector<4xindex>
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[TRUE:.*]] = arith.constant true
-// CHECK: %[[COLLAPSED:.*]] = memref.collapse_shape %[[base]] {{\[\[}}0, 1]] : memref<100x3xf32> into memref<300xf32>
-// CHECK: %[[NEW_IDXS:.*]] = arith.muli %[[IDXS]], %[[CST_3]] : vector<4xindex>
-
-// CHECK: %[[IDX_0:.*]] = vector.extract %[[NEW_IDXS]][0] : index from vector<4xindex>
+// CHECK: %[[IDX_0:.*]] = vector.extract %[[IDXS]][0] : index from vector<4xindex>
// CHECK: scf.if %[[TRUE]] -> (vector<4xf32>)
-// CHECK: %[[M_0:.*]] = vector.load %[[COLLAPSED]][%[[IDX_0]]] {alignment = 8 : i64} : memref<300xf32>, vector<1xf32>
+// CHECK: %[[M_0:.*]] = vector.load %[[base]][%[[IDX_0]], %[[C0]]] {alignment = 8 : i64} : memref<100x3xf32>, vector<1xf32>
// CHECK: %[[V_0:.*]] = vector.extract %[[M_0]][0] : f32 from vector<1xf32>
-// CHECK: %[[IDX_1:.*]] = vector.extract %[[NEW_IDXS]][1] : index from vector<4xindex>
+// CHECK: %[[IDX_1:.*]] = vector.extract %[[IDXS]][1] : index from vector<4xindex>
// CHECK: scf.if %[[TRUE]] -> (vector<4xf32>)
-// CHECK: %[[M_1:.*]] = vector.load %[[COLLAPSED]][%[[IDX_1]]] {alignment = 8 : i64} : memref<300xf32>, vector<1xf32>
+// CHECK: %[[M_1:.*]] = vector.load %[[base]][%[[IDX_1]], %[[C0]]] {alignment = 8 : i64} : memref<100x3xf32>, vector<1xf32>
// CHECK: %[[V_1:.*]] = vector.extract %[[M_1]][0] : f32 from vector<1xf32>
-// CHECK: %[[IDX_2:.*]] = vector.extract %[[NEW_IDXS]][2] : index from vector<4xindex>
+// CHECK: %[[IDX_2:.*]] = vector.extract %[[IDXS]][2] : index from vector<4xindex>
// CHECK: scf.if %[[TRUE]] -> (vector<4xf32>)
-// CHECK: %[[M_2:.*]] = vector.load %[[COLLAPSED]][%[[IDX_2]]] {alignment = 8 : i64} : memref<300xf32>, vector<1xf32>
+// CHECK: %[[M_2:.*]] = vector.load %[[base]][%[[IDX_2]], %[[C0]]] {alignment = 8 : i64} : memref<100x3xf32>, vector<1xf32>
// CHECK: %[[V_2:.*]] = vector.extract %[[M_2]][0] : f32 from vector<1xf32>
-// CHECK: %[[IDX_3:.*]] = vector.extract %[[NEW_IDXS]][3] : index from vector<4xindex>
+// CHECK: %[[IDX_3:.*]] = vector.extract %[[IDXS]][3] : index from vector<4xindex>
// CHECK: scf.if %[[TRUE]] -> (vector<4xf32>)
-// CHECK: %[[M_3:.*]] = vector.load %[[COLLAPSED]][%[[IDX_3]]] {alignment = 8 : i64} : memref<300xf32>, vector<1xf32>
+// CHECK: %[[M_3:.*]] = vector.load %[[base]][%[[IDX_3]], %[[C0]]] {alignment = 8 : i64} : memref<100x3xf32>, vector<1xf32>
// CHECK: %[[V_3:.*]] = vector.extract %[[M_3]][0] : f32 from vector<1xf32>
+// Same as @strided_gather but with non-zero subview offsets. Both subview
+// offsets must end up in the lowered loads: the outer one composed into the
+// first index, the inner one into the second.
+func.func @strided_gather_nonzero_subview_offsets(
+ %base: memref<100x3xf32>, %idxs: vector<4xindex>, %row: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %subview = memref.subview %base[%row, %c2][95, 1][1, 1]
+ : memref<100x3xf32> to memref<95xf32, strided<[3], offset: ?>>
+ %mask = arith.constant dense<true> : vector<4xi1>
+ %pass = arith.constant dense<0.000000e+00> : vector<4xf32>
+ %res = vector.gather %subview[%c0] [%idxs], %mask, %pass
+ : memref<95xf32, strided<[3], offset: ?>>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ return %res : vector<4xf32>
+}
+// CHECK-LABEL: func.func @strided_gather_nonzero_subview_offsets(
+// CHECK-SAME: %[[BASE:.*]]: memref<100x3xf32>,
+// CHECK-SAME: %[[IDXS:.*]]: vector<4xindex>,
+// CHECK-SAME: %[[ROW:.*]]: index)
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[I0:.*]] = vector.extract %[[IDXS]][0]
+// CHECK: %[[O0:.*]] = arith.addi %[[ROW]], %[[I0]]
+// CHECK: vector.load %[[BASE]][%[[O0]], %[[C2]]]
+// CHECK: %[[I1:.*]] = vector.extract %[[IDXS]][1]
+// CHECK: %[[O1:.*]] = arith.addi %[[ROW]], %[[I1]]
+// CHECK: vector.load %[[BASE]][%[[O1]], %[[C2]]]
+
+// Dynamic outer subview stride: the rewrite splats the dynamic stride across
+// the index vector and multiplies the original indices by it before lowering
+// to the 2-D gather, so no `subview` survives in the output.
+func.func @strided_gather_dynamic_subview_stride(
+ %base: memref<100x3xf32>, %idxs: vector<4xindex>, %s: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %subview = memref.subview %base[0, 0][100, 1][%s, 1]
+ : memref<100x3xf32> to memref<100xf32, strided<[?]>>
+ %mask = arith.constant dense<true> : vector<4xi1>
+ %pass = arith.constant dense<0.000000e+00> : vector<4xf32>
+ %res = vector.gather %subview[%c0] [%idxs], %mask, %pass
+ : memref<100xf32, strided<[?]>>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ return %res : vector<4xf32>
+}
+// CHECK-LABEL: func.func @strided_gather_dynamic_subview_stride(
+// CHECK-SAME: %[[BASE:.*]]: memref<100x3xf32>,
+// CHECK-SAME: %[[IDXS:.*]]: vector<4xindex>,
+// CHECK-SAME: %[[S:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[S]] : index to vector<4xindex>
+// CHECK: %[[STRIDED:.*]] = arith.muli %[[BCAST]], %[[IDXS]] : vector<4xindex>
+// CHECK-NOT: memref.subview
+// CHECK: %[[I0:.*]] = vector.extract %[[STRIDED]][0]
+// CHECK: vector.load %[[BASE]][%[[I0]], %[[C0]]]
+// CHECK: %[[I1:.*]] = vector.extract %[[STRIDED]][1]
+// CHECK: vector.load %[[BASE]][%[[I1]], %[[C0]]]
+
+// Static non-unit outer subview stride: same shape as the dynamic case, but
+// the multiply is by a constant splat.
+func.func @strided_gather_static_outer_stride(
+ %base: memref<100x3xf32>, %idxs: vector<4xindex>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %subview = memref.subview %base[0, 0][50, 1][2, 1]
+ : memref<100x3xf32> to memref<50xf32, strided<[6]>>
+ %mask = arith.constant dense<true> : vector<4xi1>
+ %pass = arith.constant dense<0.000000e+00> : vector<4xf32>
+ %res = vector.gather %subview[%c0] [%idxs], %mask, %pass
+ : memref<50xf32, strided<[6]>>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ return %res : vector<4xf32>
+}
+// CHECK-LABEL: func.func @strided_gather_static_outer_stride(
+// CHECK-SAME: %[[BASE:.*]]: memref<100x3xf32>,
+// CHECK-SAME: %[[IDXS:.*]]: vector<4xindex>)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<2> : vector<4xindex>
+// CHECK: %[[STRIDED:.*]] = arith.muli %[[IDXS]], %[[CST_2]] : vector<4xindex>
+// CHECK-NOT: memref.subview
+// CHECK: %[[I0:.*]] = vector.extract %[[STRIDED]][0]
+// CHECK: vector.load %[[BASE]][%[[I0]], %[[C0]]]
+// CHECK: %[[I1:.*]] = vector.extract %[[STRIDED]][1]
+// CHECK: vector.load %[[BASE]][%[[I1]], %[[C0]]]
+
+// Regression: with a non-unit outer stride AND a non-zero gather offset, the
+// gather's offset is in subview-element units, so it must be scaled by the
+// outer stride before adding to the subview's outer offset. Previously, this
+// pattern silently dropped the scaling.
+func.func @strided_gather_static_stride_nonzero_offset(
+ %base: memref<100x3xf32>, %idxs: vector<4xindex>) -> vector<4xf32> {
+ %c5 = arith.constant 5 : index
+ %subview = memref.subview %base[0, 0][50, 1][2, 1]
+ : memref<100x3xf32> to memref<50xf32, strided<[6]>>
+ %mask = arith.constant dense<true> : vector<4xi1>
+ %pass = arith.constant dense<0.000000e+00> : vector<4xf32>
+ %res = vector.gather %subview[%c5] [%idxs], %mask, %pass
+ : memref<50xf32, strided<[6]>>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+ return %res : vector<4xf32>
+}
+// CHECK-LABEL: func.func @strided_gather_static_stride_nonzero_offset(
+// CHECK-SAME: %[[BASE:.*]]: memref<100x3xf32>,
+// CHECK-SAME: %[[IDXS:.*]]: vector<4xindex>)
+// `5 * 2 = 10` source rows skipped at the gather offset, plus the outer
+// subview offset of 0.
+// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+// CHECK: %[[I0:.*]] = vector.extract %{{.*}}[0]
+// CHECK: %[[O0:.*]] = arith.addi %[[I0]], %[[C10]]
+// CHECK: vector.load %[[BASE]][%[[O0]], %{{.*}}]
+
+// Verify that multi-index gather on a 2D memref correctly offsets each
+// dimension independently.
+// CHECK-LABEL: @gather_memref_2d_multi_index
+// CHECK-SAME: (%[[BASE:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[IDX0:.+]]: vector<2xindex>, %[[IDX1:.+]]: vector<2xindex>,
+// CHECK-SAME: %[[MASK:.+]]: vector<2xi1>, %[[PASS:.+]]: vector<2xf32>)
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK: %[[M0:.+]] = vector.extract %[[MASK]][0]
+// CHECK: %[[I0_0:.+]] = vector.extract %[[IDX0]][0]
+// CHECK: %[[I1_0:.+]] = vector.extract %[[IDX1]][0]
+// CHECK: %[[OFF1_0:.+]] = arith.addi %[[I1_0]], %[[C1]]
+// CHECK: scf.if %[[M0]]
+// CHECK: vector.load %[[BASE]][%[[I0_0]], %[[OFF1_0]]]
+func.func @gather_memref_2d_multi_index(
+ %base: memref<?x?xf32>,
+ %idx0: vector<2xindex>, %idx1: vector<2xindex>,
+ %mask: vector<2xi1>, %pass_thru: vector<2xf32>) -> vector<2xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = vector.gather %base[%c0, %c1][%idx0, %idx1], %mask, %pass_thru
+ : memref<?x?xf32>, vector<2xindex>, vector<2xi1>, vector<2xf32> into vector<2xf32>
+ return %0 : vector<2xf32>
+}
+
// CHECK-LABEL: @scalable_gather_1d
// CHECK-NOT: extract
// CHECK: vector.gather
@@ -292,35 +415,32 @@ func.func @scalable_gather_1d(%base: tensor<?xf32>, %v: vector<[2]xindex>, %mask
return %0 : vector<[2]xf32>
}
-// Verify that gather on a 2D memref delinearizes the gather index.
-// With zero base offsets, the linearize and addi fold away.
+// Verify that gather on a 2D memref with zero base offsets directly
+// adds each index element to its corresponding offset.
// CHECK-LABEL: @gather_memref_2d_delinearize
// CHECK-SAME: (%[[BASE:.+]]: memref<4x2xf32>,
// CHECK-SAME: %[[IDXVEC:.+]]: vector<4xi32>,
// CHECK-SAME: %[[MASK:.+]]: vector<4xi1>,
// CHECK-SAME: %[[PASS:.+]]: vector<4xf32>)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[IDXS:.+]] = arith.index_cast %[[IDXVEC]]
//
-// CHECK-DAG: %[[IDX0:.+]] = vector.extract %[[IDXS]][0]
-// CHECK: %[[DL0:.+]]:2 = affine.delinearize_index %[[IDX0]] into (4, 2)
+// CHECK: %[[IDX0:.+]] = vector.extract %[[IDXS]][0]
// CHECK: scf.if
-// CHECK: vector.load %[[BASE]][%[[DL0]]#0, %[[DL0]]#1] : memref<4x2xf32>, vector<1xf32>
+// CHECK: vector.load %[[BASE]][%[[C0]], %[[IDX0]]] : memref<4x2xf32>, vector<1xf32>
//
// CHECK: %[[IDX1:.+]] = vector.extract %[[IDXS]][1]
-// CHECK: affine.delinearize_index %[[IDX1]] into (4, 2)
// CHECK: scf.if
-// CHECK: vector.load %[[BASE]][%{{.+}}, %{{.+}}] : memref<4x2xf32>, vector<1xf32>
+// CHECK: vector.load %[[BASE]][%{{.+}}, %[[IDX1]]] : memref<4x2xf32>, vector<1xf32>
//
// CHECK: %[[IDX2:.+]] = vector.extract %[[IDXS]][2]
-// CHECK: affine.delinearize_index %[[IDX2]] into (4, 2)
// CHECK: scf.if
-// CHECK: vector.load %[[BASE]][%{{.+}}, %{{.+}}] : memref<4x2xf32>, vector<1xf32>
+// CHECK: vector.load %[[BASE]][%{{.+}}, %[[IDX2]]] : memref<4x2xf32>, vector<1xf32>
//
// CHECK: %[[IDX3:.+]] = vector.extract %[[IDXS]][3]
-// CHECK: affine.delinearize_index %[[IDX3]] into (4, 2)
// CHECK: scf.if
-// CHECK: vector.load %[[BASE]][%{{.+}}, %{{.+}}] : memref<4x2xf32>, vector<1xf32>
+// CHECK: vector.load %[[BASE]][%{{.+}}, %[[IDX3]]] : memref<4x2xf32>, vector<1xf32>
func.func @gather_memref_2d_delinearize(
%base: memref<4x2xf32>,
%v: vector<4xi32>, %mask: vector<4xi1>,
@@ -335,7 +455,7 @@ func.func @gather_memref_2d_delinearize(
// -----
// Verify that gather on a 2D memref with non-zero base offsets correctly
-// incorporates the offsets via linearize + add + delinearize.
+// adds each index element to the last-dimension offset.
// CHECK-LABEL: @gather_memref_2d_delinearize_nonzero_offsets
// CHECK-SAME: (%[[BASE:.+]]: memref<4x2xf32>,
@@ -344,12 +464,10 @@ func.func @gather_memref_2d_delinearize(
// CHECK-SAME: %[[MASK:.+]]: vector<2xi1>,
// CHECK-SAME: %[[PASS:.+]]: vector<2xf32>)
// CHECK-DAG: %[[IDXS:.+]] = arith.index_cast %[[IDXVEC]]
-// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[OFF0]], %[[OFF1]]] by (4, 2)
// CHECK: %[[IDX0:.+]] = vector.extract %[[IDXS]][0]
-// CHECK: %[[FLAT:.+]] = arith.addi %[[LIN]], %[[IDX0]]
-// CHECK: %[[DL:.+]]:2 = affine.delinearize_index %[[FLAT]] into (4, 2)
+// CHECK: %[[SUM0:.+]] = arith.addi %[[OFF1]], %[[IDX0]]
// CHECK: scf.if
-// CHECK: vector.load %[[BASE]][%[[DL]]#0, %[[DL]]#1]
+// CHECK: vector.load %[[BASE]][%[[OFF0]], %[[SUM0]]]
func.func @gather_memref_2d_delinearize_nonzero_offsets(
%base: memref<4x2xf32>,
%off0: index, %off1: index,
diff --git a/mlir/test/Dialect/XeGPU/xegpu-vector-linearize.mlir b/mlir/test/Dialect/XeGPU/xegpu-vector-linearize.mlir
index e4244b8071860..8ca82591d5475 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-vector-linearize.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-vector-linearize.mlir
@@ -165,14 +165,11 @@ func.func @broadcast_stretch_in_middle(%arg0: vector<4x1x2xf32>) -> vector<4x3x2
// First shuffle + if ladder for row 0
// CHECK: %[[ROW0_INIT:.*]] = vector.shuffle %[[PASS_CAST]], %[[POISON]] [0, 1, 2]
-// CHECK: %[[DIM0:.*]] = memref.dim %[[BASE]], %[[C0]]
-// CHECK: %[[DIM1:.*]] = memref.dim %[[BASE]], %[[C1]]
// CHECK: %[[MASK_0_0:.*]] = vector.extract %[[MASK]][0, 0]
// CHECK: %[[IDX_0_0:.*]] = vector.extract %[[IDX]][0, 0]
// CHECK: %[[OFF_0_0:.*]] = arith.addi %[[IDX_0_0]], %[[C1]]
-// CHECK: %[[DL_0_0:.*]]:2 = affine.delinearize_index %[[OFF_0_0]] into (%[[DIM0]], %[[DIM1]])
// CHECK: %[[IF_0_0:.*]] = scf.if %[[MASK_0_0]] -> (vector<3xf32>) {
-// CHECK: %[[LOAD_0_0:.*]] = vector.load %[[BASE]][%[[DL_0_0]]#0, %[[DL_0_0]]#1] : memref<?x?xf32>, vector<1xf32>
+// CHECK: %[[LOAD_0_0:.*]] = vector.load %[[BASE]][%[[C0]], %[[OFF_0_0]]] : memref<?x?xf32>, vector<1xf32>
// CHECK: %[[ELEM_0_0:.*]] = vector.extract %[[LOAD_0_0]][0] : f32
// CHECK: %[[INS_0_0:.*]] = vector.insert %[[ELEM_0_0]], %[[ROW0_INIT]] [0] : f32 into vector<3xf32>
// CHECK: scf.yield %[[INS_0_0]] : vector<3xf32>
@@ -183,10 +180,9 @@ func.func @broadcast_stretch_in_middle(%arg0: vector<4x1x2xf32>) -> vector<4x3x2
// CHECK: %[[MASK_0_1:.*]] = vector.extract %[[MASK]][0, 1]
// CHECK: %[[IDX_0_1:.*]] = vector.extract %[[IDX]][0, 1]
// CHECK: %[[OFF_0_1:.*]] = arith.addi %[[IDX_0_1]], %[[C1]]
-// CHECK: %[[DL_0_1:.*]]:2 = affine.delinearize_index %[[OFF_0_1]] into (%[[DIM0]], %[[DIM1]])
// CHECK: %[[IF_0_1:.*]] = scf.if %[[MASK_0_1]] -> (vector<3xf32>)
-// … (similar checks for the rest of row 0, then row 1)
+// ... (similar checks for the rest of row 0, then row 1)
// CHECK: %[[ROW_SHUFFLE:.*]] = vector.shuffle %[[POISON]], {{.*}} [6, 7, 8, -1, -1, -1]
// CHECK: %[[ROW1_INIT:.*]] = vector.shuffle %[[PASS_CAST]], %[[POISON]] [3, 4, 5]
@@ -195,10 +191,9 @@ func.func @broadcast_stretch_in_middle(%arg0: vector<4x1x2xf32>) -> vector<4x3x2
// CHECK: %[[MASK_1_0:.*]] = vector.extract %[[MASK]][1, 0]
// CHECK: %[[IDX_1_0:.*]] = vector.extract %[[IDX]][1, 0]
// CHECK: %[[OFF_1_0:.*]] = arith.addi %[[IDX_1_0]], %[[C1]]
-// CHECK: %[[DL_1_0:.*]]:2 = affine.delinearize_index %[[OFF_1_0]] into
// CHECK: %[[IF_1_0:.*]] = scf.if %[[MASK_1_0]] -> (vector<3xf32>)
-// … (similar checks for remaining row 1 inserts)
+// ... (similar checks for remaining row 1 inserts)
// Final reshuffle and cast
// CHECK: %[[FINAL_SHUFFLE:.*]] = vector.shuffle %[[ROW_SHUFFLE]], {{.*}} [0, 1, 2, 6, 7, 8]
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/gather.mlir b/mlir/test/Integration/Dialect/Vector/CPU/gather.mlir
index 110a46e2d89d8..91aab3b2190a7 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/gather.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/gather.mlir
@@ -14,7 +14,7 @@
/// TEST 3. Verify that `test-vector-gather-lowering` will indeed produce
/// `vector.load`
// REDEFINE: %{compile} = mlir-opt %s --test-vector-gather-lowering
-// RUN: %{compile} | FileCheck %s -check-prefix CHECK-IR
+// RUN: %{compile} | FileCheck %s -check-prefix CHECK-IR
func.func @gather8(%base: memref<?x?xf32>, %indices: vector<8xi32>,
%mask: vector<8xi1>, %pass_thru: vector<8xf32>) -> vector<8xf32> {
@@ -26,6 +26,15 @@ func.func @gather8(%base: memref<?x?xf32>, %indices: vector<8xi32>,
return %g : vector<8xf32>
}
+func.func @gather8_2d(%base: memref<3x3xf32, strided<[?, 1]>>,
+ %indices0: vector<8xi32>, %indices1: vector<8xi32>,
+ %mask: vector<8xi1>, %pass_thru: vector<8xf32>) -> vector<8xf32> {
+ %c0 = arith.constant 0: index
+ %g = vector.gather %base[%c0, %c0][%indices0, %indices1], %mask, %pass_thru
+ : memref<3x3xf32, strided<[?, 1]>>, vector<8xi32>, vector<8xi1>, vector<8xf32> into vector<8xf32>
+ return %g : vector<8xf32>
+}
+
func.func @entry() {
// Set up memory.
%c0 = arith.constant 0: index
@@ -102,6 +111,26 @@ func.func @entry() {
vector.print %g4 : vector<8xf32>
// CHECK: ( 0, 31, 21, 63, -7, -7, -7, 42 )
+ // Multi-dimensional gather tests.
+ %idxs_flat0 = arith.constant dense<0> : vector<8xi32>
+ %idxs_flat1 = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xi32>
+ %idxs_2d0 = arith.constant dense<[0, 0, 0, 1, 1, 1, 2, 2]> : vector<8xi32>
+ %idxs_2d1 = arith.constant dense<[0, 1, 2, 0, 1, 2, 0, 1]> : vector<8xi32>
+
+ %A_part = memref.subview %A [0, 0] [3, 3] [1, 1] : memref<?x?xf32> to memref<3x3xf32, strided<[?, 1]>>
+
+ %g2d1 = call @gather8_2d(%A_part, %idxs_2d0, %idxs_2d1, %all, %pass)
+ : (memref<3x3xf32, strided<[?, 1]>>, vector<8xi32>, vector<8xi32>, vector<8xi1>, vector<8xf32>)
+ -> (vector<8xf32>)
+ // CHECK: ( 0, 1, 2, 10, 11, 12, 20, 22 )
+ vector.print %g2d1 : vector<8xf32>
+
+ %g2d2 = call @gather8_2d(%A_part, %idxs_flat0, %idxs_flat1, %all, %pass)
+ : (memref<3x3xf32, strided<[?, 1]>>, vector<8xi32>, vector<8xi32>, vector<8xi1>, vector<8xf32>)
+ -> (vector<8xf32>)
+ // CHECK: ( 0, 1, 2, 3, 4, 10, 11, 12 )
+ vector.print %g2d2 : vector<8xf32>
+
memref.dealloc %A : memref<?x?xf32>
return
}
More information about the Mlir-commits
mailing list