Add event decorators

pyop2/caching.py

@@ -48,6 +48,7 @@
 from pyop2.mpi import (
     MPI, COMM_WORLD, comm_cache_keyval, temp_internal_comm
 )
+from petsc4py import PETSc
 
 
 # Caches created here are registered as a tuple of
@@ -403,6 +404,7 @@ def parallel_cache(
         subcommunicators.
     """
     def decorator(func):
+        @PETSc.Log.EventDecorator("PyOP2 Cache Wrapper")
         @wraps(func)
         def wrapper(*args, **kwargs):
             """ Extract the key and then try the memory cache before falling back
@@ -430,6 +432,10 @@ def wrapper(*args, **kwargs):
                 if (comm, comm.name, func, local_cache) not in [k[1:] for k in _KNOWN_CACHES]:
                     _KNOWN_CACHES.append((next(_CACHE_CIDX), comm, comm.name, func, local_cache))
 
+                # JBTODO: Replace everything below here with:
+                # value = local_cache.get(key, CACHE_MISS)
+                # and add an optional PYOP2_SPMD_STRICT environment variable
+
                 if broadcast:
                     # Grab value from rank 0 memory cache and broadcast result
                     if comm.rank == 0:

pyop2/compilation.py

@@ -425,6 +425,7 @@ def load_hashkey(*args, **kwargs):
 # ...benchmarking disagrees with my assessment
 @mpi.collective
 @memory_cache(hashkey=load_hashkey, broadcast=False)
+@PETSc.Log.EventDecorator()
 def load(jitmodule, extension, fn_name, cppargs=(), ldargs=(),
          argtypes=None, restype=None, comm=None):
     """Build a shared library and return a function pointer from it.
@@ -556,6 +557,7 @@ def check_source_hashes(compiler, jitmodule, extension, comm):
     hashkey=_make_so_hashkey,
     cache_factory=lambda: CompilerDiskAccess(configuration['cache_dir'], extension=".so")
 )
+@PETSc.Log.EventDecorator()
 def make_so(compiler, jitmodule, extension, comm, filename=None):
     """Build a shared library and load it