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
@@ -752,8 +753,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

@@ -31,6 +31,7 @@
     'Block',
     'Break',
     'BusyWait',
+    'BusyWaitCall',
     'Call',
     'Callable',
     'CallableBody',
@@ -1534,6 +1535,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,9 +5,10 @@
 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, ThreadFence, Transformer, While,
+    make_callable, maybe_alias
 )
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.engine import iet_pass
@@ -70,12 +71,12 @@ 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)]
+            activation = [BusyWait(condition)]
         else:
             d = Temp(name=sregistry.make_name(prefix=sdata.index.name))
             condition = CondNe(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 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:
@@ -98,7 +99,7 @@ def _lower_async_objs(iet, tracker=None, sregistry=None, **kwargs):
         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)
 
@@ -222,7 +223,7 @@ 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)),
+        BusyWait(CondEq(FieldFromPointer(sdata.symbolic_flag, sbase), 2)),
         DummyExpr(FieldFromPointer(sdata.symbolic_flag, sbase), 0),
         Call('pthread_join', (threads[d], Null))
     ]))
@@ -252,7 +253,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(

devito/passes/iet/instrument.py

@@ -1,7 +1,8 @@
 from itertools import groupby
 
 from devito.ir.iet import (
-    BusyWait, FindNodes, FindSymbols, Iteration, MapNodes, Section, TimedList, Transformer
+    BusyWait, BusyWaitCall, EntryFunction, FindNodes, FindSymbols, Iteration, MapNodes,
+    Section, TimedList, Transformer
 )
 from devito.mpi.routines import (
     ComputeCall, HaloUpdateCall, HaloUpdateList, HaloWaitCall, HaloWaitList, MPICall,
@@ -30,10 +31,10 @@ def instrument(graph, **kwargs):
 @iet_pass
 def track_subsections(iet, **kwargs):
     """
-    Add sub-Sections to the `profiler`. Sub-Sections include:
+    Add synchronization Sections and sub-Sections to the `profiler`. These include:
 
         * MPI Calls (e.g., HaloUpdateCall and HaloUpdateWait)
-        * Busy-waiting on While(lock) (e.g., from host-device orchestration)
+        * Busy-waiting loops and calls (e.g., from host-device orchestration)
         * Multi-pass implementations -- one sub-Section for each pass, within one
           macro Section
     """
@@ -47,28 +48,35 @@ def track_subsections(iet, **kwargs):
         ComputeCall: 'compute',
         HaloUpdateList: 'haloupdate',
         HaloWaitList: 'halowait',
-        BusyWait: 'busywait'
+        BusyWait: 'busywait',
+        BusyWaitCall: 'busywait'
     }
 
     verbosity_mapper = {
         0: (),
-        1: (MPIList, RemainderCall, BusyWait),
-        2: (MPICall, ComputeCall, BusyWait),
+        1: (MPIList, RemainderCall, BusyWait, BusyWaitCall),
+        2: (MPICall, ComputeCall, BusyWait, BusyWaitCall),
     }
 
     mapper = {}
 
     # Enable/disable profiling of sub-Sections
     for NodeType in verbosity_mapper[profiler._verbosity]:
-        for k, v in MapNodes(Section, NodeType).visit(iet).items():
+        for k, v in MapNodes(Section, NodeType, mode='immediate').visit(iet).items():
+            if k is None and not isinstance(iet, EntryFunction):
+                # Unsectioned wait helpers are timed at their `BusyWaitCall` sites.
+                continue
             for i in v:
                 if i in mapper:
                     continue
                 name = sregistry.make_name(prefix=name_mapper[i.__class__])
-                mapper[i] = Section(name, body=i, is_subsection=True)
-                profiler.track_subsection(k.name, name)
+                mapper[i] = Section(name, body=i, is_subsection=k is not None)
+                if k is not None:
+                    profiler.track_subsection(k.name, name)
 
     iet = Transformer(mapper).visit(iet)
+    if any(not i.is_subsection for i in mapper.values()):
+        profiler.analyze(iet)
 
     # Multi-pass implementations
     mapper = {}

devito/passes/iet/langbase.py

@@ -574,7 +574,9 @@ def make_sections_from_imask(f, imask=None):
         # that the symbol `start` is actually a temporary whose value already
         # represents the unrolled size
         if not isinstance(start, Wildcard):
-            start *= remainder_size
+            # Slot offsets follow the allocation, not the transferred payload size
+            remainder_stride = prod(f.symbolic_shape[len(imask):])
+            start *= remainder_stride
         size *= remainder_size
         sections.append((start, size))
 

devito/passes/iet/orchestration.py

@@ -6,8 +6,9 @@
 
 from devito.exceptions import CompilationError
 from devito.ir.iet import (
-    AsyncCall, AsyncCallable, BlankLine, Block, BusyWait, Call, Callable, Conditional,
-    DummyExpr, List, SyncSpot, ThreadFence, Transformer, derive_parameters, make_callable
+    AsyncCall, AsyncCallable, BlankLine, Block, BusyWait, BusyWaitCall, Call, Callable,
+    Conditional, DummyExpr, List, SyncSpot, ThreadFence, Transformer, derive_parameters,
+    make_callable
 )
 from devito.ir.iet.visitors import Visitor
 from devito.ir.support import (
@@ -66,7 +67,7 @@ def _make_releaselock(self, iet, sync_ops, *args):
         if isinstance(iet, SyncSpot) and not iet.body:
             iet = List()
 
-        iet = List(body=[Call(name, efunc.parameters)] + [iet])
+        iet = List(body=[BusyWaitCall(name, efunc.parameters)] + [iet])
 
         return iet, [efunc]
 

tests/test_arch.py

@@ -1,7 +1,11 @@
 import pytest
+from packaging.version import Version
 
 from devito import configuration, switchconfig
-from devito.arch.compiler import GNUCompiler, compiler_registry, sniff_compiler_version
+from devito.arch import BLACKWELL, NVIDIAX
+from devito.arch.compiler import (
+    CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version
+)
 
 
 @pytest.mark.parametrize("cc", [
@@ -29,3 +33,26 @@ def test_switcharch():
     assert isinstance(tmp_comp, old_compiler.__class__)
     assert old_compiler.suffix == tmp_comp.suffix
     assert old_compiler.name == tmp_comp.name
+
+
+@pytest.mark.parametrize('platform,cc,target', [
+    (NVIDIAX, None, '-arch=native'),
+    (NVIDIAX, 80, '-arch=sm_80'),
+    (NVIDIAX, 90, '-arch=sm_90'),
+    (NVIDIAX, 100, '-gencode=arch=compute_100a,code=sm_100a'),
+    (NVIDIAX, 103, '-arch=sm_103'),
+    (NVIDIAX, 120, '-arch=sm_120'),
+    (BLACKWELL, None, '-gencode=arch=compute_100a,code=sm_100a'),
+    (BLACKWELL, 100, '-gencode=arch=compute_100a,code=sm_100a'),
+    (BLACKWELL, 103, '-arch=sm_103'),
+    (BLACKWELL, 120, '-arch=sm_120')
+])
+def test_cuda_target(monkeypatch, platform, cc, target):
+    with switchconfig(mpi=False), monkeypatch.context() as m:
+        m.setattr('devito.arch.compiler.get_nvidia_cc', lambda: cc)
+        m.setattr('devito.arch.compiler.get_cuda_version', lambda: Version('13.0'))
+        m.setattr('devito.arch.compiler.check_cuda_runtime', lambda: None)
+        compiler = CudaCompiler(platform=platform)
+
+    flags = [i for i in compiler.cflags if i.startswith(('-arch=', '-gencode='))]
+    assert flags == [target]