[Mlir-commits] [mlir] [mlir][vector][NFC] Document vector contract reshape helpers (PR #200544)

Fri May 29 23:00:07 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-mlir

Author: SeongJaePark (SeongjaeP)

<details>
<summary>Changes</summary>

## 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>`).
  

---
Full diff: https://github.com/llvm/llvm-project/pull/200544.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp (+35-4) 
- (modified) mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir (+31) 


``````````diff

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

``````````

</details>


https://github.com/llvm/llvm-project/pull/200544
