compiler: Revamp to enable blocking before CIRE

Author: FabioLuporini

devito/core/cpu.py

@@ -5,8 +5,8 @@
 from devito.operator.operator import rcompile
 from devito.passes import stream_dimensions
 from devito.passes.clusters import (
-    Lift, blocking, buffering, cire, cse, factorize, fission, fuse, optimize_hyperplanes,
-    optimize_pows
+    Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse,
+    optimize_hyperplanes, optimize_pows
 )
 from devito.passes.equations import collect_derivatives
 from devito.passes.iet import (
@@ -67,6 +67,7 @@ def _normalize_kwargs(cls, **kwargs):
                                 reduce=oo.pop('par-tile-reduce', None))
 
         # CIRE
+        o['cire-block-temps'] = oo.pop('cire-block-temps', cls.CIRE_BLOCK_TEMPS)
         o['min-storage'] = oo.pop('min-storage', False)
         o['cire-rotate'] = oo.pop('cire-rotate', False)
         o['cire-maxpar'] = oo.pop('cire-maxpar', False)
@@ -198,6 +199,9 @@ def _specialize_clusters(cls, clusters, **kwargs):
         if options['blocklazy']:
             clusters = blocking(clusters, sregistry, options)
 
+        # Unfold the `par-tile`s, if any
+        clusters = apply_par_tiles(clusters, **kwargs)
+
         return clusters
 
     @classmethod

devito/core/gpu.py

@@ -7,8 +7,8 @@
 from devito.operator.operator import rcompile
 from devito.passes import is_on_device, stream_dimensions
 from devito.passes.clusters import (
-    Lift, blocking, buffering, cire, cse, factorize, fission, fuse, memcpy_prefetch,
-    optimize_pows, tasking
+    Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse,
+    memcpy_prefetch, optimize_pows, tasking
 )
 from devito.passes.equations import collect_derivatives
 from devito.passes.iet import (
@@ -38,7 +38,9 @@
 
 class DeviceOperatorMixin:
 
+    # Overrides the default values in the main Operator class
     BLOCK_LEVELS = 0
+    CIRE_BLOCK_TEMPS = False
     MPI_MODES = (True, 'basic',)
 
     GPU_FIT = 'all-fallback'
@@ -76,6 +78,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['skewing'] = oo.pop('skewing', False)
 
         # CIRE
+        o['cire-block-temps'] = oo.pop('cire-block-temps', cls.CIRE_BLOCK_TEMPS)
         o['min-storage'] = False
         o['cire-rotate'] = False
         o['cire-maxpar'] = oo.pop('cire-maxpar', 'basic')
@@ -239,6 +242,9 @@ def _specialize_clusters(cls, clusters, **kwargs):
         if options['blocklazy']:
             clusters = blocking(clusters, sregistry, options)
 
+        # Unfold the `par-tile`s, if any
+        clusters = apply_par_tiles(clusters, **kwargs)
+
         return clusters
 
     @classmethod

devito/core/operator.py

@@ -58,6 +58,13 @@ class BasicOperator(Operator):
     situations where the performance impact might be detrimental.
     """
 
+    CIRE_BLOCK_TEMPS = True
+    """
+    If an aliasing expression is computed within a blocked loop nest, all CIRE-
+    generated temporaries will inherit the block shape. If set to False, the
+    temporaries shape will systematically be defined by the root Dimensions.
+    """
+
     CIRE_MINGAIN = 10
     """
     Minimum operation count reduction for a redundant expression to be optimized

devito/ir/clusters/cluster.py

@@ -13,7 +13,7 @@
 )
 from devito.mpi.halo_scheme import HaloScheme, HaloTouch
 from devito.mpi.reduction_scheme import DistReduce
-from devito.symbolics import estimate_cost
+from devito.symbolics import estimate_cost, uxreplace
 from devito.tools import as_tuple, filter_ordered, flatten, infer_dtype
 from devito.types import (
     CriticalRegion, Fence, Indexed, PhaseMarker, TensorMove, ThreadArrive, ThreadCommit,
@@ -128,6 +128,33 @@ def rebuild(self, *args, **kwargs):
                               syncs=kwargs.get('syncs', self.syncs),
                               halo_scheme=kwargs.get('halo_scheme', self.halo_scheme))
 
