[Mlir-commits] [mlir] [MLIR][Python] Fix detached operation coming from `IfOp` constructor (PR #107286)

[Matt Hofmann]

https://github.com/matth2k updated https://github.com/llvm/llvm-project/pull/107286

>From 4b45423e7c23ff4fd72945aa00031597b4c0cbe6 Mon Sep 17 00:00:00 2001
From: Matt Hofmann <47065711+matth2k at users.noreply.github.com>
Date: Wed, 4 Sep 2024 14:29:21 -0400
Subject: [PATCH 1/2] Fix detached operation from IfOp constructor

---
 mlir/python/mlir/dialects/scf.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/python/mlir/dialects/scf.py b/mlir/python/mlir/dialects/scf.py
index 7025f6e0f1a166..2d0047b76c7022 100644
--- a/mlir/python/mlir/dialects/scf.py
+++ b/mlir/python/mlir/dialects/scf.py
@@ -87,7 +87,7 @@ def __init__(self, cond, results_=None, *, hasElse=False, loc=None, ip=None):
         operands.append(cond)
         results = []
         results.extend(results_)
-        super().__init__(results, cond)
+        super().__init__(results, cond, loc=loc, ip=ip)
         self.regions[0].blocks.append(*[])
         if hasElse:
             self.regions[1].blocks.append(*[])

>From e6108a14c1651098d5406d1cc2cb1e88cc21ff53 Mon Sep 17 00:00:00 2001
From: Matt Hofmann <mrh259 at cornell.edu>
Date: Wed, 4 Sep 2024 16:03:39 -0400
Subject: [PATCH 2/2] Update scf.py

---
 mlir/test/python/dialects/scf.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/mlir/test/python/dialects/scf.py b/mlir/test/python/dialects/scf.py
index 95a6de86b670d5..4ece6eef4e3f1c 100644
--- a/mlir/test/python/dialects/scf.py
+++ b/mlir/test/python/dialects/scf.py
@@ -277,6 +277,30 @@ def simple_if(cond):
 # CHECK:   %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
 # CHECK: return
 
+ at constructAndPrintInModule
+def testNestedIf():
+  bool = IntegerType.get_signless(1)
+  i32 = IntegerType.get_signless(32)
+
+  @func.FuncOp.from_py_func(bool, bool)
+  def nested_if(b, c):
+    if_op = scf.IfOp(b)
+    with InsertionPoint(if_op.then_block) as ip:
+      if_op = scf.IfOp(c, ip=ip)
+      with InsertionPoint(if_op.then_block):
+        one = arith.ConstantOp(i32, 1)
+        add = arith.AddIOp(one, one)
+        scf.YieldOp([])
+      scf.YieldOp([])
+    return
+
+
+# CHECK: func @nested_if(%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
+# CHECK: scf.if %[[ARG0:.*]]
+# CHECK:   scf.if %[[ARG1:.*]]
+# CHECK:     %[[ONE:.*]] = arith.constant 1
+# CHECK:     %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
+# CHECK: return
 
 @constructAndPrintInModule
 def testIfWithElse():