compiler: Add support to generate switch-case

Author: FabioLuporini

devito/ir/iet/algorithms.py

@@ -1,8 +1,11 @@
 from collections import OrderedDict
 
-from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot,
-                           Section, HaloSpot, ExpressionBundle)
-from devito.tools import timed_pass
+from devito.ir.iet import (
+    Expression, Increment, Iteration, List, Conditional, SyncSpot, Section,
+    HaloSpot, ExpressionBundle, Switch
+)
+from devito.ir.support import GuardSwitch, GuardCaseSwitch
+from devito.tools import as_mapper, timed_pass
 
 __all__ = ['iet_build']
 
@@ -29,7 +32,12 @@ def iet_build(stree):
             body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs)
 
         elif i.is_Conditional:
-            body = Conditional(i.guard, queues.pop(i))
+            if isinstance(i.guard, GuardSwitch):
+                bundle, = queues.pop(i)
+                cases, nodes = _unpack_switch_case(bundle)
+                body = Switch(i.guard.arg, cases, nodes)
+            else:
+                body = Conditional(i.guard, queues.pop(i))
 
         elif i.is_Iteration:
             if i.dim.is_Virtual:
@@ -55,3 +63,26 @@ def iet_build(stree):
         queues.setdefault(i.parent, []).append(body)
 
     assert False
+
+
+def _unpack_switch_case(bundle):
+    """
+    Helper to unpack an ExpressionBundle containing GuardCaseSwitch expressions
+    into Switch cases and corresponding IET nodes.
+    """
+    assert bundle.is_ExpressionBundle
+    assert all(isinstance(e.rhs, GuardCaseSwitch) for e in bundle.body)
+
+    mapper = as_mapper(bundle.body, key=lambda e: e.rhs.case)
+
+    cases = list(mapper)
+
+    nodes = []
+    for v in mapper.values():
+        exprs = [e._rebuild(expr=e.expr._subs(e.rhs, e.rhs.arg)) for e in v]
+        if len(exprs) > 1:
+            nodes.append(List(body=exprs))
+        else:
+            nodes.append(*exprs)
+
+    return cases, nodes

devito/ir/iet/nodes.py

@@ -30,7 +30,7 @@
            'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration',
            'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma',
            'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace',
-           'Using', 'CallableBody', 'Transfer', 'EmptyList']
+           'Using', 'CallableBody', 'Transfer', 'EmptyList', 'Switch']
 
 # First-class IET nodes
 
@@ -406,6 +406,11 @@ def output(self):
         """The Symbol/Indexed this Expression writes to."""
         return self.expr.lhs
 
+    @property
+    def rhs(self):
+        """The right-hand side of the underlying expression."""
+        return self.expr.rhs
+
     @cached_property
     def reads(self):
         """The Functions read by the Expression."""
@@ -892,6 +897,50 @@ def __repr__(self):
             return "<[%s] ? [%s]" % (ccode(self.condition), repr(self.then_body))
 
 
+class Switch(DoIf):
+
+    """
+    A node to express switch-case blocks.
+
+    Parameters
+    ----------
+    condition : expr-like
+        The switch condition.
+    cases : expr-like or list of expr-like
+        One or more case conditions; there must be one case per node in `nodes`,
+        plus an optional default case.
+    nodes : Node or list of Node
+        One or more Case nodes.
+    default : Node, optional
+        The default case of the switch, if any.
+    """
+
+    _traversable = ['nodes', 'default']
+
+    def __init__(self, condition, cases, nodes, default=None):
+        super().__init__(condition)
+
+        self.cases = as_tuple(cases)
+        self.nodes = as_tuple(nodes)
+        self.default = default
+
+        assert len(self.cases) == len(self.nodes)
+
+    def __repr__(self):
+        return "f<Switch {ccode(self.condition)}; {self.ncases} cases>"
+
+    @property
+    def ncases(self):
+        return len(self.cases) + int(self.default is not None)
+
+    @property
+    def as_mapper(self):
+        retval = dict(zip(self.cases, self.nodes))
+        if self.default:
+            retval['default'] = self.default
+        return retval
+
+
 # Second level IET nodes
 
 class TimedList(List):

devito/ir/iet/visitors.py

@@ -622,6 +622,12 @@ def visit_Conditional(self, o):
         else:
             return c.If(self.ccode(o.condition), then_body)
 
+    def visit_Switch(self, o):
+        condition = self.ccode(o.condition)
+        mapper = {k: self._visit(v) for k, v in o.as_mapper.items()}
+
+        return Switch(condition, mapper)
+
     def visit_Iteration(self, o):
         body = flatten(self._visit(i) for i in self._blankline_logic(o.children))
 
@@ -1426,6 +1432,12 @@ def visit_Conditional(self, o):
         return o._rebuild(condition=condition, then_body=then_body,
                           else_body=else_body)
 
