[Mlir-commits] [mlir] [mlir][vector][NFC] Document vector contract reshape helpers (PR #200544)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri May 29 22:59:34 PDT 2026
https://github.com/SeongjaeP created https://github.com/llvm/llvm-project/pull/200544
## Summary
Replace the bare `// TODO` placeholders above `reshapeLoad` and `reshapeStore`
in `LowerVectorContract.cpp` with doc comments that describe each function's
semantics, the `index == -1` / `index == 0` / `index > 0` cases, and how
they are used by `ContractionOpLowering::lowerParallel` and `lowerReduction`
to peel one iterator at a time from a `vector.contract`.
Add a LIT test that exercises the recursive unroll path of `reshapeLoad`.
The existing parallelarith tests use shapes covered by the specialized
`ContractOpTo{OuterProduct,Dot,Elementwise}` patterns and never reach the
generic `lowerParallel` / `reshapeLoad` path. The new test maps the
parallel iterator to a non-leading dim of `vector<3x2xf32>`, falling
through to the generic path and forcing `reshapeLoad` to recursively
unroll the leading dim of size 3 to build per-lane `vector<3xf32>` slabs
before each inner reduction.
No functional change.
## Test plan
- `ninja check-mlir-dialect-vector` → 98/98 passed.
- New test `@parallel_contract_lowering_non_unit_parallel` verified to
exercise the `reshapeLoad`/`reshapeStore` recursive path by inspecting
the lowered IR (per-lane `vector<3xf32>` slab construction followed by
`arith.mulf` + `vector.reduction<add>`).
>From 840f2fc5a66586eadbb5d8e29aa1b57f9288a3bb Mon Sep 17 00:00:00 2001
From: "sjae.park" <dev at opt-ai.kr>
Date: Fri, 29 May 2026 06:17:09 +0000
Subject: [PATCH] [mlir][vector][NFC] Document vector contract reshape helpers
Replace the bare `// TODO` placeholders above `reshapeLoad` and
`reshapeStore` in `LowerVectorContract.cpp` with doc comments that
describe each function's semantics, the `index == -1` / `index == 0` /
`index > 0` cases, and how they are used by
`ContractionOpLowering::lowerParallel` and `lowerReduction` to peel
one iterator at a time from a `vector.contract`.
Add a LIT test that exercises the recursive unroll path of
`reshapeLoad`. The existing parallelarith tests use shapes covered by
the specialized `ContractOpTo{OuterProduct,Dot,Elementwise}` patterns
and never reach the generic `lowerParallel` / `reshapeLoad` path. The
new test maps the parallel iterator to a non-leading dim of
`vector<3x2xf32>`, falling through to the generic path and forcing
`reshapeLoad` to recursively unroll the leading dim of size 3 to
build per-lane `vector<3xf32>` slabs before each inner reduction.
No functional change.
---
.../Vector/Transforms/LowerVectorContract.cpp | 39 +++++++++++++++++--
...contract-to-parallel-arith-transforms.mlir | 31 +++++++++++++++
2 files changed, 66 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index eaf7bb8109514..0ba08881522dd 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -71,8 +71,23 @@ static AffineMap adjustMap(AffineMap map, int64_t index,
return AffineMap::get(map.getNumDims() - 1, 0, results, ctx);
}
-// Helper method to possibly drop a dimension in a load.
-// TODO
+/// Returns a value of the same element type as `val`, but with the dimension
+/// at position `index` dropped, by selecting the slab at coordinate `pos`
+/// along that dimension. `val` must have type `type`.
+///
+/// The `index == -1` case is treated as a pass-through: `val` is returned
+/// unchanged. This supports contract operands whose indexing map does not
+/// reference the iterator currently being unrolled (e.g. a broadcast operand
+/// or a unit-size reduction dim that only appears on one side).
+///
+/// For `index == 0` the result is a single `vector.extract %val[pos]`. For
+/// `index > 0` the function recursively unrolls the leading dimension and
+/// drops dimension `index - 1` from each unrolled sub-slab, then reassembles
+/// the slabs into a result whose `index`-th dimension has been removed.
+///
+/// Used by `ContractionOpLowering::lowerParallel` and `lowerReduction` to
+/// pick out the `pos`-th slab of an operand along the iterator being
+/// unrolled, before recursing into a lower-rank `vector.contract`.
static Value reshapeLoad(Location loc, Value val, VectorType type,
int64_t index, int64_t pos,
PatternRewriter &rewriter) {
@@ -96,8 +111,24 @@ static Value reshapeLoad(Location loc, Value val, VectorType type,
return result;
}
-// Helper method to possibly drop a dimension in a store.
-// TODO
+/// Dual of `reshapeLoad`. Inserts `val` (whose type matches `type` with the
+/// dimension at position `index` dropped) back into the accumulator `result`
+/// along dimension `index` at position `pos`. `result` must have type `type`.
+/// All other slabs of `result` are preserved.
+///
+/// The `index == -1` case is treated as a pass-through: `val` is returned
+/// directly, mirroring the `reshapeLoad` convention. This is only meaningful
+/// when the iterator being unrolled has unit size (one assignment overwrites
+/// the entire accumulator), which is also how `lowerParallel` invokes it.
+///
+/// For `index == 0` the result is a single `vector.insert %val, %result
+/// [pos]`. For `index > 0` the function recursively unrolls the leading
+/// dimension of both `result` and `val`, inserts each pair into dimension
+/// `index - 1` of the corresponding sub-slab, and reassembles the slabs.
+///
+/// Used by `ContractionOpLowering::lowerParallel` to write the result of
+/// each unrolled lower-dimensional `vector.contract` back into the
+/// accumulator along the iterator that is being unrolled.
static Value reshapeStore(Location loc, Value val, Value result,
VectorType type, int64_t index, int64_t pos,
PatternRewriter &rewriter) {
diff --git a/mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir
index e93c5a08bdc7c..fb89f042a3a5a 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir
@@ -51,6 +51,37 @@ func.func @parallel_contract_lowering_scalar(%arg0: vector<1x1xf32>, %arg1: vect
return %0 : f32
}
+// The parallel iterator (d0, size 2) maps to the *non-leading* dim of LHS/RHS,
+// exercising the recursive `reshapeLoad` path: it unrolls the leading dim of
+// size 3 to extract a per-lane `vector<3xf32>` slab for each parallel position,
+// then reduces it back to a scalar that is stored into the result via
+// `reshapeStore`.
+//
+// CHECK-LABEL: func @parallel_contract_lowering_non_unit_parallel(
+// CHECK-SAME: %[[LHS:.+]]: vector<3x2xf32>, %[[RHS:.+]]: vector<3x2xf32>, %[[ACC:.+]]: vector<2xf32>
+// CHECK: %[[LSLAB0:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[RSLAB0:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[ACC0:.+]] = vector.extract %[[ACC]][0] : f32 from vector<2xf32>
+// CHECK: %[[MUL0:.+]] = arith.mulf %[[LSLAB0]], %[[RSLAB0]] : vector<3xf32>
+// CHECK: %[[RED0:.+]] = vector.reduction <add>, %[[MUL0]], %[[ACC0]] : vector<3xf32> into f32
+// CHECK: %[[OUT0:.+]] = vector.insert %[[RED0]], %{{.*}} [0] : f32 into vector<2xf32>
+// CHECK: %[[LSLAB1:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[RSLAB1:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[ACC1:.+]] = vector.extract %[[ACC]][1] : f32 from vector<2xf32>
+// CHECK: %[[MUL1:.+]] = arith.mulf %[[LSLAB1]], %[[RSLAB1]] : vector<3xf32>
+// CHECK: %[[RED1:.+]] = vector.reduction <add>, %[[MUL1]], %[[ACC1]] : vector<3xf32> into f32
+// CHECK: %[[OUT1:.+]] = vector.insert %[[RED1]], %[[OUT0]] [1] : f32 into vector<2xf32>
+// CHECK: return %[[OUT1]] : vector<2xf32>
+func.func @parallel_contract_lowering_non_unit_parallel(%arg0: vector<3x2xf32>, %arg1: vector<3x2xf32>, %arg2: vector<2xf32>) -> vector<2xf32> {
+ %0 = vector.contract {
+ indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>,
+ affine_map<(d0, d1) -> (d1, d0)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"], kind = #vector.kind<add>
+ } %arg0, %arg1, %arg2 : vector<3x2xf32>, vector<3x2xf32> into vector<2xf32>
+ return %0 : vector<2xf32>
+}
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%f = transform.structured.match ops{["func.func"]} in %module_op
More information about the Mlir-commits
mailing list