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

[llvmlistbot at llvm.org]

Author: Matt Hofmann
Date: 2024-09-05T00:12:03-04:00
New Revision: ad89e617c703239518187912540b8ea811dc2eda

URL: https://github.com/llvm/llvm-project/commit/ad89e617c703239518187912540b8ea811dc2eda
DIFF: https://github.com/llvm/llvm-project/commit/ad89e617c703239518187912540b8ea811dc2eda.diff

LOG: [MLIR][Python] Fix detached operation coming from `IfOp` constructor (#107286)

Without this fix, `scf.if` operations would be created without a parent.
Since `scf.if` operations often have no results, this caused silent bugs
where the generated code was straight-up missing the operation.

Added: 
    

Modified: 
    mlir/python/mlir/dialects/scf.py
    mlir/test/python/dialects/scf.py

Removed: 
    


################################################################################
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(*[])

diff  --git a/mlir/test/python/dialects/scf.py b/mlir/test/python/dialects/scf.py
index 95a6de86b670d5..de61f4613868f7 100644
--- a/mlir/test/python/dialects/scf.py
+++ b/mlir/test/python/dialects/scf.py
@@ -278,6 +278,32 @@ def simple_if(cond):
 # 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():
     bool = IntegerType.get_signless(1)