[Mlir-commits] [mlir] [mlir][x86] Support rank-reducing RHS reads in AMX online packing (PR #232279)
Adam Siemieniuk
llvmlistbot at llvm.org
Fri Oct 16 05:23:18 PDT 2026
https://github.com/adam-smnk updated https://github.com/llvm/llvm-project/pull/232279
>From b1bb9fdcebfa5628a2da3c580ef915f9866617f1 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Wed, 7 Oct 2026 11:30:36 -0700
Subject: [PATCH 1/2] [mlir][x86] Support rank-reducing RHS reads in AMX online
packing
Allows AMX online packing to handle rank-reducing RHS reads.
For a rank-reducing read, the sizes miss the dropped leading unit dims,
which hit an assertion when creating the subview. The packing loads
also indexed the subview with the read's leading offsets twice.
Assisted-by: Claude
---
.../VectorContractToAMXDotProduct.cpp | 14 ++++-
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 63 +++++++++++++++++++
2 files changed, 74 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 23ea062c649a5..f16c435d2e080 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -981,15 +981,23 @@ struct VectorContractToAMXDotProduct
});
auto one = rewriter.getIndexAttr(1);
SmallVector<OpFoldResult> strides(indexVals.size(), one);
+ // A rank-reducing read drops leading unit dims. Append them back for
+ // the subview.
+ SmallVector<int64_t> subviewShape(indexVals.size() - vecTy.getRank(),
+ 1);
+ llvm::append_range(subviewShape, vecTy.getShape());
SmallVector<OpFoldResult> sizes = getAsIndexOpFoldResult(
- contractOp.getRhs().getDefiningOp()->getContext(),
- vecTy.getShape());
+ contractOp.getRhs().getDefiningOp()->getContext(), subviewShape);
auto subview = memref::SubViewOp::create(rewriter, loc, srcBuffRhs,
indexVals, sizes, strides);
+ // The source subview already accounts for the leading unit dimensions.
+ // Set them to zero for the later loads.
+ for (Value &index : MutableArrayRef<Value>(indicesRhs).drop_back(2))
+ index = c0;
auto bufferType = MemRefType::get({16, (16 * blockingFactor)}, ipType);
auto packedBuffer = memref::AllocaOp::create(rewriter, loc, bufferType);
- // create a loop that does online packing.
+ // Create a loop that does online packing.
Value step =
arith::ConstantIndexOp::create(rewriter, loc, blockingFactor);
Value uBound = arith::ConstantIndexOp::create(rewriter, loc,
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 f1d2b969ae303..b71ad093a9ce1 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
@@ -625,6 +625,69 @@ module attributes {transform.with_named_sequence} {
// -----
+// The RHS read drops a leading dim of its buffer and starts at non-zero
+// offsets: the packing subview covers the read's offsets, and the packing loads
+// index it from zero.
+
+!vecA = vector<16x32xbf16>
+!vecB = vector<32x16xbf16>
+!vecC = vector<16x16xf32>
+!memrefA = memref<32x32xbf16>
+!memrefB = memref<2x32x32xbf16>
+!memrefC = memref<32x32xf32>
+#map = affine_map<(d1, d2, d3) -> (d1, d3)>
+#map1 = affine_map<(d1, d2, d3) -> (d3, d2)>
+#map2 = affine_map<(d1, d2, d3) -> (d1, d2)>
+func.func @online_packing_bf16_rank_reducing_rhs(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c1, %c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+
+ %3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+
+ %4 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %3 : !vecA, !vecB into !vecC
+
+ vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @online_packing_bf16_rank_reducing_rhs
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[SUBVIEW:.+]] = memref.subview %arg1[%[[C1]], %[[C0]], %[[C16]]] [1, 32, 16] [1, 1, 1]
+// CHECK: scf.for %[[IV:.+]] = %[[C0]]
+// CHECK: vector.load %[[SUBVIEW]][%[[C0]], %[[IV]], %[[C0]]]
+// CHECK: vector.load %[[SUBVIEW]][%[[C0]], %{{.+}}, %[[C0]]]
+// CHECK: x86.amx.tile_mulf
+// 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
+ }
+}
+
+// -----
+
!vecAB = vector<1x16x16x2xbf16>
!vecC = vector<16x16xf32>
!memrefA = memref<1x32x16x2xbf16, strided<[8192, 128, 2, 1], offset: ?>>
>From e5b1f68d11d0b7a593dac031d342f065558af7d8 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Fri, 16 Oct 2026 14:22:16 +0200
Subject: [PATCH 2/2] Simplify indices + ctx
---
.../VectorContractToAMXDotProduct.cpp | 21 +++++++++----------
1 file changed, 10 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index f16c435d2e080..ae7f3763c37f2 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -986,14 +986,10 @@ struct VectorContractToAMXDotProduct
SmallVector<int64_t> subviewShape(indexVals.size() - vecTy.getRank(),
1);
llvm::append_range(subviewShape, vecTy.getShape());
- SmallVector<OpFoldResult> sizes = getAsIndexOpFoldResult(
- contractOp.getRhs().getDefiningOp()->getContext(), subviewShape);
+ SmallVector<OpFoldResult> sizes =
+ getAsIndexOpFoldResult(contractOp->getContext(), subviewShape);
auto subview = memref::SubViewOp::create(rewriter, loc, srcBuffRhs,
indexVals, sizes, strides);
- // The source subview already accounts for the leading unit dimensions.
- // Set them to zero for the later loads.
- for (Value &index : MutableArrayRef<Value>(indicesRhs).drop_back(2))
- index = c0;
auto bufferType = MemRefType::get({16, (16 * blockingFactor)}, ipType);
auto packedBuffer = memref::AllocaOp::create(rewriter, loc, bufferType);
@@ -1014,16 +1010,19 @@ struct VectorContractToAMXDotProduct
Value i1_load =
arith::AddIOp::create(rewriter, loc, nextLoadIndx, iv);
- indicesRhs[indicesRhs.size() - 2] = iv;
- indicesRhs[indicesRhs.size() - 1] = c0;
- ValueRange range1(indicesRhs);
+ // The source subview already accounts for the leading unit
+ // dimensions.
+ size_t numOffsets = indicesRhs.size();
+ SmallVector<Value> indices(numOffsets, c0);
+ indices[numOffsets - 2] = iv;
+ ValueRange range1(indices);
auto vec1 = vector::LoadOp::create(
rewriter, loc,
VectorType::get(16 * (blockingFactor / 2), ipType), subview,
range1);
- indicesRhs[indicesRhs.size() - 2] = i1_load;
- ValueRange range2(indicesRhs);
+ indices[numOffsets - 2] = i1_load;
+ ValueRange range2(indices);
auto vec2 = vector::LoadOp::create(
rewriter, loc,
VectorType::get(16 * (blockingFactor / 2), ipType), subview,
More information about the Mlir-commits
mailing list