[MLIR][Python] Add SCFIfOp Python binding

[chhzh123]

Summary

Current generated Python binding for the SCF dialect does not allow
users to call IfOp to create if-else branches on their own.
This PR sets up the default binding generation for scf.if operation
to address this problem.

[ftynse]

Thanks! This looks almost good. Could you please add a test?

mlir/python/mlir/dialects/_scf_ops_ext.py

72–78

I would rather use the following signature: self, cond, results_ = [], *, hasElse=False, loc=None, ip=None). This will (1) make it easy to create if operations that don't return values (most cases I've seen so far) and (2) require passing hasElse using kwargs syntax as opposed to just having magic True in the call.

81

"a boolean value" can be misinterpreted as Python True/False.

[chhzh123]

Fix signature and add tests

[chhzh123]

Oh seems I made a wrong merge from the master branch.

[chhzh123]

Revert

[chhzh123]

Fix function signature and add tests

[chhzh123]

Fix

[chhzh123]

@ftynse Sorry for messing up this PR a bit. I've made the changes you suggested. Please have a check. Thanks!

[chhzh123]

Hi @ftynse , could you please take a look at it and see if there is anything that should be changed? Thanks!

[ftynse]

Thanks! Let me know if you need help landing this.

[chhzh123]

@ftynse Thanks for reviewing! Could you help land the PR? Seems I do not have access rights to push to the main repository.

[mehdi_amini]

Can you rebase? Your patch does not apply apparently

[chhzh123]

Rebase

[chhzh123]

Rebase 3 commits

[chhzh123]

Try rebasing again

[chhzh123]

I've rebased to the latest commit in the main branch, but the pre-merge checks still failed.

[mehdi_amini]

Seems like there is a local base commit in your repository: you're not uploading a diff that applies on HEAD.

[chhzh123]

Rebase

[chhzh123]

Passed pre-merge checks. Should be able to be merged. Thanks!

mlir/python/mlir/dialects/_scf_ops_ext.py

class ForOp:
 
   @property
   def inner_iter_args(self):
     """Returns the loop-carried arguments usable within the loop.
 
     To obtain the loop-carried operands, use `iter_args`.
     """
     return self.body.arguments[1:]
+
+
+class IfOp:
+  """Specialization for the SCF if op class."""
+
+  def __init__(self,
+               cond,
+               results_=[],
+               *,
+               hasElse=False,
+               loc=None,
+               ip=None):
+    """Creates an SCF `if` operation.
+
+    - `cond` is a MLIR value of 'i1' type to determine which regions of code will be executed.
+    - `hasElse` determines whether the if operation has the else branch.
+    """
+    operands = []
+    operands.append(cond)
+    results = []
+    results.extend(results_)
+    super().__init__(
+        self.build_generic(
+            regions=2,
+            results=results,
+            operands=operands,
+            loc=loc,
+            ip=ip))
+    self.regions[0].blocks.append(*[])
+    if hasElse:
+        self.regions[1].blocks.append(*[])
+
+  @property
+  def then_block(self):
+    """Returns the then block of the if operation."""
+    return self.regions[0].blocks[0]
+
+  @property
+  def else_block(self):
+    """Returns the else block of the if operation."""
+    return self.regions[1].blocks[0]

mlir/test/python/dialects/scf.py

 # CHECK:   %[[LB:.*]] = arith.constant 0
 # CHECK:   %[[UB:.*]] = arith.constant 42
 # CHECK:   %[[STEP:.*]] = arith.constant 2
 # CHECK:   %[[ARGS:.*]]:2 = call @callee()
 # CHECK:   scf.for %arg0 = %c0 to %c42 step %c2
 # CHECK:   iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
 # CHECK:     scf.yield %{{.*}}, %{{.*}}
 # CHECK:   return
+
+
+@constructAndPrintInModule
+def testIfWithoutElse():
+  bool = IntegerType.get_signless(1)
+  i32 = IntegerType.get_signless(32)
+
+  @builtin.FuncOp.from_py_func(bool)
+  def simple_if(cond):
+    if_op = scf.IfOp(cond)
+    with InsertionPoint(if_op.then_block):
+      one = arith.ConstantOp(i32, 1)
+      add = arith.AddIOp(one, one)
+      scf.YieldOp([])
+    return
+
+
+# CHECK: func @simple_if(%[[ARG0:.*]]: i1)
+# CHECK: scf.if %[[ARG0:.*]]
+# CHECK:   %[[ONE:.*]] = arith.constant 1
+# CHECK:   %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
+# CHECK: return
+
+
+@constructAndPrintInModule
+def testIfWithElse():
+  bool = IntegerType.get_signless(1)
+  i32 = IntegerType.get_signless(32)
+
+  @builtin.FuncOp.from_py_func(bool)
+  def simple_if_else(cond):
+    if_op = scf.IfOp(cond, [i32, i32], hasElse=True)
+    with InsertionPoint(if_op.then_block):
+      x_true = arith.ConstantOp(i32, 0)
+      y_true = arith.ConstantOp(i32, 1)
+      scf.YieldOp([x_true, y_true])
+    with InsertionPoint(if_op.else_block):
+      x_false = arith.ConstantOp(i32, 2)
+      y_false = arith.ConstantOp(i32, 3)
+      scf.YieldOp([x_false, y_false])
+    add = arith.AddIOp(if_op.results[0], if_op.results[1])
+    return
+
+
+# CHECK: func @simple_if_else(%[[ARG0:.*]]: i1)
+# CHECK: %[[RET:.*]]:2 = scf.if %[[ARG0:.*]]
+# CHECK:   %[[ZERO:.*]] = arith.constant 0
+# CHECK:   %[[ONE:.*]] = arith.constant 1
+# CHECK:   scf.yield %[[ZERO]], %[[ONE]]
+# CHECK: } else {
+# CHECK:   %[[TWO:.*]] = arith.constant 2
+# CHECK:   %[[THREE:.*]] = arith.constant 3
+# CHECK:   scf.yield %[[TWO]], %[[THREE]]
+# CHECK: arith.addi %[[RET]]#0, %[[RET]]#1
+# CHECK: return