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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-amx

Author: Adam Siemieniuk (adam-smnk)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/232279.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp (+11-3) 
- (modified) mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir (+63) 


``````````diff
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: ?>>

``````````

</details>


https://github.com/llvm/llvm-project/pull/232279