[Mlir-commits] [mlir] [mlir][x86] Fix - Allow AMX lowering not to care about epilogue ops. (PR #215459)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Arun Thangamani (arun-thmn)

<details>
<summary>Changes</summary>

`x86` AMX lowering looks to patten match with `transfer_write` or `store` and hence it doesn't re-write for epilogue ops.
This patch fixes/extend the lower pattern of AMX to not care about epilogue ops.

Issue: https://github.com/libxsmm/tpp-mlir/issues/1171

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


2 Files Affected:

- (modified) mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp (+12-8) 
- (modified) mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir (+60) 


``````````diff
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index c5ec8c8b660e5..365a3fb8e523b 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -802,14 +802,18 @@ struct VectorContractToAMXDotProduct
     Operation *accReadOp =
         traceToVectorReadLikeParentOperation(contractOp.getAcc());
 
-    Operation *resultWriteOp =
-        traceToVectorWriteLikeUserOperation(contractOp.getResult());
+    // Only the contract result's first consumer is needed, not the final
+    // store. This keeps the lowering independent of the epilogue ops (truncf,
+    // bias add, ReLU, ...) that sit between the contraction and the write.
+    Value resultChainEnd = contractionUsersAfterYield(contractOp.getResult());
 
-    if (!accReadOp || !resultWriteOp)
+    if (!accReadOp || !resultChainEnd)
       return rewriter.notifyMatchFailure(
           contractOp, "The ACC operand of the vector.contract should be a "
-                      "transfer_read or a load. And, the result should be "
-                      "stored using transfer_write or store.");
+                      "transfer_read or a load. And, the result should have a "
+                      "single-use chain to its consumer.");
+
+    Block *resultBlock = resultChainEnd.user_begin()->getBlock();
 
     Type ipType = rewriter.getBF16Type();
     Type opType = rewriter.getF32Type();
@@ -826,12 +830,12 @@ struct VectorContractToAMXDotProduct
       ipType = rewriter.getF8E5M2Type();
 
     if (accReadOp->getBlock() == contractOp->getBlock() &&
-        resultWriteOp->getBlock() != contractOp->getBlock())
+        resultBlock != contractOp->getBlock())
       return rewriter.notifyMatchFailure(
           contractOp, "The accumulator store is in different block.");
 
     if (accReadOp->getBlock() != contractOp->getBlock() &&
-        resultWriteOp->getBlock() == contractOp->getBlock())
+        resultBlock == contractOp->getBlock())
       return rewriter.notifyMatchFailure(
           contractOp, "The accumulator read is in different block.");
 
@@ -847,7 +851,7 @@ struct VectorContractToAMXDotProduct
     // Case 1: For just one VC rewrite. Where all accumulator read/write
     // within the same block.
     if (accReadOp->getBlock() == contractOp->getBlock() &&
-        resultWriteOp->getBlock() == contractOp->getBlock()) {
+        resultBlock == contractOp->getBlock()) {
 
       if (!isReadSrcMemref(contractOp.getAcc()))
         return rewriter.notifyMatchFailure(contractOp,
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 2cada1d853772..093e97af0520e 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
@@ -162,6 +162,66 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+!vecA = vector<1x16x16x2xbf16>
+!vecB = vector<1x16x16x2xbf16>
+!vecC = vector<16x16xf32>
+!memrefA = memref<1x32x16x2xbf16>
+!memrefB = memref<1x16x32x2xbf16>
+!memrefC = memref<32x32xf32>
+!memrefO = memref<32x32xbf16>
+#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
+#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
+#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
+func.func @brgemm_bf16_epilogue(
+  %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC, %arg3: !memrefO) -> !memrefO
+{
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant dense<0.000000e+00> : vector<16x16xbf16>
+  %0 = ub.poison : bf16
+  %32 = ub.poison : f32
+
+  %1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
+        !memrefA, !vecA
+  %2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, 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 = ["reduction", "reduction", "parallel", "parallel", "reduction"],
+    kind = #vector.kind<add>}
+    %1, %2, %3 : !vecA, !vecB into !vecC
+
+  %5 = arith.truncf %4 : vector<16x16xf32> to vector<16x16xbf16>
+  %6 = arith.cmpf ugt, %5, %cst : vector<16x16xbf16>
+  %7 = arith.select %6, %5, %cst : vector<16x16xi1>, vector<16x16xbf16>
+  vector.transfer_write %7, %arg3[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xbf16>, !memrefO
+
+  return %arg3 : !memrefO
+}
+
+// CHECK-LABEL: @brgemm_bf16_epilogue
+// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
+// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
+// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xf32>
+// CHECK: x86.amx.tile_mulf
+// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xf32>
+// CHECK: arith.select
+// 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
+  }
+}
+
+// -----
+
 !vecA = vector<1x16x16x2xbf16>
 !vecB = vector<1x16x16x2xbf16>
 !vecC = vector<1x16x16xf32>

``````````

</details>


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