compiler: Fix performance profiling of busy-waits

devito/arch/compiler.py

@@ -15,8 +15,9 @@
 from packaging.version import Version
 
 from devito.arch import (
-    AMDGPUX, POWER8, POWER9, AppleArm, Cortex, Cpu64, Graviton, IntelDevice, NvidiaArm,
-    NvidiaDevice, check_cuda_runtime, get_cuda_version, get_m1_llvm_path, get_nvidia_cc
+    AMDGPUX, BLACKWELL, POWER8, POWER9, AppleArm, Cortex, Cpu64, Graviton, IntelDevice,
+    NvidiaArm, NvidiaDevice, check_cuda_runtime, get_cuda_version, get_m1_llvm_path,
+    get_nvidia_cc
 )
 from devito.exceptions import CompilationError
 from devito.logger import debug, warning
@@ -763,8 +764,16 @@ def __init_finalize__(self, **kwargs):
                     proc_link_flags.append(i)
             self.ldflags.extend(proc_link_flags)
 
+        platform = kwargs.pop('platform', configuration['platform'])
         cc = get_nvidia_cc()
-        if cc:
+        if cc is None and platform is BLACKWELL:
+            cc = 100
+
+        if cc == 100:
+            # `-arch=sm_100a` also validates generic compute_100 PTX, which
+            # rejects architecture-specific instructions such as tcgen05
+            self.cflags.append('-gencode=arch=compute_100a,code=sm_100a')
+        elif cc:
             self.cflags.append(f'-arch=sm_{cc}')
         else:
             self.cflags.append('-arch=native')

devito/ir/iet/nodes.py

@@ -18,7 +18,7 @@
     AFFINE, INBOUND, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, SEQUENTIAL,
     VECTORIZED, Forward, PrefetchUpdate, Property, WithLock
 )
-from devito.symbolics import CallFromPointer, ListInitializer, Macro
+from devito.symbolics import CallFromPointer, ListInitializer
 from devito.tools import (
     Signer, as_tuple, ctypes_to_cstr, filter_ordered, filter_sorted, flatten
 )
@@ -31,6 +31,7 @@
     'Block',
     'Break',
     'BusyWait',
+    'BusyWaitCall',
     'Call',
     'Callable',
     'CallableBody',
@@ -61,7 +62,6 @@
     'Section',
     'Switch',
     'SyncSpot',
-    'ThreadFence',
     'TimedList',
     'Transfer',
     'Using',
@@ -1303,31 +1303,6 @@ def size(self):
         return self.ispace.size
 
 
-class ThreadFence(Call):
-
-    """
-    A memory fence, ordering a thread's accesses either side of it.
-
-    Threads exchanging work through a shared flag need memory ordering on both
-    sides: a release before publishing the flag, and an acquire after observing it.
-    This ensures the state associated with the flag is visible before it is used.
-    Without this ordering, updates may be observed in the wrong order and lost.
-
-    `__atomic_thread_fence` is a compiler builtin, so this renders the same in C
-    and in C++ and needs no header.
-    """
-
-    def __init__(self, order):
-        assert order in ('acquire', 'release')
-        super().__init__('__atomic_thread_fence',
-                         [Macro(f'__ATOMIC_{order.upper()}')])
-        self._order = order
-
-    @property
-    def order(self):
-        return self._order
-
-
 class Prodder(Call):
 
     """
@@ -1534,6 +1509,17 @@ class BusyWait(While):
     pass
 
 
+class BusyWaitCall(Call):
+
+    """
+    A call to a helper implementing busy waiting and synchronization.
+
+    Its wait is profiled at the call site, rather than within the helper.
+    """
+
+    pass
+
+
 class SyncSpot(List):
 
     """

devito/ir/support/syncs.py

@@ -141,10 +141,8 @@ def imask(self):
         ret = []
         for d in self.target.dimensions:
             if d.root is self.dim.root:
-                if self.target.is_regular:
-                    ret.append((self.tindex, 1))
-                else:
-                    ret.append((0, 1))
+                tindex = 0 if self.target.shape_allocated[d] == 1 else self.tindex
+                ret.append((tindex, 1))
             else:
                 ret.append(FULL)
 

devito/operator/profiling.py

@@ -86,27 +86,27 @@ def analyze(self, iet):
         """
         sections = FindNodes(Section).visit(iet)
         for s in sections:
-            if s.name in self._sections:
+            if s.is_subsection or s.name in self._sections:
                 continue
 
             bundles = FindNodes(ExpressionBundle).visit(s)
 
             # Total operation count
-            ops = sum(i.ops*i.ispace.size for i in bundles)
+            ops = sum((i.ops*i.ispace.size for i in bundles), S.Zero)
 
             # Operation count at each section iteration
             # NOTE: for practical reasons, it makes much more sense to "flatten"
             # the StencilDimensions, because one expects to see the number of
             # operations per time- or space-Dimension
-            sops = sum(i.ops*max(i.ispace.project(lambda d: d.is_Stencil).size, 1)
-                       for i in bundles)
+            sops = sum((i.ops*max(i.ispace.project(lambda d: d.is_Stencil).size, 1)
+                        for i in bundles), S.Zero)
 
             # Total memory traffic
             mapper = {}
             for i in bundles:
                 for k, v in i.traffic.items():
                     mapper.setdefault(k, []).append(v)
-            traffic = 0
+            traffic = S.Zero
             for i in mapper.values():
                 try:
                     traffic += IntervalGroup.generate('union', *i).size
@@ -253,7 +253,7 @@ class AdvancedProfiler(Profiler):
 
     def _evaluate_section(self, name, data, args, dtype):
         # Time to run the section
-        time = max(getattr(args[self.name]._obj, name), 10e-7)
+        time = getattr(args[self.name]._obj, name)
 
         # Number of FLOPs performed
         try:
@@ -443,10 +443,13 @@ def add(self, name, rank, time,
         Add performance data for a given code section. With MPI enabled, the
         performance data is local, that is "per-rank".
         """
