[Mlir-commits] [mlir] [mlir][x86] Support rank-reducing RHS reads in AMX online packing (PR #232279)

[Adam Siemieniuk]

https://github.com/adam-smnk created https://github.com/llvm/llvm-project/pull/232279

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

>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] [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: ?>>