[Mlir-commits] [mlir] [mlir][linalg] Fix getSourceSkipUnary to only skip cast-like ops (PR #198725)

[llvmlistbot at llvm.org]

https://github.com/Chennesxu updated https://github.com/llvm/llvm-project/pull/198725

>From 31330df6b313adce3a5ba6ca14a109a318c31ccd Mon Sep 17 00:00:00 2001
From: Chennes Xu <xuchen359 at gmail.com>
Date: Wed, 20 May 2026 16:01:26 +0800
Subject: [PATCH] [mlir][linalg] Fix getSourceSkipUnary to only skip cast-like
 ops

getSourceSkipUnary was skipping all unary side-effect-free ops when
matching contraction bodies, but its callers expect it to only skip
cast-like ops (as stated in the error messages: "modulo unary casts").
This caused non-cast unary ops like arith.negf to be silently ignored,
leading to incorrect specialization of linalg.generic to linalg.matmul.

Restrict the skip condition to ops implementing CastOpInterface, which
correctly identifies type conversion ops (extf, trunci, sitofp, etc.)
while rejecting semantics-changing ops (negf, absf, etc.).

Fixes #197178.
---
 mlir/lib/Dialect/Linalg/IR/CMakeLists.txt     |  1 +
 .../Dialect/Linalg/IR/LinalgInterfaces.cpp    | 11 +++++-----
 .../Linalg/specialize-generic-ops-fail.mlir   | 21 +++++++++++++++++++
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
index ec433284e17ad..45744550f5173 100644
--- a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
@@ -21,6 +21,7 @@ add_mlir_dialect_library(MLIRLinalgDialect
   MLIRArithDialect
   MLIRArithUtils
   MLIRBufferizationDialect
+  MLIRCastInterfaces
   MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRFunctionInterfaces
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index 2ba77cea8f16e..0f864590d5e74 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -19,6 +19,7 @@
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/CastInterfaces.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallBitVector.h"
@@ -295,16 +296,14 @@ bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {
 // ContractionOpInterface implementation
 //===----------------------------------------------------------------------===//
 
-/// If the value is defined by a chain of unary side effect-free, go up the
-/// use-def chain until the first value that isn't defined by such an op.
+/// If the value is defined by a chain of unary cast-like ops (as identified
+/// by CastOpInterface), go up the use-def chain until the first value that
+/// isn't defined by such an op.
 // TODO: relax to multi-operands with constants, which are technically unary ops
 // as needed (e.g. add5).
 static Value getSourceSkipUnary(Value value) {
   Operation *op = value.getDefiningOp();
-  while (op && op->getNumOperands() == 1) {
-    auto iface = dyn_cast<MemoryEffectOpInterface>(op);
-    if (!iface || !iface.hasNoEffect())
-      break;
+  while (op && op->getNumOperands() == 1 && isa<CastOpInterface>(op)) {
     value = op->getOperand(0);
     op = value.getDefiningOp();
   }
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
index 5d66837fca510..a4ef6efddcae8 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
@@ -46,3 +46,24 @@ func.func @not_copy(%input: tensor<8xi32>, %init: tensor<8xi32>) -> tensor<8xi32
   } -> tensor<8xi32>
   return %res : tensor<8xi32>
 }
+
+// -----
+
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+// This test checks that linalg.generic with a negf between mulf and addf
+// does not get incorrectly specialized to matmul.
+// CHECK-LABEL: @contraction_with_negf
+//  CHECK-NOT:    linalg.matmul
+//      CHECK:    linalg.generic
+func.func @contraction_with_negf(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
+  %0 = linalg.generic {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %1 = arith.mulf %in, %in_0 : f32
+    %2 = arith.negf %1 : f32
+    %3 = arith.addf %out, %2 : f32
+    linalg.yield %3 : f32
+  } -> tensor<3x3xf32>
+  return %0 : tensor<3x3xf32>
+}