Merge pull request #2840 from devitocodes/the-TMA

compiler: Enhance IR to support more advanced parlang (CUDA/HIP/SYCL) features

Author: mloubout

devito/arch/archinfo.py

@@ -1139,7 +1139,7 @@ def supports(self, query, language=None):
             warning(f"Couldn't establish if `query={query}` is supported on this "
                     "system. Assuming it is not.")
             return False
-        elif query == 'async-loads' and cc >= 80:
+        elif query == 'async-pipe' and cc >= 80:
             # Asynchronous pipeline loads -- introduced in Ampere
             return True
         elif query in ('tma', 'thread-block-cluster') and cc >= 90:  # noqa: SIM103
@@ -1156,7 +1156,7 @@ class Volta(NvidiaDevice):
 class Ampere(Volta):
 
     def supports(self, query, language=None):
-        if query == 'async-loads':
+        if query == 'async-pipe':
             return True
         else:
             return super().supports(query, language)

devito/ir/cgen/printer.py

@@ -281,6 +281,13 @@ def _print_BitwiseNot(self, expr):
 
     def _print_BitwiseBinaryOp(self, expr):
         arg0, arg1 = expr.args
+
+        prec = precedence(expr)
+        if not arg0.is_Atom:
+            arg0 = self.parenthesize(arg0, prec)
+        if not arg1.is_Atom:
+            arg1 = self.parenthesize(arg1, prec)
+
         return f'{self._print(arg0)} {expr.op} {self._print(arg1)}'
 
     def _print_Add(self, expr, order=None):

devito/ir/clusters/cluster.py

@@ -171,16 +171,35 @@ def free_symbols(self):
     def dimensions(self):
         return set().union(*[i._defines for i in self.ispace.dimensions])
 
+    @cached_property
+    def exprs_dimensions(self):
+        """
+        The Dimensions that appear explicitly in the Cluster expressions.
+        """
+        dims_explicit = {i for i in self.free_symbols if i.is_Dimension}
+        dims_implicit = {d for e in self.exprs for d in e.implicit_dims}
+        return dims_explicit | dims_implicit
+
+    @cached_property
+    def guards_dimensions(self):
+        """
+        The Dimensions that appear explicitly in the Cluster guards.
+        """
+        syms_guards = {d for e in self.guards.values() for d in e.free_symbols}
+        dims_guards = {i for i in syms_guards if i.is_Dimension}
+        return dims_guards
+
     @cached_property
     def used_dimensions(self):
         """
-        The Dimensions that *actually* appear among the expressions in ``self``.
-        These do not necessarily coincide the IterationSpace Dimensions; for
-        example, reduction or redundant (i.e., invariant) Dimensions won't
-        appear in an expression.
+        All the Dimensions that appear explicitly either within the expressions
+        or the guards.
+
+        Note that, in some cases, some of the IterationSpace Dimensions might
+        not appear here among the used Dimensions -- for example, reduction or
+        redundant (i.e., invariant) Dimensions.
         """
-        idims = set.union(*[set(e.implicit_dims) for e in self.exprs])
-        return {i for i in self.free_symbols if i.is_Dimension} | idims
+        return self.exprs_dimensions | self.guards_dimensions
 
     @cached_property
     def dist_dimensions(self):

devito/ir/equations/equation.py

@@ -92,9 +92,11 @@ def __repr__(self):
         if not self.is_Reduction:
             return super().__repr__()
         elif self.operation is OpInc:
-            return f'{self.lhs} += {self.rhs}'
+            return f'Inc({self.lhs}, {self.rhs})'
         else:
-            return f'{self.lhs} = {self.operation}({self.rhs})'
+            return f'Eq({self.lhs}, {self.operation}({self.rhs}))'
+
+    __str__ = __repr__
 
     # Pickling support
     __reduce_ex__ = Pickable.__reduce_ex__

devito/ir/iet/visitors.py

@@ -321,6 +321,9 @@ def _gen_value(self, obj, mode=1, masked=()):
         qualifiers = [v for k, v in self._qualifiers_mapper.items()
                       if getattr(obj.function, k, False) and v not in masked]
 
