[Mlir-commits] [mlir] [mlir][x86] Bail out gracefully if accumulator is initialized from block arg (PR #228080)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Julian Oppermann (jopperm)

<details>
<summary>Changes</summary>

`traceToVectorReadLikeParentOperation` may return `nullptr`, so defer further checks until we know that a suitable op was found.

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


2 Files Affected:

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


``````````diff
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 32ab0500a3921..e8475c2ed4a3f 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -818,8 +818,6 @@ struct VectorContractToAMXDotProduct
     Operation *accReadOp =
         traceToVectorReadLikeParentOperation(contractOp.getAcc());
 
-    bool isAccZeroVectorConstant = isZeroVectorConstant(accReadOp);
-
     Operation *accWrite =
         traceToVectorWriteLikeUserOperation(contractOp.getResult());
 
@@ -834,6 +832,8 @@ struct VectorContractToAMXDotProduct
                       "transfer_read or a load. And, the result should have a "
                       "single-use chain to its consumer.");
 
+    bool isAccZeroVectorConstant = isZeroVectorConstant(accReadOp);
+
     Block *resultBlock = resultChainEnd.user_begin()->getBlock();
 
     Type ipType = rewriter.getBF16Type();
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 8e2791ed9153c..3e6d53fd4f145 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
@@ -2391,6 +2391,54 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+!vecA = vector<16x16x4xi8>
+!vecB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<32x16x4xi8>
+!memrefB = memref<16x32x4xi8>
+!memrefC = memref<32x32xi32>
+#map = affine_map<(d4, d1, d2, d3) -> (d1, d3, d4)>
+#map1 = affine_map<(d4, d1, d2, d3) -> (d3, d2, d4)>
+#map2 = affine_map<(d4, d1, d2, d3) -> (d1, d2)>
+
+func.func @negative_no_acc_op(
+  %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC, %arg3: !vecC) -> !memrefC
+{
+  %c0 = arith.constant 0 : index
+  %0 = ub.poison : i8
+  %32 = ub.poison : i32
+
+  %1 = vector.transfer_read %arg0[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+        !memrefA, !vecA
+  %2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+        !memrefB, !vecB
+  
+  %4 = vector.contract {
+    indexing_maps = [#map, #map1, #map2],
+    iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+    kind = #vector.kind<add>}
+    %1, %2, %arg3 : !vecA, !vecB into !vecC
+
+  vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+  return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_no_acc_op
+// CHECK-NOT: x86.amx
+// CHECK: vector.contract
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %0 {
+      transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+    } : !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
 !vecA = vector<16x16x4xi8>
 !vecB = vector<16x16x4xi8>
 !vecC = vector<16x16xi32>

``````````

</details>


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