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

[llvmlistbot at llvm.org]

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

>From eea7468a02f7adc2572ead4b2fd550b6e13bdd80 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 and its `index == -1` / `index == 0`
/ `index > 0` cases, with short before/after IR examples.

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>` sub-vectors before each inner reduction.

No functional change.
---
 .../Vector/Transforms/LowerVectorContract.cpp | 28 ++++++++++++++---
 ...contract-to-parallel-arith-transforms.mlir | 31 +++++++++++++++++++
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index eaf7bb8109514..7a2e63b230d94 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -71,8 +71,19 @@ 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 `val` with the dimension at position `index` dropped by indexing
+/// that dimension with `pos`.
+///
+/// If `index == -1`, returns `val` unchanged. If `index == 0`, the result is
+/// a single `vector.extract %val[pos]`.
+///
+/// Example (`index == 0`): extract the sub-vector at `pos` along the leading
+/// dimension.
+///   // val : vector<4x8xf32>, pos = 2
+///   %res = vector.extract %val[2] : vector<8xf32> from vector<4x8xf32>
+///
+/// For `index > 0`, recursively applies the same drop to each sub-vector of
+/// the leading dimension and reassembles the result.
 static Value reshapeLoad(Location loc, Value val, VectorType type,
                          int64_t index, int64_t pos,
                          PatternRewriter &rewriter) {
@@ -96,8 +107,17 @@ static Value reshapeLoad(Location loc, Value val, VectorType type,
   return result;
 }
 
-// Helper method to possibly drop a dimension in a store.
-// TODO
+/// Inserts `val` into `result` at position `pos` along dimension `index`.
+///
+/// This is the inverse of `reshapeLoad`. If `index == -1`, returns `val`. If
+/// `index == 0`, the result is a single `vector.insert %val, %result [pos]`.
+///
+/// Example (`index == 0`): insert `val` at `pos` along the leading dimension.
+///   // val : vector<4xf32>, acc : vector<2x4xf32>, pos = 1
+///   %res = vector.insert %val, %acc [1] : vector<4xf32> into vector<2x4xf32>
+///
+/// For `index > 0`, recursively applies the same insertion to each sub-vector
+/// of the leading dimension and reassembles the result.
 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