[mlir][python] Wrappers for scf.index_switch

The C++ index switch op has utilies for getCaseBlock(int i)
and getDefaultBlock(), so these have been added.

Optional builder args have been added for the default case
and each switch case.

The list comprehensions for accessing case regions are due
to what appears to be a bug in RegionSequence; using a comprehension
with explicit indices circumvents this.
The same paradigm is used for get_case_block(i: int), but this
is unavoidable.

Author: ashermancinelli

mlir/python/mlir/dialects/scf.py

@@ -6,6 +6,7 @@
 from ._scf_ops_gen import *
 from ._scf_ops_gen import _Dialect
 from .arith import constant
+import builtins
 
 try:
     from ..ir import *
@@ -19,7 +20,6 @@
 
 from typing import List, Optional, Sequence, Tuple, Union
 
-
 @_ods_cext.register_operation(_Dialect, replace=True)
 class ForOp(ForOp):
     """Specialization for the SCF for op class."""
@@ -254,3 +254,76 @@ def for_(
             yield iv, iter_args[0], for_op.results[0]
         else:
             yield iv
+
+@_ods_cext.register_operation(_Dialect, replace=True)
+class IndexSwitchOp(IndexSwitchOp):
+    __doc__ = IndexSwitchOp.__doc__
+
+    def __init__(
+        self,
+        results_,
+        arg,
+        cases,
+        case_body_builder=None,
+        default_body_builder=None,
+        loc=None,
+        ip=None,
+    ):
+        cases = DenseI64ArrayAttr.get(cases)
+        super().__init__(
+            results_, arg, cases, num_caseRegions=len(cases), loc=loc, ip=ip
+        )
+        for region in self.regions:
+            region.blocks.append()
+
+        if default_body_builder is not None:
+            with InsertionPoint(self.default_block):
+                default_body_builder(self)
+
+        if case_body_builder is not None:
+            for i, case in enumerate(cases):
+                with InsertionPoint(self.case_block(i)):
+                    case_body_builder(self, i, self.cases[i])
+
+    @builtins.property
+    def default_region(self) -> Region:
+        return self.regions[0]
+
+    @builtins.property
+    def default_block(self) -> Block:
+        return self.default_region.blocks[0]
+
+    @builtins.property
+    def case_regions(self) -> Sequence[Region]:
+        return [self.regions[1 + i] for i in range(len(self.cases))]
+
+    def case_region(self, i: int) -> Region:
+        return self.case_regions[i]
+
+    @builtins.property
+    def case_blocks(self) -> Sequence[Block]:
+        return [region.blocks[0] for region in self.case_regions]
+
+    def case_block(self, i: int) -> Block:
+        return self.case_regions[i].blocks[0]
+
+def index_switch(
+    results_,
+    arg,
+    cases,
+    case_body_builder=None,
+    default_body_builder=None,
+    loc=None,
+    ip=None,
+) -> Union[OpResult, OpResultList, IndexSwitchOp]:
+    op = IndexSwitchOp(
+        results_=results_,
+        arg=arg,
+        cases=cases,
+        case_body_builder=case_body_builder,
+        default_body_builder=default_body_builder,
+        loc=loc,
+        ip=ip,
+    )
+    results = op.results
+    return results if len(results) > 1 else (results[0] if len(results) == 1 else op)

mlir/test/python/dialects/scf.py

@@ -1,10 +1,14 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
-from mlir.dialects import arith
-from mlir.dialects import func
-from mlir.dialects import memref
-from mlir.dialects import scf
+from mlir.extras import types as T
+from mlir.dialects import (
+    arith,
+    func,
+    memref,
+    scf,
+    cf,
+)
 from mlir.passmanager import PassManager
 
 
@@ -355,3 +359,63 @@ def simple_if_else(cond):
 # CHECK:   scf.yield %[[TWO]], %[[THREE]]
 # CHECK: arith.addi %[[RET]]#0, %[[RET]]#1
 # CHECK: return
+
+
+@constructAndPrintInModule
+def testIndexSwitch():
+
+    i32 = T.i32()
+    @func.FuncOp.from_py_func(T.index(), results=[i32])
+    def index_switch(index):
+        c1 = arith.constant(i32, 1)
+        c0 = arith.constant(i32, 0)
+        value = arith.constant(i32, 5)
+        switch_op = scf.IndexSwitchOp([i32], index, range(3))
+
+        assert switch_op.regions[0] == switch_op.default_region
+        assert switch_op.regions[1] == switch_op.case_regions[0]
+        assert switch_op.regions[1] == switch_op.case_region(0)
+        assert len(switch_op.case_regions) == 3
+        assert len(switch_op.regions) == 4
+
+        with InsertionPoint(switch_op.default_block):
+            cf.assert_(arith.constant(T.bool(), 0), "Whoops!")
+            scf.yield_([c1])
+
+        for i, block in enumerate(switch_op.case_blocks):
+            with InsertionPoint(block):
+                scf.YieldOp([arith.constant(i32, i)])
+
+        func.return_([switch_op.results[0]])
+
+    return index_switch
+
+
+@constructAndPrintInModule
+def testIndexSwitchWithBodyBuilders():
+
+    i32 = T.i32()
+    @func.FuncOp.from_py_func(T.index(), results=[i32])
+    def index_switch(index):
+        c1 = arith.constant(i32, 1)
+        c0 = arith.constant(i32, 0)
+        value = arith.constant(i32, 5)
+
+        def default_body_builder(switch_op):
+            cf.assert_(arith.constant(T.bool(), 0), "Whoops!")
+            scf.yield_([c1])
+
+        def case_body_builder(switch_op, case_index: int, case_value: int):
+            scf.YieldOp([arith.constant(i32, case_value)])
+
+        result = scf.index_switch(
+            results_=[i32],
+            arg=index,
+            cases=range(3),
+            case_body_builder=case_body_builder,
+            default_body_builder=default_body_builder,
+        )
+
+        func.return_([result])
+
+    return index_switch