Add PETSc event logging to generated code, reuses logic from PyOP2

Author: connorjward

pyop3/ir/lower.py

@@ -17,32 +17,15 @@
 from pyrsistent import freeze, pmap
 
 from pyop3.array import HierarchicalArray
-from pyop3.array.harray import CalledMapVariable
 from pyop3.array.petsc import AbstractMat, Sparsity
 from pyop3.axtree import Axis, AxisComponent, AxisTree, AxisVariable, ContextFree
-from pyop3.axtree.tree import subst_layouts
 from pyop3.buffer import DistributedBuffer, NullBuffer, PackedBuffer
 from pyop3.config import config
 from pyop3.dtypes import IntType
-from pyop3.ir.transform import add_likwid_markers
-from pyop3.itree import (
-    AffineSliceComponent,
-    CalledMap,
-    Index,
-    IndexTree,
-    LocalLoopIndex,
-    LoopIndex,
-    Map,
-    Slice,
-    Subset,
-    TabulatedMapComponent,
-)
+from pyop3.ir.transform import with_likwid_markers, with_petsc_event
 from pyop3.itree.tree import (
-    ContextFreeLoopIndex,
-    IndexExpressionReplacer,
     LocalLoopIndexVariable,
     LoopIndexVariable,
-    collect_shape_index_callback,
 )
 from pyop3.lang import (
     INC,
@@ -323,12 +306,12 @@ def set_temporary_shapes(self, shapes):
 
 
 class CodegenResult:
-    def __init__(self, expr, ir, arg_replace_map):
+    def __init__(self, expr, ir, arg_replace_map, *, compiler_parameters):
         self.expr = as_tuple(expr)
         self.ir = ir
         self.arg_replace_map = arg_replace_map
 
-        self._exec = compile_loopy(self.ir)
+        self._exec = compile_loopy(self.ir, pyop3_compiler_parameters=compiler_parameters)
 
     @cached_property
     def datamap(self):
@@ -401,6 +384,7 @@ class CompilerParameters:
     # NOTE: This sort of thing could have a default set from the config
     # dict (but do not use PYOP3_USE_LIKWID as that's a separate option).
     add_likwid_markers: bool = False
+    add_petsc_event: bool = False
 
 
 def parse_compiler_parameters(compiler_parameters) -> CompilerParameters:
@@ -412,16 +396,18 @@ def parse_compiler_parameters(compiler_parameters) -> CompilerParameters:
 
 
 # prefer generate_code?
-def compile(expr: Instruction, name="mykernel", compiler_parameters=None):
+def compile(expr: Instruction, compiler_parameters=None):
     compiler_parameters = parse_compiler_parameters(compiler_parameters)
 
     # preprocess expr before lowering
     from pyop3.transform import expand_implicit_pack_unpack, expand_loop_contexts
 
+    function_name = expr.name
+
     cs_expr = expand_loop_contexts(expr)
     ctx = LoopyCodegenContext()
-    for context, expr in cs_expr:
-        expr = expand_implicit_pack_unpack(expr)
+    for context, ex in cs_expr:
+        ex = expand_implicit_pack_unpack(ex)
 
         # add external loop indices as kernel arguments
         loop_indices = {}
@@ -441,7 +427,7 @@ def compile(expr: Instruction, name="mykernel", compiler_parameters=None):
             # FIXME currently assume that source and target exprs are the same, they are not!
             loop_indices[index] = (replace_map, replace_map)
 
-        for e in as_tuple(expr):
+        for e in as_tuple(ex):
             # context manager?
             ctx.set_temporary_shapes(_collect_temporary_shapes(e))
             _compile(e, loop_indices, ctx)
@@ -480,24 +466,29 @@ def compile(expr: Instruction, name="mykernel", compiler_parameters=None):
         ctx.domains,
         ctx.instructions,
         ctx.arguments,
-        name=name,
+        name=function_name,
         target=LOOPY_TARGET,
         lang_version=LOOPY_LANG_VERSION,
         preambles=preambles,
         # options=lp.Options(check_dep_resolution=False),
     )
 
+    entrypoint = translation_unit.default_entrypoint
     if compiler_parameters.add_likwid_markers:
-        translation_unit = add_likwid_markers(translation_unit)
+        entrypoint = with_likwid_markers(entrypoint)
+    if compiler_parameters.add_petsc_event:
+        entrypoint = with_petsc_event(entrypoint)
+    translation_unit = translation_unit.with_kernel(entrypoint)
 
-    tu = lp.merge((translation_unit, *ctx.subkernels))
+    translation_unit = lp.merge((translation_unit, *ctx.subkernels))
 
     # add callables
     # tu = lp.register_callable(tu, "bsearch", BinarySearchCallable())
 
-    tu = tu.with_entrypoints(name)
+    # needed?
+    translation_unit = translation_unit.with_entrypoints(entrypoint.name)
 
