diff --git a/arraycontext/impl/pyopencl/__init__.py b/arraycontext/impl/pyopencl/__init__.py index 84d5f483..36721013 100644 --- a/arraycontext/impl/pyopencl/__init__.py +++ b/arraycontext/impl/pyopencl/__init__.py @@ -56,6 +56,26 @@ # {{{ PyOpenCLArrayContext + +class PyOpenCLProfilingExecutor: + def __init__(self, actx, executor): + self.executor = executor + self.t_unit = executor.t_unit + self.actx = actx + + def __call__(self, queue, **kwargs): + if self.actx.profile_kernels: + import pyopencl as cl + start_evt = cl.enqueue_marker(queue) + + evt, result = self.executor(queue, **kwargs) + + if self.actx.profile_kernels: + self.actx._add_profiling_events(start_evt, evt, t_unit.entrypoint) + + return evt, result + + class PyOpenCLArrayContext(ArrayContext): """ A :class:`ArrayContext` that uses :class:`pyopencl.array.Array` instances @@ -246,7 +266,8 @@ def call_loopy(self, t_unit, **kwargs): executor = self._loopy_transform_cache[t_unit] except KeyError: orig_t_unit = t_unit - executor = self.transform_loopy_program(t_unit).executor(self.context) + executor = PyOpenCLProfilingExecutor(self, + self.transform_loopy_program(t_unit).executor(self.context)) self._loopy_transform_cache[orig_t_unit] = executor del orig_t_unit