[Mlir-commits] [mlir] 4881326 - [mlir][vector][NFC] Document and simplify vector contract reshape helpers (#200544)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 11 16:46:39 PDT 2026
Author: SeongJaePark
Date: 2026-06-11T16:46:35-07:00
New Revision: 48813264376be9868bec78fa1653bc9003db512f
URL: https://github.com/llvm/llvm-project/commit/48813264376be9868bec78fa1653bc9003db512f
DIFF: https://github.com/llvm/llvm-project/commit/48813264376be9868bec78fa1653bc9003db512f.diff
LOG: [mlir][vector][NFC] Document and simplify vector contract reshape helpers (#200544)
Replace the bare `// TODO` placeholders above `reshapeLoad` and
`reshapeStore`
in `LowerVectorContract.cpp` with doc comments describing each helper's
semantics (the `index == -1` / `index == 0` / `index > 0` cases) and
short
before/after IR examples.
Also drop the redundant `VectorType` parameter from both helpers; the
type is
now inferred from the value operand (`val` for `reshapeLoad`, `result`
for
`reshapeStore`).
Add a LIT test that exercises the recursive `reshapeLoad` unroll path by
mapping a parallel iterator to a non-leading dimension, which falls
through
to the generic `lowerParallel` / `reshapeLoad` path.
No functional change.
---------
Co-authored-by: sjae.park <dev at opt-ai.kr>
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index eaf7bb8109514..f6fff0ffe323a 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -71,10 +71,20 @@ 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
-static Value reshapeLoad(Location loc, Value val, VectorType type,
- int64_t index, int64_t pos,
+/// 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, int64_t index, int64_t pos,
PatternRewriter &rewriter) {
if (index == -1)
return val;
@@ -84,23 +94,31 @@ static Value reshapeLoad(Location loc, Value val, VectorType type,
return vector::ExtractOp::create(rewriter, loc, val, pos);
// Unroll leading dimensions.
- VectorType vType = VectorType::Builder(type).dropDim(0);
+ VectorType type = cast<VectorType>(val.getType());
VectorType resType = VectorType::Builder(type).dropDim(index);
Value result = arith::ConstantOp::create(rewriter, loc, resType,
rewriter.getZeroAttr(resType));
for (int64_t d = 0, e = resType.getDimSize(0); d < e; d++) {
Value ext = vector::ExtractOp::create(rewriter, loc, val, d);
- Value load = reshapeLoad(loc, ext, vType, index - 1, pos, rewriter);
+ Value load = reshapeLoad(loc, ext, index - 1, pos, rewriter);
result = vector::InsertOp::create(rewriter, loc, load, result, d);
}
return result;
}
-// Helper method to possibly drop a dimension in a store.
-// TODO
-static Value reshapeStore(Location loc, Value val, Value result,
- VectorType type, int64_t index, int64_t pos,
- PatternRewriter &rewriter) {
+/// 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, int64_t index,
+ int64_t pos, PatternRewriter &rewriter) {
// Unmodified?
if (index == -1)
return val;
@@ -109,11 +127,11 @@ static Value reshapeStore(Location loc, Value val, Value result,
return vector::InsertOp::create(rewriter, loc, val, result, pos);
// Unroll leading dimensions.
- VectorType vType = VectorType::Builder(type).dropDim(0);
+ VectorType type = cast<VectorType>(result.getType());
for (int64_t d = 0, e = type.getDimSize(0); d < e; d++) {
Value ext = vector::ExtractOp::create(rewriter, loc, result, d);
Value ins = vector::ExtractOp::create(rewriter, loc, val, d);
- Value sto = reshapeStore(loc, ins, ext, vType, index - 1, pos, rewriter);
+ Value sto = reshapeStore(loc, ins, ext, index - 1, pos, rewriter);
result = vector::InsertOp::create(rewriter, loc, sto, result, d);
}
return result;
@@ -1049,21 +1067,20 @@ FailureOr<Value> ContractionOpLowering::lowerParallel(PatternRewriter &rewriter,
rewriter.getZeroAttr(resType));
for (int64_t d = 0; d < dimSize; ++d) {
- auto lhs = reshapeLoad(loc, op.getLhs(), lhsType, lhsIndex, d, rewriter);
- auto rhs = reshapeLoad(loc, op.getRhs(), rhsType, rhsIndex, d, rewriter);
- auto acc = reshapeLoad(loc, op.getAcc(), resType, resIndex, d, rewriter);
+ auto lhs = reshapeLoad(loc, op.getLhs(), lhsIndex, d, rewriter);
+ auto rhs = reshapeLoad(loc, op.getRhs(), rhsIndex, d, rewriter);
+ auto acc = reshapeLoad(loc, op.getAcc(), resIndex, d, rewriter);
Value lowMask;
if (mask)
- lowMask = reshapeLoad(loc, mask, cast<VectorType>(mask.getType()),
- iterIndex, d, rewriter);
+ lowMask = reshapeLoad(loc, mask, iterIndex, d, rewriter);
Operation *lowContract =
vector::ContractionOp::create(rewriter, loc, lhs, rhs, acc, lowAffine,
lowIter, op.getKind(), op.getFastmath());
lowContract = maskOperation(rewriter, lowContract, lowMask);
- result = reshapeStore(loc, lowContract->getResult(0), result, resType,
- resIndex, d, rewriter);
+ result = reshapeStore(loc, lowContract->getResult(0), result, resIndex, d,
+ rewriter);
}
return result;
}
@@ -1131,12 +1148,11 @@ FailureOr<Value> ContractionOpLowering::lowerReduction(
// the sum of all reductions is computed.
Value result = op.getAcc();
for (int64_t d = 0; d < dimSize; ++d) {
- auto lhs = reshapeLoad(loc, op.getLhs(), lhsType, lhsIndex, d, rewriter);
- auto rhs = reshapeLoad(loc, op.getRhs(), rhsType, rhsIndex, d, rewriter);
+ auto lhs = reshapeLoad(loc, op.getLhs(), lhsIndex, d, rewriter);
+ auto rhs = reshapeLoad(loc, op.getRhs(), rhsIndex, d, rewriter);
Value newMask;
if (mask)
- newMask = reshapeLoad(loc, mask, cast<VectorType>(mask.getType()),
- iterIndex, d, rewriter);
+ newMask = reshapeLoad(loc, mask, iterIndex, d, rewriter);
Operation *newContract = vector::ContractionOp::create(
rewriter, loc, lhs, rhs, result, lowAffine, lowIter, op.getKind(),
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..c5e116d326a78 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>` sub-vector 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: %[[LSUB0:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[RSUB0:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[ACC0:.+]] = vector.extract %[[ACC]][0] : f32 from vector<2xf32>
+// CHECK: %[[MUL0:.+]] = arith.mulf %[[LSUB0]], %[[RSUB0]] : vector<3xf32>
+// CHECK: %[[RED0:.+]] = vector.reduction <add>, %[[MUL0]], %[[ACC0]] : vector<3xf32> into f32
+// CHECK: %[[OUT0:.+]] = vector.insert %[[RED0]], %{{.*}} [0] : f32 into vector<2xf32>
+// CHECK: %[[LSUB1:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[RSUB1:.+]] = vector.insert %{{.*}}, %{{.*}} [2] : f32 into vector<3xf32>
+// CHECK: %[[ACC1:.+]] = vector.extract %[[ACC]][1] : f32 from vector<2xf32>
+// CHECK: %[[MUL1:.+]] = arith.mulf %[[LSUB1]], %[[RSUB1]] : 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