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

[Julian Oppermann]

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

>From 72d94d86cf25c965ef8f94c0899d211dffc18a60 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Thu, 1 Oct 2026 06:33:43 -0700
Subject: [PATCH 1/2] [mlir][x86] Bail out gracefully if accumulator is
 initialized from block arg

---
 .../VectorContractToAMXDotProduct.cpp         |  4 +-
 .../X86/AMX/vector-contract-to-tiled-dp.mlir  | 48 +++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 32ab0500a39210e..e8475c2ed4a3f3c 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 8e2791ed9153c7c..3e6d53fd4f14572 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>

>From 6b3f888453406cb8665a1400aedb52f1d8c66afe Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Fri, 2 Oct 2026 00:02:53 -0700
Subject: [PATCH 2/2] Use dyn_cast_if_present

---
 .../Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index e8475c2ed4a3f3c..ec3c14fd3e70816 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -30,7 +30,7 @@ namespace {
 // Return true if the operation is a constant dense element
 // attribute.
 static bool isZeroVectorConstant(Operation *op) {
-  auto constantOp = dyn_cast<arith::ConstantOp>(op);
+  auto constantOp = dyn_cast_if_present<arith::ConstantOp>(op);
   if (!constantOp)
     return false;