From 026d35447b4cb599907f879e62e877c701e35597 Mon Sep 17 00:00:00 2001
From: George Bisbas <g.bisbas18@imperial.ac.uk>
Date: Fri, 9 Jun 2023 15:46:58 +0100
Subject: [PATCH] compiler: Add instrumentation for compute0 core area

---
 devito/operator/profiling.py    |  4 ++--
 devito/passes/iet/instrument.py |  6 ++++--
 tests/test_mpi.py               | 17 +++++++++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py
index 207e44351f5..2a58677481d 100644
--- 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
index 8ede2aa0acf..db9b446d7f2 100644
--- 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,6 +37,7 @@ def track_subsections(iet, **kwargs):
         HaloUpdateCall: 'haloupdate',
         HaloWaitCall: 'halowait',
         RemainderCall: 'remainder',
+        ComputeCall: 'compute',
         HaloUpdateList: 'haloupdate',
         HaloWaitList: 'halowait',
         BusyWait: 'busywait'
@@ -43,7 +45,7 @@ def track_subsections(iet, **kwargs):
 
     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
index 8b253f2ab8e..a6d236a1f31 100644
--- 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))