-        # Do not show unexecuted Sections (i.e., loop trip count was 0)
-        if traffic == 0:
+        # A synchronization Section may execute without arithmetic or traffic.
+        # Only discard zero-traffic Sections whose timer was never incremented.
+        if traffic == 0 and time == 0:
             return
 
+        time = max(time, 10e-7)
+
         k = PerfKey(name, rank)
 
         if not ops or any(not np.isfinite(i) for i in [ops, points, traffic]):
@@ -464,7 +467,8 @@ def add(self, name, rank, time,
 
     def add_subsection(self, sname, name, rank, time, *args):
         k0 = PerfKey(sname, rank)
-        assert k0 in self
+        if k0 not in self:
+            return
 
         self.subsections[sname][name] = PerfEntry(time, None, None, None, None, [])
 

devito/passes/clusters/buffering.py

@@ -179,7 +179,8 @@ def callback(self, clusters, prefix):
 
         # Create all the necessary ModuloDimensions to step through the buffer
         mds = make_mds(descriptors, prefix, self.sregistry)
-        subiters = as_mapper(mds.values(), lambda i: i.root)
+        subiters = as_mapper([i for i in mds.values() if isinstance(i, ModuloDimension)],
+                             lambda i: i.root)
 
         # Substitution rules to replace buffered Functions with buffers
         # E.g., `usave[time+1, x+1, y+1] -> ub0[t1, x+1, y+1]`
@@ -874,16 +875,26 @@ def first_idx(self):
 
 def make_mds(descriptors, prefix, sregistry):
     """
-    Create the ModuloDimensions to step through the buffers. This is done
-    inspecting all buffers so that ModuloDimensions are reused when possible.
+    Map buffered accesses to slot indices, reusing existing indices across
+    buffering levels and creating ModuloDimensions when necessary.
     """
     mds = defaultdict(int)
     for v in descriptors.flat_values():
-        size = v.xd.symbolic_size
+        size = v.b.shape[v.xd]
 
         if size == 1:
             continue
         if v.is_double_buffering:
+            # E.g., `f[time] = b[t0]` must become `b1[t0] = b[t0]`
+            # Initialization may use a plain index expression instead of `t0`
+            for e in v.exprs:
+                if e.lhs.function is v.f:
+                    target, buffer = e.lhs, e.rhs
+                elif e.rhs.function is v.f:
+                    target, buffer = e.rhs, e.lhs
+                else:
+                    continue
+                mds[(v.xd, target.indices[v.dim])] = buffer.indices[v.xd]
             continue
 
         p, _ = offset_from_centre(v.dim, v.indices)

devito/passes/iet/asynchrony.py

@@ -5,13 +5,16 @@
 import cgen as c
 
 from devito.ir import (
-    AsyncCall, AsyncCallable, BlankLine, Call, Callable, Conditional, DummyEq, DummyExpr,
-    EntryFunction, FindNodes, FindSymbols, Increment, Iteration, List, PointerCast,
-    Return, ThreadCallable, ThreadFence, Transformer, While, make_callable, maybe_alias
+    AsyncCall, AsyncCallable, BlankLine, BusyWait, BusyWaitCall, Call, Callable,
+    Conditional, DummyEq, DummyExpr, EntryFunction, FindNodes, FindSymbols, Increment,
+    Iteration, List, PointerCast, Return, ThreadCallable, Transformer, While,
+    make_callable, maybe_alias
 )
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.engine import iet_pass
-from devito.symbolics import CondEq, CondNe, FieldFromComposite, FieldFromPointer, Null
+from devito.symbolics import (
+    CondEq, CondNe, FieldFromComposite, FieldFromPointer, Null, atomic_load, atomic_store
+)
 from devito.tools import split
 from devito.types import (
     Lock, Pointer, PThreadArray, QueueID, SharedData, Temp, VolatileInt
@@ -69,13 +72,15 @@ def _lower_async_objs(iet, tracker=None, sregistry=None, **kwargs):
 
         if sdata.size == 1:
             d = sdata.index
-            condition = CondNe(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 1)
-            activation = [While(condition)]
+            flag = FieldFromComposite(sdata.symbolic_flag, sdata[d])
+            condition = CondNe(atomic_load(flag), 1)
+            activation = [BusyWait(condition)]
         else:
             d = Temp(name=sregistry.make_name(prefix=sdata.index.name))
-            condition = CondNe(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 1)
+            flag = FieldFromComposite(sdata.symbolic_flag, sdata[d])
+            condition = CondNe(atomic_load(flag), 1)
             activation = [DummyExpr(d, 0),
-                          While(condition, DummyExpr(d, (d + 1) % sdata.size))]
+                          BusyWait(condition, DummyExpr(d, (d + 1) % sdata.size))]
         arguments = []
         for i in sdata.ncfields:
             for a in n.arguments:
@@ -87,18 +92,13 @@ def _lower_async_objs(iet, tracker=None, sregistry=None, **kwargs):
         activation.extend([DummyExpr(FieldFromComposite(i.base, sdata[d]), i)
                            for i in arguments])
 
-        # Publishing the request must happen last. Everything the thread depends on
-        # must be visible before this flag is set. `volatile` does not provide that
-        # ordering, so weakly ordered targets could otherwise observe stale state.
-        activation.append(ThreadFence('release'))
-        activation.append(
-            DummyExpr(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 2)
-        )
+        # The request is published last, so the thread sees everything set above
+        activation.append(Call(*atomic_store(flag, 2).args))
         name = sregistry.make_name(prefix='activate')
         efunc = make_callable(name, activation)
 
         efuncs.append(efunc)
-        subs[n] = Call(name, efunc.parameters)
+        subs[n] = BusyWaitCall(name, efunc.parameters)
 
     iet = Transformer(subs).visit(iet)
 
@@ -133,6 +133,7 @@ def _(iet, key=None, tracker=None, sregistry=None, **kwargs):
         pname=sregistry.make_name(prefix='tsdata')
     )
     sbase = sdata.indexed