-    return CodegenResult(expr, tu, ctx.kernel_to_actual_rename_map)
+    return CodegenResult(expr, translation_unit, ctx.kernel_to_actual_rename_map, compiler_parameters=compiler_parameters)
 
 
 # put into a class in transform.py?

pyop3/ir/transform.py

@@ -1,23 +1,55 @@
+import textwrap
+
 import loopy as lp
 
 
-def add_likwid_markers(knl):
+def with_likwid_markers(knl):
     """
     See https://github.com/RRZE-HPC/likwid/wiki/TutorialMarkerC
     """
     import pylikwid
 
-    preambles = knl.preambles + (("99_likwid", "#include <likwid-marker.h>"),)
-
     marker_name = knl.name
     pylikwid.markerregisterregion(marker_name)
 
+    preambles = [("99_likwid", "#include <likwid-marker.h>")]
     start_insn = lp.CInstruction((), f"LIKWID_MARKER_START(\"{marker_name}\");", id="likwid_start")
+    stop_insn = lp.CInstruction((), f"LIKWID_MARKER_STOP(\"{marker_name}\");", id="likwid_stop")
+
+    return _with_region_markers(knl, start_insn, stop_insn, preambles)
+
+
+def with_petsc_event(knl):
+    event_name = knl.name
+
+
+    preambles = [
+        (
+            "99_petsc",
+            textwrap.dedent(f"""
+                #include <petsclog.h>
+
+                // Prepare a dummy event so that things compile. This is overwridden using
+                // the object file.
+                PetscLogEvent id_{event_name} = -1;
+            """)
+        )
+    ]
+
+    start_insn = lp.CInstruction((), f"PetscLogEventBegin(id_{event_name}, 0, 0, 0, 0);", id="petsc_log_begin")
+    stop_insn = lp.CInstruction((), f"PetscLogEventEnd(id_{event_name}, 0, 0, 0, 0);", id="petsc_log_end")
+
+    return _with_region_markers(knl, start_insn, stop_insn, preambles)
+
+
+def _with_region_markers(knl, start_insn, stop_insn, preambles):
+    preambles = knl.preambles + tuple(preambles)
+
+    assert start_insn.id is not None
     insns = (
-        [start_insn]
-        + [insn.copy(depends_on=insn.depends_on | {"likwid_start"}) for insn in knl.instructions]
+        start_insn,
+        *(insn.copy(depends_on=insn.depends_on | {start_insn.id}) for insn in knl.instructions),
+        stop_insn.copy(depends_on=frozenset(insn.id for insn in knl.instructions)),
     )
-    stop_insn = lp.CInstruction((), f"LIKWID_MARKER_STOP(\"{marker_name}\");", id="likwid_stop", depends_on=frozenset(insn.id for insn in insns))
-    insns.append(stop_insn)
 
     return knl.copy(preambles=preambles, instructions=insns)

pyop3/lang.py

@@ -116,8 +116,11 @@ def datamap(self):
     #     pass
 
 
+_DEFAULT_LOOP_NAME = "pyop3_loop"
+
+
 class Loop(Instruction):
-    fields = Instruction.fields | {"index", "statements", "compiler_parameters"}
+    fields = Instruction.fields | {"index", "statements", "compiler_parameters", "name"}
 
     # doubt that I need an ID here
     id_generator = pytools.UniqueNameGenerator()
