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

[Hocky Yudhiono]

================
@@ -296,19 +296,20 @@ 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.
-// TODO: relax to multi-operands with constants, which are technically unary ops
-// as needed (e.g. add5).
-static Value getSourceSkipUnary(Value value) {
+/// Return true for scalar arith cast ops modeled by current linalg contraction
+/// cast semantics.
+static bool isSupportedContractionCast(Operation *op) {
+  return isa<arith::ExtFOp, arith::TruncFOp, arith::ExtSIOp, arith::ExtUIOp,
+             arith::TruncIOp, arith::SIToFPOp, arith::UIToFPOp, arith::FPToSIOp,
+             arith::FPToUIOp>(op);
+}
+
+/// If the value is defined by a supported contraction cast op, return its
+/// source. Otherwise, return the value unchanged.
+static Value getSourceSkipCast(Value value) {
   Operation *op = value.getDefiningOp();
-  while (op && op->getNumOperands() == 1) {
-    auto iface = dyn_cast<MemoryEffectOpInterface>(op);
-    if (!iface || !iface.hasNoEffect())
-      break;
-    value = op->getOperand(0);
-    op = value.getDefiningOp();
-  }
+  if (op && op->getNumOperands() == 1 && isSupportedContractionCast(op))
----------------
hockyy wrote:

nit: can just use `isa_and_present<arith::ExtFOp, arith::TruncFOp, arith::ExtSIOp, arith::ExtUIOp,
             arith::TruncIOp, arith::SIToFPOp, arith::UIToFPOp, arith::FPToSIOp,
             arith::FPToUIOp>(op));`

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