[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/5] 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/5] 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():

>From c90b246264b4f00fbaf34c787970ed74aca1c01a Mon Sep 17 00:00:00 2001
From: Matt Hofmann <mrh259 at cornell.edu>
Date: Wed, 4 Sep 2024 16:15:41 -0400
Subject: [PATCH 3/5] Use 4 spaces

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

diff --git a/mlir/test/python/dialects/scf.py b/mlir/test/python/dialects/scf.py
index 4ece6eef4e3f1c..04a548e84c07aa 100644
--- a/mlir/test/python/dialects/scf.py
+++ b/mlir/test/python/dialects/scf.py
@@ -279,20 +279,20 @@ def simple_if(cond):
 
 @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
+    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)

>From a52f7f0404e220a130c9c6393a6ff4cd38e3246c Mon Sep 17 00:00:00 2001
From: Matt Hofmann <mrh259 at cornell.edu>
Date: Wed, 4 Sep 2024 16:22:06 -0400
Subject: [PATCH 4/5] Format scf.py

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

diff --git a/mlir/test/python/dialects/scf.py b/mlir/test/python/dialects/scf.py
index 04a548e84c07aa..63cb33b48179d6 100644
--- a/mlir/test/python/dialects/scf.py
+++ b/mlir/test/python/dialects/scf.py
@@ -277,6 +277,7 @@ def simple_if(cond):
 # CHECK:   %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
 # CHECK: return
 
+
 @constructAndPrintInModule
 def testNestedIf():
     bool = IntegerType.get_signless(1)

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

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

diff --git a/mlir/test/python/dialects/scf.py b/mlir/test/python/dialects/scf.py
index 63cb33b48179d6..de61f4613868f7 100644
--- a/mlir/test/python/dialects/scf.py
+++ b/mlir/test/python/dialects/scf.py
@@ -303,6 +303,7 @@ def nested_if(b, c):
 # CHECK:     %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
 # CHECK: return
 
+
 @constructAndPrintInModule
 def testIfWithElse():
     bool = IntegerType.get_signless(1)