[Mlir-commits] [mlir] 6af0cbf - [mlir][vector] Account for subview offset in gather lowering. (#195359)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 5 08:55:34 PDT 2026
Author: Han-Chung Wang
Date: 2026-05-05T08:55:28-07:00
New Revision: 6af0cbfe009a693f1fdf557bc0f31a36a73cc659
URL: https://github.com/llvm/llvm-project/commit/6af0cbfe009a693f1fdf557bc0f31a36a73cc659
DIFF: https://github.com/llvm/llvm-project/commit/6af0cbfe009a693f1fdf557bc0f31a36a73cc659.diff
LOG: [mlir][vector] Account for subview offset in gather lowering. (#195359)
Strided vector.gather on a column subview was reading the wrong column
because the rewrite to a collapsed gather dropped the subview's static
offset.
---------
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
mlir/test/Dialect/Vector/vector-gather-lowering.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
index 7194d41d60df7..5a8e473d39360 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorGather.cpp
@@ -79,23 +79,29 @@ 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.
+/// MemRef with updated offsets/indices that model the strided access.
///
/// ```mlir
-/// %subview = memref.subview %M (...)
-/// : memref<100x3xf32> to memref<100xf32, strided<[3]>>
-/// %gather = vector.gather %subview[%idxs] (...)
-/// : memref<100xf32, strided<[3]>>
+/// %subview = memref.subview %M[%i, %j] [100, 1] [1, 1]
+/// : memref<100x3xf32> to memref<100xf32, strided<[3], offset: ?>>
+/// %gather = vector.gather %subview[%c0] [%idxs] (...)
+/// : memref<100xf32, strided<[3], offset: ?>>
/// ```
/// ==>
/// ```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] (...)
+/// %new_off = arith.addi %c0_scaled, %subview_offset : index
+/// %gather = vector.gather %collapse_shape[%new_off] [%new_idxs] (...)
/// : memref<300xf32> (...)
/// ```
///
+/// The subview's static offset (the linearized position of the first element
+/// in the source memref) must be folded into the gather's base offsets, so a
+/// subview that selects e.g. column `j_sub` of a row-major `MxN` memref still
+/// reads from `M_base + j_sub + idx * N` instead of `M_base + idx * N`.
+///
/// ATM this is effectively limited to reading a 1D Vector from a 2D MemRef,
/// but should be fairly straightforward to extend beyond that.
struct RemoveStrideFromGatherSource : OpRewritePattern<vector::GatherOp> {
@@ -134,27 +140,56 @@ struct RemoveStrideFromGatherSource : OpRewritePattern<vector::GatherOp> {
if (stridedLayoutAttr.getStrides()[0] != srcTrailingDim)
return failure();
+ // The result memref's offset is the linearized position of the subview's
+ // first element within the source memref. Bail out on dynamic offsets so
+ // we don't have to materialize them; the conditional-load fallback will
+ // still produce correct code.
+ // TODO: Support dynamic offsets.
+ int64_t subviewOffset = stridedLayoutAttr.getOffset();
+ if (ShapedType::isDynamic(subviewOffset))
+ 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.
+ // 2. Generate new gather indices that will model the strided access.
+ // Take `memref<4xf32, strided<[3], offset: 1>>` and lane k as an example.
+ // For the rewrite to be correct, the flat positions must match:
+ // new_off + new_idxs[k] = 1 + (base_off + idxs[k]) * 3
+ // = 1 + base_off * 3 + idxs[k] * 3
+ // So the newIdxs is scaled with the stride.
IntegerAttr stride = rewriter.getIndexAttr(srcTrailingDim);
VectorType vType = op.getIndices().getType();
Value mulCst = arith::ConstantOp::create(
rewriter, op.getLoc(), vType, DenseElementsAttr::get(vType, stride));
-
Value newIdxs =
arith::MulIOp::create(rewriter, op.getLoc(), op.getIndices(), mulCst);
- // 3. Create an updated gather op with the collapsed input memref and the
- // updated indices.
+ // 3. Linearize the gather's base offsets through the source memref. On the
+ // collapsed memref the trailing offset must be scaled by the source's
+ // trailing dim and shifted by the subview's static offset.
+ // Pick new_idxs[k] = idxs[k] * 3 (that's step 2), and solve for new_off:
+ // new_off = 1 + base_off * 3
+ // = subview_offset + base_off * stride
+ // Note that createOrFold collapses the muli/addi when the trailing offset
+ // is a constant zero or the subview offset is zero.
+ SmallVector<Value> newOffsets(op.getOffsets());
+ Value strideVal =
+ arith::ConstantIndexOp::create(rewriter, op.getLoc(), srcTrailingDim);
+ newOffsets.back() = rewriter.createOrFold<arith::MulIOp>(
+ op.getLoc(), newOffsets.back(), strideVal);
+ Value subviewOffsetValue =
+ arith::ConstantIndexOp::create(rewriter, op.getLoc(), subviewOffset);
+ newOffsets.back() = rewriter.createOrFold<arith::AddIOp>(
+ op.getLoc(), newOffsets.back(), subviewOffsetValue);
+
+ // 4. Create an updated gather op with the collapsed input memref and the
+ // updated offsets/indices.
Value newGather = vector::GatherOp::create(
- rewriter, op.getLoc(), op.getResult().getType(), collapsed,
- op.getOffsets(), newIdxs, op.getMask(), op.getPassThru(),
- op.getAlignmentAttr());
+ rewriter, op.getLoc(), op.getResult().getType(), collapsed, newOffsets,
+ newIdxs, op.getMask(), op.getPassThru(), op.getAlignmentAttr());
rewriter.replaceOp(op, newGather);
return success();
diff --git a/mlir/test/Dialect/Vector/vector-gather-lowering.mlir b/mlir/test/Dialect/Vector/vector-gather-lowering.mlir
index 59b13e300e5e5..cee504078fe6d 100644
--- a/mlir/test/Dialect/Vector/vector-gather-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-gather-lowering.mlir
@@ -360,3 +360,64 @@ func.func @gather_memref_2d_delinearize_nonzero_offsets(
vector<2xi1>, vector<2xf32> into vector<2xf32>
return %0 : vector<2xf32>
}
+
+// -----
+
+// CHECK-LABEL: func.func @strided_gather_with_non_zero_offset(
+// CHECK-SAME: %[[BASE:.+]]: memref<6x5xf32>,
+// CHECK-SAME: %[[IDXS:.+]]: vector<2xindex>
+// CHECK-DAG: %[[OFFSET:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[STRIDE:.+]] = arith.constant dense<5> : vector<2xindex>
+// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[BASE]] {{\[\[}}0, 1]] : memref<6x5xf32> into memref<30xf32>
+// Scale each gather lane by the source's trailing dim (stride 5).
+// CHECK: %[[NEW_IDXS:.+]] = arith.muli %[[IDXS]], %[[STRIDE]]
+// CHECK: %[[IDX_0:.+]] = vector.extract %[[NEW_IDXS]][0]
+// Shift by the subview's static offset (3) so the flat load address matches
+// the original strided access.
+// CHECK: %[[ADDR_0:.+]] = arith.addi %[[IDX_0]], %[[OFFSET]]
+// CHECK: scf.if
+// CHECK: vector.load %[[COLLAPSED]][%[[ADDR_0]]] : memref<30xf32>
+// CHECK: %[[IDX_1:.+]] = vector.extract %[[NEW_IDXS]][1]
+// CHECK: %[[ADDR_1:.+]] = arith.addi %[[IDX_1]], %[[OFFSET]]
+// CHECK: scf.if
+// CHECK: vector.load %[[COLLAPSED]][%[[ADDR_1]]] : memref<30xf32>
+func.func @strided_gather_with_non_zero_offset(%base: memref<6x5xf32>,
+ %idxs: vector<2xindex>,
+ %mask: vector<2xi1>,
+ %pass_thru: vector<2xf32>)
+ -> vector<2xf32> {
+ %c0 = arith.constant 0 : index
+ %sub = memref.subview %base[0, 3] [6, 1] [1, 1]
+ : memref<6x5xf32> to memref<6xf32, strided<[5], offset: 3>>
+ %0 = vector.gather %sub[%c0] [%idxs], %mask, %pass_thru
+ : memref<6xf32, strided<[5], offset: 3>>, vector<2xindex>,
+ vector<2xi1>, vector<2xf32> into vector<2xf32>
+ return %0 : vector<2xf32>
+}
+
+// -----
+
+// TODO: Support dynamic offsets.
+// CHECK-LABEL: func.func @negative_strided_gather_with_dynamic_offset(
+// CHECK-SAME: %[[BASE:.+]]: memref<4x3xf32>,
+// CHECK-SAME: %[[COL:.+]]: index,
+// CHECK-NOT: memref.collapse_shape
+// CHECK: %[[SUB:.+]] = memref.subview %[[BASE]][0, %[[COL]]] [4, 1] [1, 1]
+// CHECK-SAME: : memref<4x3xf32> to memref<4xf32, strided<[3], offset: ?>>
+// CHECK: %[[RES:.+]] = vector.gather %[[SUB]]
+// CHECK-SAME: : memref<4xf32, strided<[3], offset: ?>>
+// CHECK: return %[[RES]]
+func.func @negative_strided_gather_with_dynamic_offset(
+ %base: memref<4x3xf32>,
+ %col: index,
+ %idxs: vector<2xindex>,
+ %mask: vector<2xi1>,
+ %pass_thru: vector<2xf32>) -> vector<2xf32> {
+ %c0 = arith.constant 0 : index
+ %sub = memref.subview %base[0, %col] [4, 1] [1, 1]
+ : memref<4x3xf32> to memref<4xf32, strided<[3], offset: ?>>
+ %0 = vector.gather %sub[%c0] [%idxs], %mask, %pass_thru
+ : memref<4xf32, strided<[3], offset: ?>>, vector<2xindex>,
+ vector<2xi1>, vector<2xf32> into vector<2xf32>
+ return %0 : vector<2xf32>
+}
More information about the Mlir-commits
mailing list