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

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/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 = set().union(*[set(e.implicit_dims) for e in self.exprs])
+        return dims_explicit | dims_implicit
+
+    @cached_property
+    def guards_dimensions(self):
+        """
+        The Dimensions that appear explicitly in the Cluster guards.
+        """
+        syms_guards = set().union(*[e.free_symbols for e in self.guards.values()])
+        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
 
@@ -399,6 +398,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

devito/passes/iet/definitions.py

@@ -5,7 +5,6 @@
 
 from collections import OrderedDict
 from ctypes import c_uint64
-from functools import singledispatch
 from operator import itemgetter
 
 import numpy as np
@@ -98,14 +97,29 @@ def _alloc_object_on_low_lat_mem(self, site, obj, storage):
         """
         decl = Definition(obj)
 
-        definition = (decl, obj._C_init) if obj._C_init else (decl)
+        init = obj._C_init
+        if not init:
+            definition = decl
+            efuncs = ()
+        elif isinstance(init, (list, tuple)):
+            assert len(init) == 2, "Expected (efunc, call)"
+            init, definition = init
+            efuncs = (init,)
+        elif init.is_Callable:
+            definition = Call(init.name, init.parameters,
+                              retobj=obj if init.retval else None)
+            efuncs = (init,)
+        else:
+            definition = (decl, init)
+            efuncs = ()
 
         frees = obj._C_free
 
         if obj.free_symbols - {obj}:
-            storage.update(obj, site, objs=definition, frees=frees)
+            storage.update(obj, site, objs=definition, efuncs=efuncs, frees=frees)
         else:
-            storage.update(obj, site, standalones=definition, frees=frees)
+            storage.update(obj, site, standalones=definition, efuncs=efuncs,
+                           frees=frees)
 
     def _alloc_array_on_low_lat_mem(self, site, obj, storage):
         """
@@ -552,7 +566,7 @@ class DeviceAwareDataManager(DataManager):
     def __init__(self, options=None, **kwargs):
         self.gpu_fit = options['gpu-fit']
         self.gpu_create = options['gpu-create']
-        self.pmode = options.get('place-transfers')
+        self.gpu_place_transfers = options.get('place-transfers')
 
         super().__init__(**kwargs)
 
@@ -585,7 +599,8 @@ def _map_array_on_high_bw_mem(self, site, obj, storage):
 
         storage.update(obj, site, maps=mmap, unmaps=unmap)
 
-    def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm, read_only=False):
+    def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm,
+                                     read_only=False, **kwargs):
         """
         Map a Function already defined in the host memory in to the device high
         bandwidth memory.
@@ -618,42 +633,41 @@ def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm, read_only=F
         storage.update(obj, site, maps=mmap, unmaps=unmap, efuncs=efuncs)
 
     @iet_pass
-    def place_transfers(self, iet, data_movs=None, **kwargs):
+    def place_transfers(self, iet, data_movs=None, ctx=None, **kwargs):
         """
         Create a new IET with host-device data transfers. This requires mapping
         symbols to the suitable memory spaces.
         """
-        if not self.pmode:
+        if not self.gpu_place_transfers:
             return iet, {}
 
-        @singledispatch
-        def _place_transfers(iet, data_movs):
+        if not isinstance(iet, EntryFunction):
             return iet, {}
 
-        @_place_transfers.register(EntryFunction)
-        def _(iet, data_movs):
-            reads, writes = data_movs
-
-            # Special symbol which gives user code control over data deallocations
-            devicerm = DeviceRM()
+        reads, writes = data_movs
 
-            storage = Storage()
-            for i in filter_sorted(writes):
-                if i.is_Array:
-                    self._map_array_on_high_bw_mem(iet, i, storage)
-                else:
-                    self._map_function_on_high_bw_mem(iet, i, storage, devicerm)
-            for i in filter_sorted(reads - writes):
-                if i.is_Array:
-                    self._map_array_on_high_bw_mem(iet, i, storage)
-                else:
-                    self._map_function_on_high_bw_mem(iet, i, storage, devicerm, True)
+        # Special symbol which gives user code control over data deallocations
+        devicerm = DeviceRM()
 
-            iet, efuncs = self._inject_definitions(iet, storage)
+        storage = Storage()
+        for i in filter_sorted(writes):
+            if i.is_Array:
+                self._map_array_on_high_bw_mem(iet, i, storage)
+            else:
+                self._map_function_on_high_bw_mem(
+                    iet, i, storage, devicerm, ctx=ctx
+                )
+        for i in filter_sorted(reads - writes):
+            if i.is_Array:
+                self._map_array_on_high_bw_mem(iet, i, storage)
+            else:
+                self._map_function_on_high_bw_mem(
+                    iet, i, storage, devicerm, read_only=True, ctx=ctx
+                )
 
-            return iet, {'efuncs': efuncs}
+        iet, efuncs = self._inject_definitions(iet, storage)
 
-        return _place_transfers(iet, data_movs=data_movs)
+        return iet, {'efuncs': efuncs}
 
     @iet_pass
     def place_devptr(self, iet, **kwargs):