[Mlir-commits] [mlir] [mlir][x86] Support multiple result uses in AMX nanokernel (PR #236811)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 30 01:21:41 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Julian Oppermann (jopperm)
<details>
<summary>Changes</summary>
Lifts the restriction that the final (i.e. usually after the accumulation loop nest) contraction result shall only have one use.
This check was load-bearing only for selecting the block in which the contents of the tile registers are stored/shuffled/accumulated. Before, this was the block of the unique result _user_. I'm proposing here to take the block of the result _value_ instead, which is the parent block of the outermost loop currently being considered by the pattern. The pattern can still benefit from knowing that there's a single user of the result, e.g. the direct tile-store-path is still intact.
Assisted-by: Copilot
---
Full diff: https://github.com/llvm/llvm-project/pull/236811.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp (+23-19)
- (modified) mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir (+159)
``````````diff
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 600686258d988..6602e477271a3 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -45,10 +45,13 @@ static bool isZeroVectorConstant(Operation *op) {
}
// Recursively follows single-use values through scf.yield operations
-// and returns the first non-yield user result in the contraction chain.
-static Value contractionUsersAfterYield(Value v) {
- if (v.getNumUses() != 1)
+// and returns the first value that has more than one user, or is not used by a
+// yield operation. Returns nullptr if the chain is ultimately unused.
+static Value contractionResultAfterYield(Value v) {
+ if (v.hasNUses(0))
return nullptr;
+ if (v.hasNUsesOrMore(2))
+ return v;
OpOperand &use = *v.use_begin();
Operation *user = use.getOwner();
@@ -60,7 +63,7 @@ static Value contractionUsersAfterYield(Value v) {
Operation *parent = yield->getParentOp();
unsigned idx = use.getOperandNumber();
- return contractionUsersAfterYield(parent->getResult(idx));
+ return contractionResultAfterYield(parent->getResult(idx));
}
// Function to collapse the last two dimension (k or n, and vnni) to help the
@@ -212,7 +215,7 @@ static LogicalResult validateContractOps(OpBuilder &rewriter,
return failure();
}
- if (!contractionUsersAfterYield(contractOp.getResult()))
+ if (!contractionResultAfterYield(contractOp.getResult()))
return failure();
VectorType accTy = dyn_cast<VectorType>(contractOp.getAccType());
@@ -598,7 +601,7 @@ static std::pair<std::optional<SmallVector<Operation *>>,
getInPlaceAccAccesses(ArrayRef<vector::ContractionOp> ops, bool writeOnly) {
SmallVector<Operation *> reads, writes;
for (vector::ContractionOp op : ops) {
- Value result = contractionUsersAfterYield(op.getResult());
+ Value result = contractionResultAfterYield(op.getResult());
if (!result || !result.hasOneUse())
return {std::nullopt, std::nullopt};
Operation *write = *result.user_begin();
@@ -947,20 +950,21 @@ struct VectorContractToAMXDotProduct
Operation *accReadOp =
traceToVectorReadLikeParentOperation(contractOp.getAcc());
- // Only the contract result's first consumer is needed, not the final
- // store. This keeps the lowering independent of the epilogue ops (truncf,
- // bias add, ReLU, ...) that sit between the contraction and the write.
- Value resultChainEnd = contractionUsersAfterYield(contractOp.getResult());
+ // Find the value representing the final result of the accumulation by
+ // traversing a single-use chain through the loop levels. This keeps the
+ // lowering independent of the epilogue ops (truncf, bias add, ReLU, ...)
+ // that sit between the contraction and the write.
+ Value resultChainEnd = contractionResultAfterYield(contractOp.getResult());
if (!accReadOp || !resultChainEnd)
return rewriter.notifyMatchFailure(
contractOp, "The ACC operand of the vector.contract should be a "
- "transfer_read or a load. And, the result should have a "
- "single-use chain to its consumer.");
+ "transfer_read, a load, or a zero vector constant. And, "
+ "the computation should not be dead.");
bool isAccZeroVectorConstant = isZeroVectorConstant(accReadOp);
- Block *resultBlock = resultChainEnd.user_begin()->getBlock();
+ Block *resultBlock = resultChainEnd.getParentBlock();
Type ipType = rewriter.getBF16Type();
Type opType = rewriter.getF32Type();
@@ -1194,11 +1198,11 @@ struct VectorContractToAMXDotProduct
rewriter, loc, vectorType, resultBuffer, ValueRange{c0, c0}, padding,
map, inBounds);
- Value resultOp = contractionUsersAfterYield(contractOp.getResult());
- if (auto vecType = dyn_cast<VectorType>(resultOp.getType()))
+ Value result = contractionResultAfterYield(contractOp.getResult());
+ if (auto vecType = dyn_cast<VectorType>(result.getType()))
vecRow = vector::ShapeCastOp::create(rewriter, loc, vecType, vecRow);
- rewriter.replaceAllUsesWith(resultOp, vecRow);
+ rewriter.replaceAllUsesWith(result, vecRow);
return success();
}
@@ -1804,12 +1808,12 @@ struct VectorContractToAMXDotProduct
vector::ContractionOp contOp = ops[i];
Value vecRow = writeResults[i];
- Value resultWriteOp = contractionUsersAfterYield(contOp.getResult());
- if (auto vecType = dyn_cast<VectorType>(resultWriteOp.getType()))
+ Value result = contractionResultAfterYield(contOp.getResult());
+ if (auto vecType = dyn_cast<VectorType>(result.getType()))
vecRow = vector::ShapeCastOp::create(rewriter, loc, vecType,
writeResults[i]);
- rewriter.replaceAllUsesWith(resultWriteOp, vecRow);
+ rewriter.replaceAllUsesWith(result, vecRow);
}
return success();
diff --git a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
index 61b86a510f4d5..90db076bb68f6 100644
--- a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
+++ b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
@@ -108,6 +108,66 @@ module attributes {transform.with_named_sequence} {
// -----
+!vec = vector<16x16x4xi8>
+!acc = vector<16x16xi32>
+#multi_lhs = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#multi_rhs = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#multi_acc = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @single_contract_multiple_users(%lhs: memref<16x16x4xi8>, %rhs: memref<16x16x4xi8>, %acc: memref<16x16xi32>, %out_add: memref<16x16xi32>, %out_sub: memref<16x16xi32>) {
+ %c0 = arith.constant 0 : index
+ %pad_i8 = ub.poison : i8
+ %pad_i32 = ub.poison : i32
+ %bias = arith.constant dense<1> : !acc
+
+ %lhs_vec = vector.transfer_read %lhs[%c0, %c0, %c0], %pad_i8 {in_bounds = [true, true, true]} :
+ memref<16x16x4xi8>, !vec
+ %rhs_vec = vector.transfer_read %rhs[%c0, %c0, %c0], %pad_i8 {in_bounds = [true, true, true]} :
+ memref<16x16x4xi8>, !vec
+ %acc_vec = vector.transfer_read %acc[%c0, %c0], %pad_i32 {in_bounds = [true, true]} :
+ memref<16x16xi32>, !acc
+
+ %result = vector.contract {indexing_maps = [#multi_lhs, #multi_rhs, #multi_acc], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %lhs_vec, %rhs_vec, %acc_vec : !vec, !vec into !acc
+
+ %sum = arith.addi %result, %bias : !acc
+ %diff = arith.subi %result, %bias : !acc
+ vector.transfer_write %sum, %out_add[%c0, %c0] {in_bounds = [true, true]} :
+ !acc, memref<16x16xi32>
+ vector.transfer_write %diff, %out_sub[%c0, %c0] {in_bounds = [true, true]} :
+ !acc, memref<16x16xi32>
+
+ return
+}
+
+// CHECK-LABEL: @single_contract_multiple_users
+// CHECK-SAME: %[[ACC:[^:]+]]: memref<16x16xi32>, %[[OUT_ADD:[^:]+]]: memref<16x16xi32>, %[[OUT_SUB:[^:]+]]: memref<16x16xi32>
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: x86.amx.tile_load %[[ACC]][%[[C0]], %[[C0]]]
+// CHECK: %[[TILE:.+]] = x86.amx.tile_muli
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<16x16xi32>
+// CHECK: x86.amx.tile_store %[[BUF]][%[[C0]], %[[C0]]], %[[TILE]]
+// CHECK: %[[RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK: %[[SUM:.+]] = arith.addi %[[RES]], %[[BIAS:.+]] : vector<16x16xi32>
+// CHECK: %[[DIFF:.+]] = arith.subi %[[RES]], %[[BIAS]] : vector<16x16xi32>
+// CHECK: vector.transfer_write %[[SUM]], %[[OUT_ADD]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[DIFF]], %[[OUT_SUB]][%[[C0]], %[[C0]]]
+// CHECK: return
+// CHECK-NOT: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// The indices (i, vnni_index) of the two innermost dims of the operands map to
// i * vnni + vnni_index in the collapsed buffers.
@@ -2992,6 +3052,105 @@ module attributes {transform.with_named_sequence} {
// -----
+!vecA = vector<16x64xi8>
+!vecB = vector<64x16xi8>
+!vecC = vector<16x16xi32>
+!loopMemrefA = memref<16x64xi8, strided<[128, 1], offset: ?>>
+!loopMemrefB = memref<64x32xi8, strided<[32, 1], offset: ?>>
+#loop_lhs = affine_map<(d0, d1, d2) -> (d0, d2)>
+#loop_rhs = affine_map<(d0, d1, d2) -> (d2, d1)>
+#loop_acc = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @loop_contract_multiple_users(%lhs: memref<16x128xi8>, %rhs: memref<128x32xi8>, %acc: memref<16x32xi32>, %out_add: memref<16x32xi32>, %out_sub: memref<16x32xi32>) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %pad_i8 = ub.poison : i8
+ %pad_i32 = ub.poison : i32
+ %bias = arith.constant dense<1> : !vecC
+
+ %init0 = vector.transfer_read %acc[%c0, %c0], %pad_i32 {in_bounds = [true, true]} :
+ memref<16x32xi32>, !vecC
+ %init1 = vector.transfer_read %acc[%c0, %c16], %pad_i32 {in_bounds = [true, true]} :
+ memref<16x32xi32>, !vecC
+
+ %results:2 = scf.for %k = %c0 to %c128 step %c64 iter_args(%iter0 = %init0, %iter1 = %init1) -> (!vecC, !vecC) {
+ %lhs_view = memref.subview %lhs[%c0, %k] [16, 64] [1, 1] :
+ memref<16x128xi8> to !loopMemrefA
+ %rhs_view = memref.subview %rhs[%k, %c0] [64, 32] [1, 1] :
+ memref<128x32xi8> to !loopMemrefB
+ %lhs_vec = vector.transfer_read %lhs_view[%c0, %c0], %pad_i8 {in_bounds = [true, true]} :
+ !loopMemrefA, !vecA
+ %rhs0 = vector.transfer_read %rhs_view[%c0, %c0], %pad_i8 {in_bounds = [true, true]} :
+ !loopMemrefB, !vecB
+ %rhs1 = vector.transfer_read %rhs_view[%c0, %c16], %pad_i8 {in_bounds = [true, true]} :
+ !loopMemrefB, !vecB
+
+ %result0 = vector.contract {indexing_maps = [#loop_lhs, #loop_rhs, #loop_acc], iterator_types =
+ ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %lhs_vec, %rhs0, %iter0 : !vecA, !vecB into !vecC
+ %result1 = vector.contract {indexing_maps = [#loop_lhs, #loop_rhs, #loop_acc], iterator_types =
+ ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %lhs_vec, %rhs1, %iter1 : !vecA, !vecB into !vecC
+
+ scf.yield %result0, %result1 : !vecC, !vecC
+ }
+
+ %sum = arith.addi %results#0, %bias : !vecC
+ %diff = arith.subi %results#0, %bias : !vecC
+ %sum1 = arith.addi %results#1, %bias : !vecC
+ %diff1 = arith.subi %results#1, %bias : !vecC
+ vector.transfer_write %sum, %out_add[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, memref<16x32xi32>
+ vector.transfer_write %diff, %out_sub[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, memref<16x32xi32>
+ vector.transfer_write %sum1, %out_add[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, memref<16x32xi32>
+ vector.transfer_write %diff1, %out_sub[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, memref<16x32xi32>
+
+ return
+}
+
+// CHECK-LABEL: @loop_contract_multiple_users
+// CHECK-SAME: %[[ACC:[^:]+]]: memref<16x32xi32>, %[[OUT_ADD:[^:]+]]: memref<16x32xi32>, %[[OUT_SUB:[^:]+]]: memref<16x32xi32>
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>)
+// CHECK-COUNT-2: x86.amx.tile_muli
+// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>)
+// CHECK-COUNT-2: x86.amx.tile_muli
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<16x32xi32>
+// CHECK: x86.amx.tile_store %[[BUF]][%[[C0]], %[[C0]]], %[[TILES]]#0
+// CHECK: x86.amx.tile_store %[[BUF]][%[[C0]], %[[C16]]], %[[TILES]]#1
+// CHECK: scf.for
+// CHECK-COUNT-2: vector.shuffle {{.*}} : vector<16xi32>, vector<16xi32>
+// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK: %[[RES1:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C16]]]
+// CHECK: %[[SUM0:.+]] = arith.addi %[[RES0]], %[[BIAS:.+]] : vector<16x16xi32>
+// CHECK: %[[DIFF0:.+]] = arith.subi %[[RES0]], %[[BIAS]] : vector<16x16xi32>
+// CHECK: %[[SUM1:.+]] = arith.addi %[[RES1]], %[[BIAS]] : vector<16x16xi32>
+// CHECK: %[[DIFF1:.+]] = arith.subi %[[RES1]], %[[BIAS]] : vector<16x16xi32>
+// CHECK: vector.transfer_write %[[SUM0]], %[[OUT_ADD]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[DIFF0]], %[[OUT_SUB]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[SUM1]], %[[OUT_ADD]][%[[C0]], %[[C16]]]
+// CHECK: vector.transfer_write %[[DIFF1]], %[[OUT_SUB]][%[[C0]], %[[C16]]]
+// CHECK: return
+// CHECK-NOT: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
!vecA = vector<1x16x16x4xi8>
!vecB = vector<1x16x16x4xi8>
!vecC = vector<16x16xi32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/236811
More information about the Mlir-commits
mailing list