+    def subs(self, mapper, compact=()):
+        """
+        Build a new Cluster applying substitutions rules to `self`.
+        """
+        if not mapper:
+            return self
+
+        if self.halo_scheme:
+            raise NotImplementedError
+
+        key0 = lambda i: i.is_Block
+        subs0 = {d: self.ispace[d].promote(key0).dim for d in compact}
+
+        subs = {**mapper, **subs0}
+        exprs = [uxreplace(e, subs) for e in self.exprs]
+
+        ispace = self.ispace.switch(mapper)
+        key = lambda i: key0(i) and i in flatten(d._defines for d in subs0)
+        ispace = ispace.promote(key, mode='total')
+
+        guards = self.guards.subs(mapper).promote(subs0)
+        properties = self.properties.subs(mapper).promote(subs0)
+        syncs = self.syncs.subs(mapper)
+
+        return self.__class__(exprs=exprs, ispace=ispace, guards=guards,
+                              properties=properties, syncs=syncs)
+
     @property
     def exprs(self):
         return self._exprs
@@ -591,6 +618,14 @@ def dspace(self):
         """Return the DataSpace of this ClusterGroup."""
         return DataSpace.union(*[i.dspace.reset() for i in self])
 
+    @property
+    def is_dense(self):
+        return all(i.is_dense for i in self)
+
+    @property
+    def is_wild(self):
+        return all(i.is_wild for i in self)
+
     @property
     def is_halo_touch(self):
         return all(i.is_halo_touch for i in self)

devito/ir/support/guards.py

@@ -342,6 +342,31 @@ def as_map(self, d, cls):
 
         return dict(i.args for i in search(self.get(d), cls))
 
+    def subs(self, mapper):
+        m = {mapper.get(d, d): v.xreplace(mapper) for d, v in self.items()}
+
+        return Guards(m)
+
+    def promote(self, subs):
+        m = self
+        for d, v in subs.items():
+            guards = {self.get(i) for i in d._defines} - {true}
+            if len(guards) > 1:
+                raise NotImplementedError(
+                    f"Cannot promote {d} to {v} due to multiple guards: {guards}"
+                )
+            elif len(guards) == 0:
+                continue
+
+            guard = guards.pop()
+            guard = guard.xreplace({d: v})
+
+            m = m.impose(v, guard)
+
+        m = m.popany(subs)
+
+        return m
+
 
 class GuardExpr(LocalObject, BooleanFunction):
 

devito/ir/support/properties.py

@@ -39,10 +39,11 @@ def __init__(self, name, val=None):
 TILABLE = Property('tilable')
 """A fully parallel Dimension that would benefit from tiling (or "blocking")."""
 
-TILABLE_SMALL = Property('tilable*')
+NO_TUNING = Property('notuning')
 """
-Like TILABLE, but it would benefit from relatively small block, since the
-iteration space is likely to be very small.
+A Dimension that would be unlikely to benefit from tuning. For example, the
+underlying iteration space is relatively small, and/or the enclosed expressions
+are not characterized by data reuse, etc.
 """
 
 SKEWABLE = Property('skewable')
@@ -115,7 +116,6 @@ def __init__(self, name, val=None):
 
 # Bundles
 PARALLELS = {PARALLEL, PARALLEL_INDEP, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT}
-TILABLES = {TILABLE, TILABLE_SMALL}
 
 
 def normalize_properties(*args):
@@ -253,16 +253,10 @@ def prefetchable(self, dims, v=PREFETCHABLE):
             m[d] = self.get(d, set()) | {v}
         return Properties(m)
 
-    def block(self, dims, kind='default'):
-        if kind == 'default':
-            p = TILABLE
-        elif kind == 'small':
-            p = TILABLE_SMALL
-        else:
-            raise ValueError
+    def block(self, dims):
         m = dict(self)
         for d in as_tuple(dims):
-            m[d] = set(self.get(d, [])) | {p}
+            m[d] = set(self.get(d, [])) | {TILABLE}
         return Properties(m)
 
     def inbound(self, dims):
