Merge pull request #2776 from devitocodes/guard-pairwise-or

compiler: Add Guards.pairwise_or

Author: FabioLuporini

devito/arch/archinfo.py

@@ -957,6 +957,23 @@ def _detect_isa(self):
 
 class Device(Platform):
 
+    """
+    A generic Device is based on the SIMT (Single Instruction, Multiple Threads)
+    programming model. In this execution model, threads are batched together and
+    execute the same instruction at the same time, though each thread operates on
+    its own data. Intel, AMD, and Nvidia GPUs are all based on this model.
+    Unfortunately they use different terminology to refer to the same or at least
+    very similar concepts. Throughout Devito, whenever possible, we attempt to
+    adopt a neutral terminology -- the docstrings below provide some examples.
+    """
+
+    thread_group_size = None
+    """
+    A collection of threads that execute the same instruction in lockstep.
+    The group size is a hardware-specific property. For example, this is a
+    "warp" in NVidia GPUs and a "wavefront" in AMD GPUs.
+    """
+
     def __init__(self, name, cores_logical=None, cores_physical=None, isa='cpp',
                  max_threads_per_block=1024, max_threads_dimx=1024,
                  max_threads_dimy=1024, max_threads_dimz=64,
@@ -1039,6 +1056,8 @@ def march(self):
 
 class NvidiaDevice(Device):
 
+    thread_group_size = 32
+
     max_mem_trans_nbytes = 128
 
     @cached_property
@@ -1102,6 +1121,8 @@ class Blackwell(Hopper):
 
 class AmdDevice(Device):
 
+    thread_group_size = 64
+
     max_mem_trans_nbytes = 256
 
     @cached_property

devito/ir/cgen/printer.py

@@ -276,6 +276,15 @@ def _print_Abs(self, expr):
             return f"fabs({self._print(arg)})"
         return self._print_fmath_func('abs', expr)
 
+    def _print_BitwiseNot(self, expr):
+        # Unary function, single argument
+        arg = expr.args[0]
+        return f'~{self._print(arg)}'
+
+    def _print_BitwiseBinaryOp(self, expr):
+        arg0, arg1 = expr.args
+        return f'{self._print(arg0)} {expr.op} {self._print(arg1)}'
+
     def _print_Add(self, expr, order=None):
         """"
         Print an addition.

devito/ir/clusters/cluster.py

@@ -69,13 +69,23 @@ def from_clusters(cls, *clusters):
         """
         assert len(clusters) > 0
         root = clusters[0]
+
+        if len(clusters) == 1:
+            return root
+
         if not all(root.ispace.is_compatible(c.ispace) for c in clusters):
             raise ValueError("Cannot build a Cluster from Clusters with "
                              "incompatible IterationSpace")
         if not all(root.guards == c.guards for c in clusters):
             raise ValueError("Cannot build a Cluster from Clusters with "
                              "non-homogeneous guards")
 
+        writes = set().union(*[c.scope.writes for c in clusters])
+        reads = set().union(*[c.scope.reads for c in clusters])
+        if any(f._mem_shared for f in writes & reads):
+            raise ValueError("Cannot build a Cluster from Clusters with "
+                             "read-write conflicts on shared-memory Functions")
+
         exprs = chain(*[c.exprs for c in clusters])
         ispace = IterationSpace.union(*[c.ispace for c in clusters])
 

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,