Merge pull request #2538 from devitocodes/thread-block-clusters

compiler: Misc extensions towards supporting thread block clustering

Author: FabioLuporini

devito/arch/archinfo.py

@@ -684,6 +684,12 @@ def limits(self, compiler=None, language=None):
             'max-block-dims': sys.maxsize,
         }
 
+    def supports(self, query, language=None):
+        """
+        Return True if the platform supports a given feature, False otherwise.
+        """
+        return False
+
 
 class Cpu64(Platform):
 
@@ -836,7 +842,8 @@ class Device(Platform):
 
     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):
+                 max_threads_dimy=1024, max_threads_dimz=64,
+                 max_thread_block_cluster_size=8):
         super().__init__(name)
 
         cpu_info = get_cpu_info()
@@ -849,6 +856,7 @@ def __init__(self, name, cores_logical=None, cores_physical=None, isa='cpp',
         self.max_threads_dimx = max_threads_dimx
         self.max_threads_dimy = max_threads_dimy
         self.max_threads_dimz = max_threads_dimz
+        self.max_thread_block_cluster_size = max_thread_block_cluster_size
 
     @classmethod
     def _mro(cls):
@@ -897,12 +905,6 @@ def limits(self, compiler=None, language=None):
             'max-block-dims': 3,
         }
 
-    def supports(self, query, language=None):
-        """
-        Check if the device supports a given feature.
-        """
-        return False
-
 
 class IntelDevice(Device):
 
@@ -939,7 +941,7 @@ def supports(self, query, language=None):
         if query == 'async-loads' and cc >= 80:
             # Asynchronous pipeline loads -- introduced in Ampere
             return True
-        elif query == 'tma' and cc >= 90:
+        elif query in ('tma', 'thread-block-cluster') and cc >= 90:
             # Tensor Memory Accelerator -- introduced in Hopper
             return True
         else:
@@ -953,25 +955,23 @@ class Volta(NvidiaDevice):
 class Ampere(Volta):
 
     def supports(self, query, language=None):
-        if language != 'cuda':
-            return False
-
         if query == 'async-loads':
             return True
-
-        return super().supports(query, language)
+        else:
+            return super().supports(query, language)
 
 
 class Hopper(Ampere):
 
-    def supports(self, query, language=None):
-        if language != 'cuda':
-            return False
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault('max_thread_block_cluster_size', 16)
+        super().__init__(*args, **kwargs)
 
-        if query == 'tma':
+    def supports(self, query, language=None):
+        if query in ('tma', 'thread-block-cluster'):
             return True
-
-        return super().supports(query, language)
+        else:
+            return super().supports(query, language)
 
 
 class Blackwell(Hopper):

devito/ir/iet/efunc.py

@@ -159,9 +159,9 @@ class DeviceFunction(Callable):
     """
 
     def __init__(self, name, body, retval='void', parameters=None,
-                 prefix='__global__', templates=None):
+                 prefix='__global__', templates=None, attributes=None):
         super().__init__(name, body, retval, parameters=parameters, prefix=prefix,
-                         templates=templates)
+                         templates=templates, attributes=attributes)
 
 
 class DeviceCall(Call):

devito/ir/iet/nodes.py

@@ -718,17 +718,21 @@ class Callable(Node):
     parameters : list of Basic, optional
         The objects in input to the Callable.
     prefix : list of str, optional
-        Qualifiers to prepend to the Callable signature. None by defaults.
+        Qualifiers to prepend to the Callable signature. None by default.
     templates : list of Basic, optional
         The template parameters of the Callable.
+    attributes : list of str, optional
+        Additional attributes to append to the Callable signature. An
+        attributes is one or more keywords that appear in between the
+        return type and the function name. None by default.
     """
 
     is_Callable = True
 
     _traversable = ['body']
 
     def __init__(self, name, body, retval, parameters=None, prefix=None,
-                 templates=None):
+                 templates=None, attributes=None):
         self.name = name
         if not isinstance(body, CallableBody):
             self.body = CallableBody(body)