+        if obj.is_LocalObject and mode == 2:
+            qualifiers.extend(as_tuple(obj._C_tag))
+
         if (obj._mem_stack or obj._mem_constant) and mode == 1:
             strtype = self.ccode(obj._C_typedata)
             strshape = ''.join(f'[{self.ccode(i)}]' for i in obj.symbolic_shape)

devito/ir/support/guards.py

@@ -500,7 +500,9 @@ def pairwise_or(*guards):
 
     # Analysis
     for guard in guards:
-        if guard is true or guard is None:
+        if guard is true:
+            return true
+        elif guard is None:
             continue
         elif isinstance(guard, And):
             components = guard.args

devito/ir/support/space.py

@@ -789,17 +789,19 @@ def __init__(self, intervals, sub_iterators=None, directions=None):
         super().__init__(intervals)
 
         # Normalize sub-iterators
-        sub_iterators = dict([(k, tuple(filter_ordered(as_tuple(v))))
-                              for k, v in (sub_iterators or {}).items()])
+        sub_iterators = sub_iterators or {}
+        sub_iterators = {d: tuple(filter_ordered(as_tuple(v)))
+                         for d, v in sub_iterators.items() if d in self.intervals}
         sub_iterators.update({i.dim: () for i in self.intervals
                               if i.dim not in sub_iterators})
         self._sub_iterators = frozendict(sub_iterators)
 
         # Normalize directions