@@ -289,6 +283,9 @@ def init_halo_right_shm(self, dims):
                                                  INIT_HALO_LEFT_SHM})
         return properties
 
+    def notune(self, dims):
+        return self.add(dims, NO_TUNING)
+
     def is_parallel(self, dims):
         return any(len(self[d] & {PARALLEL, PARALLEL_INDEP}) > 0
                    for d in as_tuple(dims))
@@ -310,10 +307,7 @@ def is_sequential(self, dims):
         return any(SEQUENTIAL in self.get(d, ()) for d in as_tuple(dims))
 
     def is_blockable(self, d):
-        return bool(self.get(d, set()) & {TILABLE, TILABLE_SMALL})
-
-    def is_blockable_small(self, d):
-        return TILABLE_SMALL in self.get(d, set())
+        return bool(TILABLE in self.get(d, ()))
 
     def _is_property_any(self, dims, v):
         if dims is None:
@@ -335,6 +329,25 @@ def is_halo_right_init(self, dims=None):
     def is_halo_init(self, dims=None):
         return self.is_halo_left_init(dims) or self.is_halo_right_init(dims)
 
+    def avoid_tuning(self, dims):
+        return any(NO_TUNING in self.get(d, set()) for d in as_tuple(dims))
+
+    def subs(self, mapper):
+        return Properties({mapper.get(d, d): v for d, v in self.items()})
+
+    def promote(self, subs):
+        m = self
+        for d, pd in subs.items():
+            if pd not in d._defines:
+                raise ValueError(f"Cannot promote {d} to {pd} as {pd} does not "
+                                 f"belong to {d}'s hierarchy")
+
+            v = normalize_properties(*[self.get(i, set()) for i in d._defines])
+
+            m = self.drop(d._defines).add(pd, v)
+
+        return m
+
     @property
     def nblockable(self):
         return sum([self.is_blockable(d) for d in self])

devito/ir/support/space.py

@@ -478,7 +478,12 @@ def relaxed(self, d=None):
 
     def promote(self, cond, mode='unordered'):
         intervals = [i.promote(cond) for i in self]
-        intervals = IntervalGroup(intervals, relations=None, mode=mode)
+
+        if mode == 'total':
+            relations = [filter_ordered(i.dim for i in intervals)]
+        else:
+            relations = None
+        intervals = IntervalGroup(intervals, relations=relations, mode=mode)
 
         # There could be duplicate Dimensions at this point, so we sum up the
         # Intervals defined over the same Dimension to produce a well-defined
@@ -899,16 +904,29 @@ def augment(self, sub_iterators):
 
         return IterationSpace(self.intervals, items, self.directions)
 
-    def switch(self, d0, d1, direction=None):
-        intervals = self.intervals.switch(d0, d1)
+    def switch(self, d0, d1=None, direction=None):
+        """
+        Construct a new IterationSpace in which the Dimension `d0` is replaced with
+        `d1`. Optionally, `d0` could be a mapper, in which case multiple Dimensions
+        may be switched.
+        """
+        if isinstance(d0, dict):
+            mapper = d0
+        else:
+            mapper = {d0: d1}
 
+        intervals = self.intervals
         sub_iterators = dict(self.sub_iterators)
-        sub_iterators.pop(d0, None)
-        sub_iterators[d1] = ()
-
         directions = dict(self.directions)
-        v = directions.pop(d0, None)
-        directions[d1] = direction or v
+
+        for d0, d1 in mapper.items():
+            intervals = intervals.switch(d0, d1)
+
+            sub_iterators.pop(d0, None)
+            sub_iterators[d1] = ()
+
+            v = directions.pop(d0, None)
+            directions[d1] = direction or v
 
         return IterationSpace(intervals, sub_iterators, directions)
 

devito/ir/support/syncs.py

@@ -175,6 +175,9 @@ def update(self, ops):
             m[d] = set(self.get(d, [])) | set(v)
         return Ops(m)
 
+    def subs(self, mapper):
+        return Ops({mapper.get(d, d): v for d, v in self.items()})
+
     def _get_sync(self, cls, dims=None):
         if dims is None:
             dims = list(self)