[Mlir-commits] [mlir] [mlir][x86] Adds extra pattern for AMX lowering (PR #213228)
Arun Thangamani
llvmlistbot at llvm.org
Fri Jul 31 07:44:49 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/213228
>From f889774c3e00fffef7c00c4a30a6c248c2cf3c6a Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Fri, 31 Jul 2026 02:02:53 -0700
Subject: [PATCH 1/3] adds extra pattern for AMX lowering
---
.../VectorContractToAMXDotProduct.cpp | 35 ++++++++++++-------
1 file changed, 23 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 238d03ce4f6af..28b8542d79cfb 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -495,14 +495,16 @@ static SmallVector<Value> createTileZeros(OpBuilder &rewriter, Location loc,
static Value getIndxToLoadStoreFromPckBuffer(
OpBuilder &rewriter, Location loc, Value ivInnerLoop, Value ivOuterLoop,
bool isInnerLoopUBHasOddQuot, bool isInnerLoopUBLarger, bool pack,
- unsigned int blockingFactor) {
+ Value blockStride) {
Value c2 = arith::ConstantIndexOp::create(rewriter, loc, 2);
- Value packOffset =
- arith::ConstantIndexOp::create(rewriter, loc, (16 * blockingFactor));
+ // `blockStride` is the reduction (K) loop step, i.e. the amount by which the
+ // induction variable advances for one K-block. Dividing the induction value
+ // by it yields the K-block index regardless of whether the loop counts
+ // K-elements (step == 16*blockingFactor) or pre-blocked K-tiles (step == 1).
Value quotientInnerLoop =
- arith::DivUIOp::create(rewriter, loc, ivInnerLoop, packOffset);
+ arith::DivUIOp::create(rewriter, loc, ivInnerLoop, blockStride);
Value remInnerLoop = arith::RemUIOp::create(
rewriter, loc, rewriter.getIndexType(), quotientInnerLoop, c2);
@@ -587,8 +589,7 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
nLoadIndx, ivNewInnerLoop);
indxToStoreInBuffer = getIndxToLoadStoreFromPckBuffer(
rewriter, loc, ivNewInnerLoop, ivOuterLoop,
- isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack,
- blockingFactor);
+ isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack, step);
Value indxToLoadFromMatB =
arith::AddIOp::create(rewriter, loc, indxToStoreInBuffer, c1);
indxToLoadFromBuffer =
@@ -653,14 +654,16 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
if (!isVnni) {
if (outerLoop) {
if (!pack) {
- Value nLoadIndx = arith::ConstantIndexOp::create(
- rewriter, locNewInnerLoop, offset);
matB = Value();
indxToLoadFromBuffer = c0;
+ // Use the real spill-block induction value (== spillInnerLoop)
+ // together with the loop step so the computed ping-pong slot
+ // matches the prefetch store side for any number of register
+ // blocks, including odd counts (e.g. 96 = 3 blocks). Passing a
+ // constant here mis-parities the slot for odd block counts.
indxToLoadFromBuffer = getIndxToLoadStoreFromPckBuffer(
- rewriter, loc, nLoadIndx, ivOuterLoop,
- isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack,
- blockingFactor);
+ rewriter, loc, ivNewInnerLoop, ivOuterLoop,
+ isInnerLoopUBHasOddQuot, isInnerLoopUBLarger, pack, step);
}
} else {
if (!pack) {
@@ -1045,10 +1048,18 @@ struct VectorContractToAMXDotProduct
while (true) {
Operation *parent = current->getParentOfType<scf::ForOp>();
- if (!parent)
+ if (!parent) {
+ // The accumulator initialization can be hoisted above an enclosing
+ // parallel region (scf.parallel/scf.forall) when the register tile
+ // matches the problem size and the M/N register loops fold away. In
+ // that case the reduction loop(s) collected so far are still valid to
+ // rewrite, so stop climbing instead of bailing out.
+ if (!loopLists.empty())
+ break;
return rewriter.notifyMatchFailure(
contractOp,
"Accumulator read and contract op not within scf.for op");
+ }
loopLists.push_back(dyn_cast<scf::ForOp>(parent));
>From 294e884a71905877df4ae45e768fc2e13443b007 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Fri, 31 Jul 2026 02:38:38 -0700
Subject: [PATCH 2/3] fix clang format + add a new test-case
---
.../VectorContractToAMXDotProduct.cpp | 10 ++--
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 56 +++++++++++++++++++
2 files changed, 62 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 28b8542d79cfb..c5ec8c8b660e5 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -492,10 +492,12 @@ static SmallVector<Value> createTileZeros(OpBuilder &rewriter, Location loc,
return loopItrArgs;
}
-static Value getIndxToLoadStoreFromPckBuffer(
- OpBuilder &rewriter, Location loc, Value ivInnerLoop, Value ivOuterLoop,
- bool isInnerLoopUBHasOddQuot, bool isInnerLoopUBLarger, bool pack,
- Value blockStride) {
+static Value getIndxToLoadStoreFromPckBuffer(OpBuilder &rewriter, Location loc,
+ Value ivInnerLoop,
+ Value ivOuterLoop,
+ bool isInnerLoopUBHasOddQuot,
+ bool isInnerLoopUBLarger,
+ bool pack, Value blockStride) {
Value c2 = arith::ConstantIndexOp::create(rewriter, loc, 2);
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 ad983ee5d6ff1..a31fdcd80e16c 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
@@ -483,6 +483,62 @@ func.func @brgemm_bf16_loop(%arg0: memref<16x64x64x2xbf16>, %arg1: memref<16x64x
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
// 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
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @M_N_K_Can(%arg0: memref<16x32x32xbf16>, %arg1: memref<16x32x32xbf16>, %arg2: memref<32x32xf32>) {
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c1 = arith.constant 1 : index
+ %2 = vector.transfer_read %arg2[%c0, %c0], %0 {in_bounds = [true, true]} : memref<32x32xf32>, vector<16x16xf32>
+ %3 = vector.transfer_read %arg2[%c0, %c16], %0 {in_bounds = [true, true]} : memref<32x32xf32>, vector<16x16xf32>
+ %4 = vector.transfer_read %arg2[%c16, %c0], %0 {in_bounds = [true, true]} : memref<32x32xf32>, vector<16x16xf32>
+ %5 = vector.transfer_read %arg2[%c16, %c16], %0 {in_bounds = [true, true]} : memref<32x32xf32>, vector<16x16xf32>
+ %6:4 = scf.for %arg3 = %c0 to %c16 step %c1 iter_args(%arg4 = %2, %arg5 = %3, %arg6 = %4, %arg7 = %5) -> (vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>) {
+ %subview = memref.subview %arg0[%arg3, 0, 0] [1, 32, 32] [1, 1, 1] : memref<16x32x32xbf16> to memref<1x32x32xbf16, strided<[1024, 32, 1], offset: ?>>
+ %subview_0 = memref.subview %arg1[%arg3, 0, 0] [1, 32, 32] [1, 1, 1] : memref<16x32x32xbf16> to memref<1x32x32xbf16, strided<[1024, 32, 1], offset: ?>>
+ %7 = vector.transfer_read %subview[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} : memref<1x32x32xbf16, strided<[1024, 32, 1], offset: ?>>, vector<1x16x32xbf16>
+ %8 = vector.transfer_read %subview[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} : memref<1x32x32xbf16, strided<[1024, 32, 1], offset: ?>>, vector<1x16x32xbf16>
+ %9 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} : memref<1x32x32xbf16, strided<[1024, 32, 1], offset: ?>>, vector<1x32x16xbf16>
+ %10 = vector.transfer_read %subview_0[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]} : memref<1x32x32xbf16, strided<[1024, 32, 1], offset: ?>>, vector<1x32x16xbf16>
+ %11 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %7, %9, %arg4 {unroll_shape = array<i64: 1, 16, 16, 32>} : vector<1x16x32xbf16>, vector<1x32x16xbf16> into vector<16x16xf32>
+ %12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %7, %10, %arg5 {unroll_shape = array<i64: 1, 16, 16, 32>} : vector<1x16x32xbf16>, vector<1x32x16xbf16> into vector<16x16xf32>
+ %13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %8, %9, %arg6 {unroll_shape = array<i64: 1, 16, 16, 32>} : vector<1x16x32xbf16>, vector<1x32x16xbf16> into vector<16x16xf32>
+ %14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %8, %10, %arg7 {unroll_shape = array<i64: 1, 16, 16, 32>} : vector<1x16x32xbf16>, vector<1x32x16xbf16> into vector<16x16xf32>
+ scf.yield %11, %12, %13, %14 : vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
+ }
+ vector.transfer_write %6#3, %arg2[%c16, %c16] {in_bounds = [true, true]} : vector<16x16xf32>, memref<32x32xf32>
+ vector.transfer_write %6#2, %arg2[%c16, %c0] {in_bounds = [true, true]} : vector<16x16xf32>, memref<32x32xf32>
+ vector.transfer_write %6#1, %arg2[%c0, %c16] {in_bounds = [true, true]} : vector<16x16xf32>, memref<32x32xf32>
+ vector.transfer_write %6#0, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf32>, memref<32x32xf32>
+ return
+}
+
+// CHECK-LABEL: @M_N_K_Can
+// CHECK-2: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
+// CHECK-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
+// CHECK-4: x86.amx.tile_load
+// CHECK-4: x86.amx.tile_mulf
+// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
+// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
+// 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
>From 42c469ab7f9add73522df871b3ed0d91f00e3076 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Fri, 31 Jul 2026 07:44:23 -0700
Subject: [PATCH 3/3] rename test-case
---
mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
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 a31fdcd80e16c..2cada1d853772 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
@@ -499,7 +499,7 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
-func.func @M_N_K_Can(%arg0: memref<16x32x32xbf16>, %arg1: memref<16x32x32xbf16>, %arg2: memref<32x32xf32>) {
+func.func @Cano_Opt_Removed_M_N_K_loops_pattern(%arg0: memref<16x32x32xbf16>, %arg1: memref<16x32x32xbf16>, %arg2: memref<32x32xf32>) {
%0 = ub.poison : f32
%1 = ub.poison : bf16
%c0 = arith.constant 0 : index
@@ -529,7 +529,7 @@ func.func @M_N_K_Can(%arg0: memref<16x32x32xbf16>, %arg1: memref<16x32x32xbf16>,
return
}
-// CHECK-LABEL: @M_N_K_Can
+// CHECK-LABEL: @Cano_Opt_Removed_M_N_K_loops_pattern
// CHECK-2: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
// CHECK-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
// CHECK-4: x86.amx.tile_load
More information about the Mlir-commits
mailing list