-        if directions is None:
-            self._directions = frozendict([(i.dim, Any) for i in self.intervals])
-        else:
-            self._directions = frozendict(directions)
+        directions = directions or {}
+        directions = {d: v for d, v in directions.items() if d in self.intervals}
+        directions.update({i.dim: Any for i in self.intervals
+                           if i.dim not in directions})
+        self._directions = frozendict(directions)
 
     def __repr__(self):
         ret = ', '.join([f"{repr(i)}{repr(self.directions[i.dim])}"
@@ -821,8 +823,7 @@ def __lt__(self, other):
         return len(self.itintervals) < len(other.itintervals)
 
     def __hash__(self):
-        return hash((super().__hash__(), self.sub_iterators,
-                     self.directions))
+        return hash((super().__hash__(), self.sub_iterators, self.directions))
 
     def __contains__(self, d):
         try:

devito/passes/clusters/cse.py

@@ -13,9 +13,8 @@
 
 from devito.finite_differences.differentiable import IndexDerivative
 from devito.ir import Cluster, Scope, cluster_pass
-from devito.symbolics import estimate_cost, q_leaf, q_terminal
+from devito.symbolics import Reserved, estimate_cost, q_leaf, q_terminal, search
 from devito.symbolics.manipulation import _uxreplace
-from devito.symbolics.search import search
 from devito.tools import DAG, as_list, as_tuple, extract_dtype, frozendict
 from devito.types import Eq, Symbol, Temp
 
@@ -411,6 +410,7 @@ def _(expr):
 
 @_catch.register(Indexed)
 @_catch.register(Symbol)
+@_catch.register(Reserved)
 def _(expr):
     """
     Handler for objects preventing CSE to propagate through their arguments.

devito/passes/clusters/derivatives.py

@@ -3,7 +3,7 @@
 import numpy as np
 from sympy import S
 
-from devito.finite_differences import IndexDerivative
+from devito.finite_differences import IndexDerivative, Weights
 from devito.ir import Backward, Forward, Interval, IterationSpace, Queue
 from devito.passes.clusters.misc import fuse
 from devito.symbolics import BasicWrapperMixin, reuse_if_untouched, uxreplace
@@ -91,17 +91,39 @@ def _core(expr, c, ispace, weights, reusables, mapper, **kwargs):
 
 
 @_core.register(Symbol)
-@_core.register(Indexed)
 @_core.register(BasicWrapperMixin)
 def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     return expr, []
 
 
+@_core.register(Indexed)
+def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
+    if not isinstance(expr.function, Weights):
+        return expr, []
+
+    # Lower or reuse a previously lowered Weights array
+    sregistry = kwargs['sregistry']
+    subs_user = kwargs['subs']
+
+    w0 = expr.function
+    k = tuple(w0.weights)
+    try:
+        w = weights[k]
+    except KeyError:
+        name = sregistry.make_name(prefix='w')
+        dtype = infer_dtype([w0.dtype, c.dtype])  # At least np.float32
+        initvalue = tuple(i.subs(subs_user) for i in k)
+        w = weights[k] = w0._rebuild(name=name, dtype=dtype, initvalue=initvalue)
+
+    rebuilt = expr._subs(w0.indexed, w.indexed)
+
+    return rebuilt, []
+
+
 @_core.register(IndexDerivative)
 def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     sregistry = kwargs['sregistry']
     options = kwargs['options']
-    subs_user = kwargs['subs']
 
     try:
         cbk0 = deriv_schedule_registry[options['deriv-schedule']]
@@ -114,18 +136,10 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
 
     # Create the concrete Weights array, or reuse an already existing one
     # if possible
-    name = sregistry.make_name(prefix='w')
-    w0 = ideriv.weights.function
-    dtype = infer_dtype([w0.dtype, c.dtype])  # At least np.float32
-    k = tuple(w0.weights)
-    try:
-        w = weights[k]
-    except KeyError:
-        initvalue = tuple(i.subs(subs_user) for i in k)
-        w = weights[k] = w0._rebuild(name=name, dtype=dtype, initvalue=initvalue)
+    w, _ = _core(ideriv.weights, c, ispace, weights, reusables, mapper, **kwargs)
 
     # Replace the abstract Weights array with the concrete one
-    subs = {w0.indexed: w.indexed}
+    subs = {ideriv.weights.base: w.base}
     init = uxreplace(init, subs)
     ideriv = uxreplace(ideriv, subs)
 
@@ -152,13 +166,13 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     ispace1 = IterationSpace.union(ispace, ispace0, relations=extra)
 
     # The Symbol that will hold the result of the IndexDerivative computation
-    # NOTE: created before recurring so that we ultimately get a sound ordering
+    # NOTE: created before recursing so that we ultimately get a sound ordering
     try:
         s = reusables.pop()
-        assert np.can_cast(s.dtype, dtype)
+        assert np.can_cast(s.dtype, w.dtype)
     except KeyError:
         name = sregistry.make_name(prefix='r')
-        s = Symbol(name=name, dtype=dtype)
+        s = Symbol(name=name, dtype=w.dtype)
 
     # Go inside `expr` and recursively lower any nested IndexDerivatives
     expr, processed = _core(expr, c, ispace1, weights, reusables, mapper, **kwargs)

devito/passes/clusters/misc.py

@@ -55,7 +55,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Is `c` a real candidate -- is there at least one invariant Dimension?
-            if any(d._defines & hope_invariant for d in c.used_dimensions):
+            if any(d._defines & hope_invariant for d in c.exprs_dimensions):
                 processed.append(c)
                 continue
 
@@ -69,16 +69,16 @@ def callback(self, clusters, prefix):
             # All of the inner Dimensions must appear in the write-to region
             # otherwise we would violate data dependencies. Consider
             #
-            # 1)                 2)                          3)
-            # for i              for i                       for i
-            #   for x              for x                       for x
-            #     r = f(a[x])        for y                       for y
-            #                          r[x] = f(a[x, y])           r[x, y] = f(a[x, y])
+            # 1)                2)                        3)
+            # for i             for i                     for i
+            #   for x             for x                     for x
+            #     r = f(a[x])       for y                     for y
+            #                         r[x] = f(a[x, y])         r[x, y] = f(a[x, y])
             #
             # In 1) and 2) lifting is infeasible; in 3) the statement can
             # be lifted outside the `i` loop as `r`'s write-to region contains
             # both `x` and `y`
-            xed = {d._defines for d in c.used_dimensions if d not in outer}
+            xed = {d._defines for d in c.exprs_dimensions if d not in outer}
             if not all(i & set(w.dimensions) for i, w in product(xed, c.scope.writes)):
                 processed.append(c)
                 continue