[Mlir-commits] [mlir] [mlir][x86] Fix AMX loop carried accumulation (PR #231170)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 12 06:46:41 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Adam Siemieniuk (adam-smnk)
<details>
<summary>Changes</summary>
Fixes the AMX contraction lowering of reduction loops whose accumulators are carried in by an enclosing loop.
The values accumulated by the enclosing loop were dropped and replaced with the initial accumulator, which produced wrong results.
Assisted-by: Claude
---
Patch is 37.74 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/231170.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp (+58-2)
- (modified) mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir (+699)
``````````diff
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 62d184ed7d20f..23ea062c649a5 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -259,6 +259,22 @@ static LogicalResult validateContractOps(OpBuilder &rewriter,
return success();
}
+// Traces an accumulator back through the iter_args of the given loops and
+// returns the first value that is not one of them, e.g., the init value of the
+// outermost loop or a value computed within the loops.
+static Value getLoopNestInit(Value acc, ArrayRef<scf::ForOp> loops) {
+ while (auto arg = dyn_cast<BlockArgument>(acc)) {
+ auto loop = dyn_cast_if_present<scf::ForOp>(arg.getOwner()->getParentOp());
+ if (!loop || !llvm::is_contained(loops, loop))
+ break;
+ OpOperand *init = loop.getTiedLoopInit(arg);
+ if (!init)
+ break;
+ acc = init->get();
+ }
+ return acc;
+}
+
// Returns the loop index position to get mapped during the
// MemRef type clone.
static unsigned getIndexPosition(Value operand, scf::ForOp loop) {
@@ -1165,6 +1181,10 @@ struct VectorContractToAMXDotProduct
// Retrive all the contaction operation within the loop.
SmallVector<vector::ContractionOp> ops;
+ // Accumulators whose initial value passed into the rewritten reduction
+ // loops is neither read from memory nor zero, e.g. an iter_arg of an
+ // enclosing loop. Paired with the index of their contraction in `ops`.
+ SmallVector<std::pair<size_t, Value>> carriedAccs;
for (Operation &op : loopLists[0].getBody()->getOperations()) {
if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
@@ -1192,10 +1212,31 @@ struct VectorContractToAMXDotProduct
contractOp, "The input acc should be a zero constant or "
"transfer_read for all contracts");
+ Value accInit = getLoopNestInit(contract.getAcc(), loopLists);
+ Operation *accInitOp = accInit.getDefiningOp();
+ bool carried = !isa_and_present<TransferReadOp, LoadOp>(accInitOp) &&
+ !isZeroVectorConstant(accInitOp);
+ if (carried) {
+ if (loopLists.back().getRegion().isAncestor(
+ accInit.getParentRegion()))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The accumulator should enter the reduction loops "
+ "through their init args.");
+ carriedAccs.emplace_back(ops.size(), accInit);
+ }
+
ops.push_back(contract);
}
}
+ // The accumulators read from memory are added back to all the results at
+ // once, so they can't be mixed with carried ones.
+ if (!isAccZeroVectorConstant && !carriedAccs.empty() &&
+ carriedAccs.size() != ops.size())
+ return rewriter.notifyMatchFailure(
+ contractOp, "The accumulators read from memory can't be mixed with "
+ "carried in accumulators.");
+
if (!isVnni) {
unsigned int pairCount = 0;
for (size_t j = 0; j < ops.size(); j++) {
@@ -1537,7 +1578,10 @@ struct VectorContractToAMXDotProduct
auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
auto nBound = arith::ConstantIndexOp::create(rewriter, loc, N);
- if (!isVnni || !isAccZeroVectorConstant) {
+ // Carried-in accumulators are added to the results instead.
+ bool addAccFromMemory = !isAccZeroVectorConstant && carriedAccs.empty();
+
+ if (!isVnni || addAccFromMemory) {
// Create a loop that iterates over the MxN memerf, retrives two rows +
// shuffle them, add up the C element values and stores them to temp
// buffer.
@@ -1568,7 +1612,7 @@ struct VectorContractToAMXDotProduct
15, 28, 29, 30, 31});
}
- if (!isAccZeroVectorConstant) {
+ if (addAccFromMemory) {
indicesAcc[indicesAcc.size() - 2] = iv;
indicesAcc[indicesAcc.size() - 1] = c0;
@@ -1629,6 +1673,18 @@ struct VectorContractToAMXDotProduct
}
}
+ for (auto [idx, acc] : carriedAccs) {
+ Type resultType = writeResults[idx].getType();
+ if (acc.getType() != resultType)
+ acc = vector::ShapeCastOp::create(rewriter, loc, resultType, acc);
+ if (isa<FloatType>(opType))
+ writeResults[idx] =
+ arith::AddFOp::create(rewriter, loc, writeResults[idx], acc);
+ else
+ writeResults[idx] =
+ arith::AddIOp::create(rewriter, loc, writeResults[idx], acc);
+ }
+
// Replace use of vector.contract with dot-products.
for (size_t i = 0; i < ops.size(); i++) {
vector::ContractionOp contOp = ops[i];
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 b7ce24b54134f..f1d2b969ae303 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
@@ -1397,6 +1397,705 @@ module attributes {transform.with_named_sequence} {
// -----
+// The reduction loop accumulators are carried in by an enclosing loop.
+// The carried values are added to the AMX results instead of reloading C.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @matmul_int8_loop_carried_acc(
+ %arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>,
+ %bias: !vecC) {
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %r:2 = scf.for %b = %c0 to %c2 step %c1 iter_args(%acc0 = %2, %acc1 = %3) -> (!vecC, !vecC) {
+ %4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %acc0, %arg7 = %acc1) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ %10 = arith.addi %4#0, %bias : !vecC
+ %11 = arith.addi %4#1, %bias : !vecC
+ scf.yield %10, %11 : !vecC, !vecC
+ }
+
+ vector.transfer_write %r#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %r#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_carried_acc
+// CHECK: scf.for {{.*}} iter_args(%[[ACC0:arg[0-9]+]] = %{{.+}}, %[[ACC1:arg[0-9]+]] = %{{.+}})
+// CHECK-COUNT-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK: x86.amx.tile_store %[[BUF:.+]][%{{.+}}, %{{.+}}], %[[TILES]]#0
+// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#1
+// CHECK-NOT: vector.load
+// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[RES1:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[SUM0:.+]] = arith.addi %[[RES0]], %[[ACC0]] : vector<16x16xi32>
+// CHECK: %[[SUM1:.+]] = arith.addi %[[RES1]], %[[ACC1]] : vector<16x16xi32>
+// CHECK: arith.addi %[[SUM0]], %{{.+}} : vector<16x16xi32>
+// CHECK: arith.addi %[[SUM1]], %{{.+}} : vector<16x16xi32>
+// 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 reduction loop accumulators with a unit batch dim are carried in by an
+// enclosing loop.
+// The carried values are reshaped and added to the AMX results.
+
+!vecAB = vector<1x16x16x4xi8>
+!vecC = vector<1x16x16xi32>
+!memrefA = memref<1x16x16x4xi8, strided<[16384, 256, 4, 1], offset: ?>>
+!memrefB = memref<1x16x32x4xi8, strided<[32768, 512, 4, 1], offset: ?>>
+!memrefC = memref<1x16x32xi32, strided<[8192, 128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
+#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
+#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d3)>
+
+func.func @batch_matmul_int8_loop_carried_acc(
+ %arg0: memref<16x64x64x4xi8>, %arg1: memref<16x64x128x4xi8>, %arg2: memref<16x64x128xi32>,
+ %bias: !vecC) {
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+ scf.for %arg5 = %c0 to %c16 step %c1 {
+
+ %subview = memref.subview %arg2[%arg5, %arg3, %arg4] [1, 16, 32] [1, 1, 1] :
+ memref<16x64x128xi32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c0, %c16], %0 {in_bounds = [true, true, true]} :
+ !memrefC, !vecC
+
+ %r:2 = scf.for %b = %c0 to %c2 step %c1 iter_args(%acc0 = %2, %acc1 = %3) -> (!vecC, !vecC) {
+ %4:2 = scf.for %arg6 = %c0 to %c64 step %c16 iter_args(%arg7 = %acc0, %arg8 = %acc1) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg5, %arg3, %arg6, 0] [1, 16, 16, 4] [1, 1, 1, 1] :
+ memref<16x64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg6, %arg4, 0] [1, 16, 32, 4] [1, 1, 1, 1] :
+ memref<16x64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["parallel", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg7 {unroll_shape = array<i64: 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["parallel", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg8 {unroll_shape = array<i64: 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ %10 = arith.addi %4#0, %bias : !vecC
+ %11 = arith.addi %4#1, %bias : !vecC
+ scf.yield %10, %11 : !vecC, !vecC
+ }
+
+ vector.transfer_write %r#1, %subview[%c0, %c0, %c16] {in_bounds = [true, true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %r#0, %subview[%c0, %c0, %c0] {in_bounds = [true, true, true]} :
+ !vecC, !memrefC
+ }
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @batch_matmul_int8_loop_carried_acc
+// CHECK: scf.for {{.*}} iter_args(%[[ACC0:arg[0-9]+]] = %{{.+}}, %[[ACC1:arg[0-9]+]] = %{{.+}})
+// CHECK-COUNT-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK: x86.amx.tile_store %[[BUF:.+]][%{{.+}}, %{{.+}}], %[[TILES]]#0
+// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#1
+// CHECK-NOT: vector.load
+// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[RES1:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[CAST0:.+]] = vector.shape_cast %[[ACC0]] : vector<1x16x16xi32> to vector<16x16xi32>
+// CHECK: arith.addi %[[RES0]], %[[CAST0]] : vector<16x16xi32>
+// CHECK: %[[CAST1:.+]] = vector.shape_cast %[[ACC1]] : vector<1x16x16xi32> to vector<16x16xi32>
+// CHECK: arith.addi %[[RES1]], %[[CAST1]] : vector<16x16xi32>
+// 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 batch-reduce and K reduction loops are nested in a loop that carries in
+// the accumulators, starting from zero.
+// The carried values are added to the AMX results.
+
+!vecAB = vector<1x16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<1x16x16x4xi8, strided<[16384, 256, 4, 1], offset: ?>>
+!memrefB = memref<1x16x32x4xi8, strided<[32768, 512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
+#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
+#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
+
+func.func @brgemm_int8_loop_carried_zero_acc(
+ %arg0: memref<16x64x64x4xi8>, %arg1: memref<16x64x128x4xi8>, %arg2: memref<64x128xi32>,
+ %bias: !vecC) {
+ %cst = arith.constant dense<0> : !vecC
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> to !memrefC
+
+ %r:2 = scf.for %b = %c0 to %c2 step %c1 iter_args(%acc0 = %cst, %acc1 = %cst) -> (!vecC, !vecC) {
+ %2:2 = scf.for %arg5 = %c0 to %c16 step %c1 iter_args(%arg6 = %acc0, %arg7 = %acc1) -> (!vecC, !vecC) {
+ %3:2 = scf.for %arg8 = %c0 to %c64 step %c16 iter_args(%arg9 = %arg6, %arg10 = %arg7) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg5, %arg3, %arg8, 0] [1, 16, 16, 4] [1, 1, 1, 1] :
+ memref<16x64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg8, %arg4, 0] [1, 16, 32, 4] [1, 1, 1, 1] :
+ memref<16x64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg9 {unroll_shape = array<i64: 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg10 {unroll_shape = array<i64: 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+ scf.yield %3#0, %3#1 : !vecC, !vecC
+ }
+
+ %10 = arith.addi %2#0, %bias : !vecC
+ %11 = arith.addi %2#1, %bias : !vecC
+ scf.yield %10, %11 : !vecC, !vecC
+ }
+
+ vector.transfer_write %r#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %r#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @brgemm_int8_loop_carried_zero_acc
+// CHECK: scf.for {{.*}} iter_args(%[[ACC0:arg[0-9]+]] = %{{.+}}, %[[ACC1:arg[0-9]+]] = %{{.+}})
+// CHECK-COUNT-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK-COUNT-2: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK: x86.amx.tile_store %[[BUF:.+]][%{{.+}}, %{{.+}}], %{{.+}}#0
+// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %{{.+}}#1
+// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[RES1:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: arith.addi %[[RES0]], %[[ACC0]] : vector<16x16xi32>
+// CHECK: arith.addi %[[RES1]], %[[ACC1]] : vector<16x16xi32>
+// 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
+ }
+}
+
+// -----
+
+// Only one of the reduction loop accumulators is carried in by an enclosing
+// loop, the other one starts from zero.
+// Only the carried value is added to its AMX result.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @m...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/231170
More information about the Mlir-commits
mailing list