+    def visit_Switch(self, o):
+        condition = uxreplace(o.condition, self.mapper)
+        nodes = self._visit(o.nodes)
+        default = self._visit(o.default)
+        return o._rebuild(condition=condition, nodes=nodes, default=default)
+
     def visit_PointerCast(self, o):
         function = self.mapper.get(o.function, o.function)
         obj = self.mapper.get(o.obj, o.obj)
@@ -1500,6 +1512,31 @@ class LambdaCollection(c.Collection):
     pass
 
 
+class Switch(c.Generable):
+
+    def __init__(self, condition, mapper):
+        self.condition = condition
+
+        # If the `default` case is present, it is encoded with the key "default"
+        self.mapper = mapper
+
+    def generate(self):
+        yield f"switch ({self.condition})"
+        yield "{"
+
+        for case, body in self.mapper.items():
+            if case == "default":
+                yield "  default: {"
+            else:
+                yield f"  case {case}: {{"
+            for line in body.generate():
+                yield f"    {line}"
+            yield "    break;"
+            yield "  }"
+
+        yield "}"
+
+
 class MultilineCall(c.Generable):
 
     def __init__(self, name, arguments, is_expr=False, is_indirect=False,

devito/ir/support/guards.py

@@ -8,20 +8,25 @@
 from operator import ge, gt, le, lt
 
 from functools import singledispatch
-from sympy import And, Ge, Gt, Le, Lt, Mul, true
+from sympy import And, Expr, Ge, Gt, Le, Lt, Mul, true
 from sympy.logic.boolalg import BooleanFunction
 import numpy as np
 
 from devito.ir.support.space import Forward, IterationDirection
-from devito.symbolics import CondEq, CondNe
+from devito.symbolics import CondEq, CondNe, search
 from devito.tools import Pickable, as_tuple, frozendict, split
 from devito.types import Dimension, LocalObject
 
 __all__ = ['GuardFactor', 'GuardBound', 'GuardBoundNext', 'BaseGuardBound',
-           'BaseGuardBoundNext', 'GuardOverflow', 'Guards', 'GuardExpr']
+           'BaseGuardBoundNext', 'GuardOverflow', 'Guards', 'GuardExpr',
+           'GuardSwitch', 'GuardCaseSwitch']
 
 
-class Guard:
+class AbstractGuard:
+    pass
+
+
+class Guard(AbstractGuard):
 
     @property
     def _args_rebuild(self):
@@ -217,6 +222,35 @@ class GuardOverflowLt(BaseGuardOverflow, Lt):
 }
 
 
+class GuardSwitch(AbstractGuard, Expr):
+
+    """
+    A switch guard (akin to C's switch-case) that can be used to select
+    between multiple cases at runtime.
+    """
+
+    def __new__(cls, arg, **kwargs):
+        return Expr.__new__(cls, arg)
+
+    @property
+    def arg(self):
+        return self.args[0]
+
+
+class GuardCaseSwitch(GuardSwitch):
+
+    """
+    A case within a GuardSwitch.
+    """
+
+    def __new__(cls, arg, case, **kwargs):
+        return Expr.__new__(cls, arg, case)
+
+    @property
+    def case(self):
+        return self.args[1]
+
+
 class Guards(frozendict):
 
     """

tests/test_iet.py

@@ -8,9 +8,10 @@
 from devito import (Eq, Grid, Function, TimeFunction, Operator, Dimension,  # noqa
                     switchconfig)
 from devito.ir.iet import (
-    Call, Callable, Conditional, Definition, DeviceCall, DummyExpr, Iteration, List,
-    KernelLaunch, Dereference, Lambda, ElementalFunction, CGen, FindSymbols,
-    filter_iterations, make_efunc, retrieve_iteration_tree, Transformer
+    Call, Callable, Conditional, Definition, DeviceCall, DummyExpr, Iteration,
+    List, KernelLaunch, Dereference, Lambda, Switch, ElementalFunction, CGen,
+    FindSymbols, filter_iterations, make_efunc, retrieve_iteration_tree,
+    Transformer
 )
 from devito.ir import SymbolRegistry
 from devito.passes.iet.engine import Graph
@@ -509,3 +510,31 @@ def test_dereference_base_plus_off():
     deref = Dereference(x, ptr, offset=off)
 
     assert str(deref) == "float (*restrict x)[3] = (float (*)[3]) (p + offs);"
+
+
+def test_switch_case():
+    flag = Symbol(name='flag')
+    a = Symbol(name='a')
+
+    cases = [0, 1]
+    nodes = [DummyExpr(a, 1), DummyExpr(a, 2)]
+    default = DummyExpr(a, 0)
+
+    switch = Switch(flag, cases, nodes, default=default)
+
+    assert str(switch) == """\
+switch (flag)
+{
+  case 0: {
+    a = 1;
+    break;
+  }
+  case 1: {
+    a = 2;
+    break;
+  }
+  default: {
+    a = 0;
+    break;
+  }
+}"""