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

[llvmlistbot at llvm.org]

Author: Julian Oppermann
Date: 2026-10-05T08:02:51Z
New Revision: f96d2e83ca4667d44e71bfb242fe9639f94f437e

URL: https://github.com/llvm/llvm-project/commit/f96d2e83ca4667d44e71bfb242fe9639f94f437e
DIFF: https://github.com/llvm/llvm-project/commit/f96d2e83ca4667d44e71bfb242fe9639f94f437e.diff

LOG: [mlir][x86] Bail out gracefully if accumulator is initialized from block arg (#228080)

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

Added: 
    

Modified: 
    mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
    mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 32ab0500a3921..ec3c14fd3e708 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;
 
@@ -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>