+    flag = FieldFromPointer(sdata.symbolic_flag, sbase)
 
     # PThreadArray -- that is he pthreads that will execute the AsyncCallable
     # asynchronously
@@ -143,31 +144,25 @@ def _(iet, key=None, tracker=None, sregistry=None, **kwargs):
     tbase = threads.indexed
 
     # Prepend the SharedData fields available upon thread activation
-    preactions = [ThreadFence('acquire')]
-    preactions.extend(DummyExpr(i, FieldFromPointer(i.base, sbase))
-                      for i in ncfields)
+    preactions = [DummyExpr(i, FieldFromPointer(i.base, sbase)) for i in ncfields]
     preactions.append(BlankLine)
 
     # Append the flag reset
     postactions = [List(body=[
         BlankLine,
         # Whatever the task produced -- data in a buffer, a lock handed back --
-        # has to be visible before the flag says it is done.
-        ThreadFence('release'),
-        DummyExpr(FieldFromPointer(sdata.symbolic_flag, sbase), 1)
+        # has to be visible before the flag says it is done
+        Call(*atomic_store(flag, 1).args)
     ])]
 
     wrap = List(body=preactions + list(iet.body.body) + postactions)
 
     # The thread has work to do when it receives the signal that all locks have
     # been set to 0 by the main thread
-    wrap = Conditional(
-        CondEq(FieldFromPointer(sdata.symbolic_flag, sbase), 2),
-        wrap
-    )
+    wrap = Conditional(CondEq(atomic_load(flag), 2), wrap)
 
     # The thread keeps spinning until the alive flag is set to 0 by the main thread
-    wrap = While(CondNe(FieldFromPointer(sdata.symbolic_flag, sbase), 0), wrap)
+    wrap = While(CondNe(atomic_load(flag), 0), wrap)
 
     # pthread functions expect exactly one argument of type void*
     tparameter = Pointer(name=f'_{sdata.name}')
@@ -207,7 +202,7 @@ def _(iet, key=None, tracker=None, sregistry=None, **kwargs):
             body.append(DummyExpr(FieldFromPointer(i.base, sbase), i.base))
     body.extend([
         BlankLine,
-        DummyExpr(FieldFromPointer(sdata.symbolic_flag, sbase), 1)
+        DummyExpr(flag, 1)
     ])
     body.extend([
         BlankLine,
@@ -222,8 +217,8 @@ def _(iet, key=None, tracker=None, sregistry=None, **kwargs):
     # Create an efunc to shutdown the pthreads
     name = sregistry.make_name(prefix='shutdown')
     body = List(body=callback([
-        While(CondEq(FieldFromPointer(sdata.symbolic_flag, sbase), 2)),
-        DummyExpr(FieldFromPointer(sdata.symbolic_flag, sbase), 0),
+        BusyWait(CondEq(atomic_load(flag), 2)),
+        Call(*atomic_store(flag, 0).args),
         Call('pthread_join', (threads[d], Null))
     ]))
     shutdown = make_callable(name, body)
@@ -252,7 +247,7 @@ def inject_async_tear_updown(iet, tracker=None, **kwargs):
 
         # Tear-down
         arguments = list(shutdown.parameters)
-        teardown.append(Call(shutdown.name, arguments))
+        teardown.append(BusyWaitCall(shutdown.name, arguments))
 
     # Inject tearup and teardown
     tearup = List(