@@ -127,33 +130,35 @@ def __init__(
         index: LoopIndex,
         statements: Iterable[Instruction],
         *,
+        name: str = _DEFAULT_LOOP_NAME,
         compiler_parameters=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.index = index
         self.statements = as_tuple(statements)
+        self.name = name
         self.compiler_parameters = compiler_parameters
 
     def __call__(self, **kwargs):
         # TODO just parse into ContextAwareLoop and call that
         from pyop3.ir.lower import compile
         from pyop3.itree.tree import partition_iterset
 
+        code = compile(self, compiler_parameters=self.compiler_parameters)
+
         if self.is_parallel:
             # FIXME: The partitioning code does not seem to always run properly
             # so for now do all the transfers in advance.
             # interleave computation and communication
-            # new_index, (icore, iroot, ileaf) = partition_iterset(
-            #     self.index, [a for a, _ in self.function_arguments]
-            # )
+            new_index, (icore, iroot, ileaf) = partition_iterset(
+                self.index, [a for a, _ in self.function_arguments]
+            )
             #
             # assert self.index.id == new_index.id
             #
             # # substitute subsets into loopexpr, should maybe be done in partition_iterset
             # parallel_loop = self.copy(index=new_index)
-            # code = compile(parallel_loop)
-            code = compile(self, compiler_parameters=self.compiler_parameters)
 
             # interleave communication and computation
             initializers, reductions, broadcasts = self._array_updates()
@@ -162,40 +167,43 @@ def __call__(self, **kwargs):
                 init()
 
             # replace the parallel axis subset with one for the specific indices here
-            # extent = just_one(icore.axes.root.components).count
-            # core_kwargs = merge_dicts(
-            #     [kwargs, {icore.name: icore, extent.name: extent}]
-            # )
-            # code(**core_kwargs)
+            extent = just_one(icore.axes.root.components).count
+            core_kwargs = merge_dicts(
+                [kwargs, {icore.name: icore, extent.name: extent}]
+            )
+
+            with PETSc.Log.Event(f"compute_{self.name}_core"):
+                code(**core_kwargs)
 
             # await reductions
             for red in reductions:
                 red()
 
             # roots
             # replace the parallel axis subset with one for the specific indices here
-            # root_extent = just_one(iroot.axes.root.components).count
-            # root_kwargs = merge_dicts(
-            #     [kwargs, {icore.name: iroot, extent.name: root_extent}]
-            # )
-            # code(**root_kwargs)
+            root_extent = just_one(iroot.axes.root.components).count
+            root_kwargs = merge_dicts(
+                [kwargs, {icore.name: iroot, extent.name: root_extent}]
+            )
+            with PETSc.Log.Event(f"compute_{self.name}_root"):
+                code(**root_kwargs)
 
             # await broadcasts
             for broadcast in broadcasts:
                 broadcast()
 
             # leaves
-            # leaf_extent = just_one(ileaf.axes.root.components).count
-            # leaf_kwargs = merge_dicts(
-            #     [kwargs, {icore.name: ileaf, extent.name: leaf_extent}]
-            # )
-            # code(**leaf_kwargs)
-
-            code(**kwargs)
+            leaf_extent = just_one(ileaf.axes.root.components).count
+            leaf_kwargs = merge_dicts(
+                [kwargs, {icore.name: ileaf, extent.name: leaf_extent}]
+            )
+            with PETSc.Log.Event(f"compute_{self.name}_leaf"):
+                code(**leaf_kwargs)
 
             # also may need to eagerly assemble Mats, or be clever and spike the accessors?
         else:
-            compile(self, compiler_parameters=self.compiler_parameters)(**kwargs)
+            with PETSc.Log.Event(f"compute_{self.name}_serial"):
+                code(**kwargs)
 
     @cached_property
     def loopy_code(self):
@@ -345,6 +353,10 @@ def _init_nil():
 
         return tuple(initializers), tuple(reductions), tuple(broadcasts)
 
+    @cached_property
+    def datamap(self):
+        return self.index.datamap | merge_dicts(stmt.datamap for stmt in self.statements)
+
 
 class ContextAwareLoop(ContextAwareInstruction):
     fields = Instruction.fields | {"index", "statements"}

pyop3/target.py

@@ -54,7 +54,7 @@
 
 
 @mpi.collective
-def compile_loopy(kernel, **kwargs):
+def compile_loopy(translation_unit, *, pyop3_compiler_parameters, **kwargs):
     """Build a shared library and return a function pointer from it.
 
     :arg jitmodule: The JIT Module which can generate the code to compile, or
@@ -71,8 +71,8 @@ def compile_loopy(kernel, **kwargs):
     :kwarg comm: Optional communicator to compile the code on (only
         rank 0 compiles code) (defaults to pyop2.mpi.COMM_WORLD).
     """
-    code = lp.generate_code_v2(kernel).device_code()
-    argtypes = [ctypes.c_voidp for _ in kernel.default_entrypoint.args]
+    code = lp.generate_code_v2(translation_unit).device_code()
+    argtypes = [ctypes.c_voidp for _ in translation_unit.default_entrypoint.args]
     restype = None
 
     # ideally move this logic somewhere else
@@ -88,21 +88,30 @@ def compile_loopy(kernel, **kwargs):
         # + tuple(self.local_kernel.ldargs)
     )
 
+    # NOTE: no - instead of this inspect the compiler parameters!!!
     # TODO: Make some sort of function in config.py
     if "LIKWID_MODE" in os.environ:
         cppargs += ("-DLIKWID_PERFMON",)
         ldargs += ("-llikwid",)
 
-    return compile_c(
+    func, lib = compile_c(
         code,
-        kernel.default_entrypoint.name,
+        translation_unit.default_entrypoint.name,
         argtypes,
         restype,
         extra_compiler_flags=cppargs,
         extra_linker_flags=ldargs,
         **kwargs,
     )
 
+    if pyop3_compiler_parameters.add_petsc_event:
+        # Create the event in python and then set in the shared library to avoid
+        # allocating memory over and over again in the C kernel.
+        event_name = translation_unit.default_entrypoint.name
+        ctypes.c_int.in_dll(lib, f"id_{event_name}").value = PETSc.Log.Event(event_name).id
+
+    return func
+
 
 def compile_c(code: str, name, argtypes, restype, **kwargs):
     compiler = _compiler if _compiler else sniff_compiler(config["cc"])
@@ -511,7 +520,7 @@ def compile_library(self, code: str, name: str, argtypes, restype):
         fn = getattr(dll, name)
         fn.argtypes = argtypes
         fn.restype = restype
-        return fn
+        return fn, dll
 
 
 class MacClangCompiler(Compiler):