@@ -738,6 +742,7 @@ def __init__(self, name, body, retval, parameters=None, prefix=None,
         self.prefix = as_tuple(prefix)
         self.parameters = as_tuple(parameters)
         self.templates = as_tuple(templates)
+        self.attributes = as_tuple(attributes)
 
     def __repr__(self):
         param_types = [ctypes_to_cstr(i._C_ctype) for i in self.parameters]

devito/ir/iet/visitors.py

@@ -315,14 +315,25 @@ def _args_call(self, args):
 
     def _gen_signature(self, o, is_declaration=False):
         decls = self._args_decl(o.parameters)
+
         prefix = ' '.join(o.prefix + (self._gen_rettype(o.retval),))
-        signature = c.FunctionDeclaration(c.Value(prefix, o.name), decls)
+
+        if o.attributes:
+            # NOTE: ugly, but I can't bother extending `c.FunctionDeclaration`
+            # for such a tiny thing
+            v = f"{' '.join(o.attributes)} {o.name}"
+        else:
+            v = o.name
+
+        signature = c.FunctionDeclaration(c.Value(prefix, v), decls)
+
         if o.templates:
             tparams = ', '.join([i.inline() for i in self._args_decl(o.templates)])
             if is_declaration:
                 signature = TemplateDecl(tparams, signature)
             else:
                 signature = c.Template(tparams, signature)
+
         return signature
 
     def _blankline_logic(self, children):

devito/ir/support/properties.py

@@ -95,6 +95,21 @@ def __init__(self, name, val=None):
 A Dimension along which shared-memory prefetching is feasible and beneficial.
 """
 
+INIT_CORE_SHM = Property('init-core-shm')
+"""
+A Dimension along which the shared-memory CORE data region is initialized.
+"""
+
+INIT_HALO_LEFT_SHM = Property('init-halo-left-shm')
+"""
+A Dimension along which the shared-memory left-HALO data region is initialized.
+"""
+
+INIT_HALO_RIGHT_SHM = Property('init-halo-right-shm')
+"""
+A Dimension along which the shared-memory right-HALO data region is initialized.
+"""
+
 
 # Bundles
 PARALLELS = {PARALLEL, PARALLEL_INDEP, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT}
@@ -122,9 +137,13 @@ def normalize_properties(*args):
     else:
         drop = set()
 
-    # SEPARABLE <=> all are SEPARABLE
-    if not all(SEPARABLE in p for p in args):
-        drop.add(SEPARABLE)
+    # A property X must be dropped if not all set of properties in `args`
+    # contain X. For example, if one set of properties contains SEPARABLE and
+    # another does not, then the resulting set of properties should not contain
+    # SEPARABLE.
+    for i in (SEPARABLE, INIT_CORE_SHM, INIT_HALO_LEFT_SHM, INIT_HALO_RIGHT_SHM):
+        if not all(i in p for p in args):
+            drop.add(i)
 
     properties = set()
     for p in args:
@@ -190,6 +209,11 @@ def update_properties(properties, exprs):
     else:
         properties = properties.drop(properties=PREFETCHABLE_SHM)
 
+    # Remove properties that are trivially incompatible with `exprs`
+    if not all(e.lhs.function._mem_shared for e in as_tuple(exprs)):
+        drop = {INIT_CORE_SHM, INIT_HALO_LEFT_SHM, INIT_HALO_RIGHT_SHM}
+        properties = properties.drop(properties=drop)
+
     return properties
 
 
@@ -269,10 +293,16 @@ def block(self, dims, kind='default'):
         return Properties(m)
 
     def inbound(self, dims):
-        m = dict(self)
-        for d in as_tuple(dims):
-            m[d] = set(m.get(d, [])) | {INBOUND}
-        return Properties(m)
+        return self.add(dims, INBOUND)
+
+    def init_core_shm(self, dims):
+        return self.add(dims, INIT_CORE_SHM)
+
+    def init_halo_left_shm(self, dims):
+        return self.add(dims, INIT_HALO_LEFT_SHM)
+
+    def init_halo_right_shm(self, dims):
+        return self.add(dims, INIT_HALO_RIGHT_SHM)
 
     def is_parallel(self, dims):
         return any(len(self[d] & {PARALLEL, PARALLEL_INDEP}) > 0
@@ -299,13 +329,28 @@ def is_blockable(self, d):
     def is_blockable_small(self, d):
         return TILABLE_SMALL in self.get(d, set())
 
-    def is_prefetchable(self, dims=None, v=PREFETCHABLE):
+    def _is_property_any(self, dims, v):
         if dims is None:
             dims = list(self)
         return any(v in self.get(d, set()) for d in as_tuple(dims))
 
+    def is_prefetchable(self, dims=None, v=PREFETCHABLE):
+        return self._is_property_any(dims, PREFETCHABLE)
+
     def is_prefetchable_shm(self, dims=None):
-        return self.is_prefetchable(dims, PREFETCHABLE_SHM)
+        return self._is_property_any(dims, PREFETCHABLE_SHM)
+
+    def is_core_init(self, dims=None):
+        return self._is_property_any(dims, INIT_CORE_SHM)
+
+    def is_halo_left_init(self, dims=None):
+        return self._is_property_any(dims, INIT_HALO_LEFT_SHM)
+
+    def is_halo_right_init(self, dims=None):
+        return self._is_property_any(dims, INIT_HALO_RIGHT_SHM)
+
+    def is_halo_init(self, dims=None):
+        return self.is_halo_left_init(dims) or self.is_halo_right_init(dims)
 
     @property
     def nblockable(self):

devito/ir/support/space.py

@@ -1056,6 +1056,10 @@ def outermost(self):
     def innermost(self):
         return self[-1]
 
+    @cached_property
+    def concrete(self):
+        return self.project(lambda d: not d.is_Virtual)
+
     @cached_property
     def itintervals(self):
         return tuple(self[d] for d in self.itdims)

devito/operator/operator.py

@@ -703,6 +703,15 @@ def _postprocess_errors(self, retval):
                 "due to excessive register pressure in one of the Operator "
                 "kernels. Try supplying a smaller `par-tile` value."
             )
+        elif retval == error_mapper['KernelLaunchClusterConfig']:
+            raise ExecutionError(
+                "Kernel launch failed due to an invalid thread block cluster "
+                "configuration. This is probably due to a `tbc-tile` value that "
+                "does not perfectly divide the number of blocks launched for a "
+                "kernel. This is a known, strong limitation which effectively "
+                "prevents the use of `tbc-tile` in realistic scenarios, but it "
+                "will be removed in future versions."
+            )
         elif retval == error_mapper['KernelLaunchUnknown']:
             raise ExecutionError(
                 "Kernel launch failed due to an unknown error. This might "

devito/passes/clusters/cse.py

@@ -70,6 +70,9 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     min_cost = options['cse-min-cost']
     mode = options['cse-algo']
 
+    if cluster.is_fence:
+        return cluster
+
     make = lambda: CTemp(name=sregistry.make_name(), dtype=cluster.dtype)
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)

devito/passes/clusters/misc.py

@@ -228,11 +228,12 @@ def _key(self, c):
         # If there are writes to thread-shared object, make it part of the key.
         # This will promote fusion of non-adjacent Clusters writing to (some
         # form of) shared memory, which in turn will minimize the number of
-        # necessary barriers Same story for reads from thread-shared objects
+        # necessary barriers. Same story for reads from thread-shared objects
         weak.extend([
             any(f._mem_shared for f in c.scope.writes),
             any(f._mem_shared for f in c.scope.reads)
         ])
+        weak.append(c.properties.is_core_init())
 
         # Prefetchable Clusters should get merged, if possible
         weak.append(c.properties.is_prefetchable_shm())

devito/passes/iet/errors.py

@@ -109,5 +109,6 @@ class Retval(LocalObject, Expr):
     'Stability': 100,
     'KernelLaunch': 200,
     'KernelLaunchOutOfResources': 201,
-    'KernelLaunchUnknown': 202,
+    'KernelLaunchClusterConfig': 202,
+    'KernelLaunchUnknown': 203,
 }