compiler: Add instrumentation for compute0 core area

devitocodes · Jun 27, 2023 · 026d354 · 026d354
1 parent 93059c1
commit 026d354
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py
@@ -16,7 +16,7 @@
 from devito.ir.support import IntervalGroup
 from devito.logger import warning, error
 from devito.mpi import MPI
-from devito.mpi.routines import MPICall, MPIList, RemainderCall
+from devito.mpi.routines import MPICall, MPIList, RemainderCall, ComputeCall
 from devito.parameters import configuration
 from devito.symbolics import subs_op_args
 from devito.tools import DefaultOrderedDict, flatten
@@ -332,7 +332,7 @@ class AdvancedProfilerVerbose2(AdvancedProfilerVerbose):
 
     @property
     def trackable_subsections(self):
-        return (MPICall, BusyWait)
+        return (MPICall, BusyWait, ComputeCall)
 
 
 class AdvisorProfiler(AdvancedProfiler):

diff --git a/devito/passes/iet/instrument.py b/devito/passes/iet/instrument.py
@@ -1,7 +1,8 @@
 from devito.ir.iet import (BusyWait, FindNodes, FindSymbols, MapNodes, Section,
                            TimedList, Transformer)
 from devito.mpi.routines import (HaloUpdateCall, HaloWaitCall, MPICall, MPIList,
-                                 HaloUpdateList, HaloWaitList, RemainderCall)
+                                 HaloUpdateList, HaloWaitList, RemainderCall,
+                                 ComputeCall)
 from devito.passes.iet.engine import iet_pass
 from devito.types import Timer
 
@@ -36,14 +37,15 @@ def track_subsections(iet, **kwargs):
         HaloUpdateCall: 'haloupdate',
         HaloWaitCall: 'halowait',
         RemainderCall: 'remainder',
+        ComputeCall: 'compute',
         HaloUpdateList: 'haloupdate',
         HaloWaitList: 'halowait',
         BusyWait: 'busywait'
     }
 
     mapper = {}
 
-    for NodeType in [MPIList, MPICall, BusyWait]:
+    for NodeType in [MPIList, MPICall, BusyWait, ComputeCall]:
         for k, v in MapNodes(Section, NodeType).visit(iet).items():
             for i in v:
                 if i in mapper or not any(issubclass(i.__class__, n)

diff --git a/tests/test_mpi.py b/tests/test_mpi.py
@@ -1420,6 +1420,23 @@ def test_many_functions(self):
         assert len(calls) == 2
         assert calls[0].ncomps == 7
 
+    @switchconfig(profiling='advanced2')
+    @pytest.mark.parallel(mode=[
+        (1, 'full'),
+    ])
+    def test_profiled_regions(self):
+        grid = Grid(shape=(10, 10, 10))
+
+        f = TimeFunction(name='f', grid=grid, space_order=2)
+        g = TimeFunction(name='g', grid=grid, space_order=2)
+
+        eqns = [Eq(f.forward, f.dx2 + 1.),
+                Eq(g.forward, g.dx2 + 1.)]
+
+        op = Operator(eqns)
+        assert op._profiler.all_sections == ['section0', 'haloupdate0', 'halowait0',
+                                             'remainder0', 'compute0']
+
     @pytest.mark.parallel(mode=1)
     def test_enforce_haloupdate_if_unwritten_function(self):
         grid = Grid(shape=(16, 16))