[Mlir-commits] [mlir] [mlir] added a check in the walk to prevent catching a cos in a nested region (PR #190064)

[llvmlistbot at llvm.org]

https://github.com/yebinchon created https://github.com/llvm/llvm-project/pull/190064

The walk in SincosFusion may detect a cos within a nested region of the sin block. This triggers an assertion in `isBeforeInBlock` later on. Added a check within the walk so it filters operations in nested regions, which are not in the same block and should not be fused anyway.

>From 30c8adc4679003b688596213233d3549ff8b3692 Mon Sep 17 00:00:00 2001
From: Yebin Chon <ychon at nvidia.com>
Date: Wed, 1 Apr 2026 13:53:56 -0700
Subject: [PATCH] added a check in the walk to prevent catching an op in a
 nested region

---
 .../Dialect/Math/Transforms/SincosFusion.cpp  |  4 +++-
 mlir/test/Dialect/Math/sincos-fusion.mlir     | 23 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp b/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp
index 69407df201cfa..c6602eb293cbb 100644
--- a/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp
@@ -27,8 +27,10 @@ struct SincosFusionPattern : OpRewritePattern<math::SinOp> {
     mlir::arith::FastMathFlags sinFastMathFlags = sinOp.getFastmath();
 
     math::CosOp cosOp = nullptr;
+    Block* sinBlock = sinOp->getBlock();
     sinOp->getBlock()->walk([&](math::CosOp op) {
-      if (op.getOperand() == operand && op.getFastmath() == sinFastMathFlags) {
+      if (op->getBlock() == sinBlock && op.getOperand() == operand && 
+          op.getFastmath() == sinFastMathFlags) {
         cosOp = op;
         return WalkResult::interrupt();
       }
diff --git a/mlir/test/Dialect/Math/sincos-fusion.mlir b/mlir/test/Dialect/Math/sincos-fusion.mlir
index 29fb9f12475b8..cf16f9f02f63a 100644
--- a/mlir/test/Dialect/Math/sincos-fusion.mlir
+++ b/mlir/test/Dialect/Math/sincos-fusion.mlir
@@ -74,6 +74,29 @@ func.func @sincos_no_fusion_different_block(%arg0 : f32, %flag : i1) -> f32 {
   func.return %0 : f32
 }
 
+// CHECK-LABEL:   func.func @sincos_no_fusion_nested_region(
+// CHECK-SAME:      %[[ARG0:.*]]: f32,
+// CHECK-SAME:      %[[ARG1:.*]]: i1) -> (f32, f32) {
+// CHECK:           %[[SIN:.*]] = math.sin %[[ARG0]] : f32
+// CHECK:           %[[IF:.*]] = scf.if %[[ARG1]] -> (f32) {
+// CHECK:             %[[COS:.*]] = math.cos %[[ARG0]] : f32
+// CHECK:             scf.yield %[[COS]] : f32
+// CHECK:           } else {
+// CHECK:             scf.yield %[[SIN]] : f32
+// CHECK:           }
+// CHECK:           return %[[SIN]], %[[IF]] : f32, f32
+// CHECK:         }
+func.func @sincos_no_fusion_nested_region(%arg0 : f32, %flag : i1) -> (f32, f32) {
+    %s = math.sin %arg0 : f32
+    %r = scf.if %flag -> f32 {
+      %c = math.cos %arg0 : f32
+      scf.yield %c : f32
+    } else {
+      scf.yield %s : f32
+    }
+    func.return %s, %r : f32, f32
+}
+
 // CHECK-LABEL:   func.func @sincos_fusion_preserve_fastmath(
 // CHECK-SAME:      %[[ARG0:.*]]: f32) -> (f32, f32) {
 // CHECK:           %[[VAL_0:.*]], %[[VAL_1:.*]] = math.sincos %[[ARG0]] fastmath<contract> : f32