[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 ↗

(On Diff #413308)

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 ↗

(On Diff #413308)

"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/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