refactor to simplify API

Author: matthiasdiener

arraycontext/impl/pytato/__init__.py

@@ -245,14 +245,6 @@ class ProfileEvent:
     t_unit_name: str
 
 
-@dataclass
-class MultiCallKernelProfile:
-    """Class to hold the results of multiple kernel executions."""
-
-    num_calls: int
-    time: int
-
-
 class PytatoPyOpenCLArrayContext(_BasePytatoArrayContext):
     """
     An :class:`ArrayContext` that uses :mod:`pytato` data types to represent
@@ -289,14 +281,13 @@ def __init__(
             representation. This interface should be considered
             unstable.
         """
-        import pyopencl as cl
-
         if allocator is not None and use_memory_pool is not None:
             raise TypeError("may not specify both allocator and use_memory_pool")
 
         self.profile_kernels = profile_kernels
 
         if profile_kernels:
+            import pyopencl as cl
             if not queue.properties & cl.command_queue_properties.PROFILING_ENABLE:
                 raise RuntimeError("Profiling was not enabled in the command queue. "
                  "Please create the queue with "
@@ -379,44 +370,12 @@ def _wait_and_transfer_profile_events(self) -> None:
     def _add_profiling_events(self, start: cl._cl.Event, stop: cl._cl.Event,
                              t_unit_name: str) -> None:
         """Add profiling events to the list of profiling events."""
-        if self.profile_kernels:
-            self._profile_events.append(ProfileEvent(start, stop, t_unit_name))
-
-    def get_profiling_data_for_kernel(self, kernel_name: str) \
-          -> MultiCallKernelProfile:
-        """Return profiling data for kernel *kernel_name*."""
-        self._wait_and_transfer_profile_events()
-
-        time = 0
-        num_calls = 0
-
-        if kernel_name in self.profile_results:
-            knl_results = self.profile_results[kernel_name]
-
-            num_calls = len(knl_results)
-
-            for r in knl_results:
-                time += r
-
-        return MultiCallKernelProfile(num_calls, time)
-
-    def reset_profiling_data_for_kernel(self, kernel_name: str) -> None:
-        """Reset profiling data for kernel *kernel_name*."""
-        self.profile_results.pop(kernel_name, None)
-
-    def get_and_reset_profiling_data(self) -> dict[str, MultiCallKernelProfile]:
-        """Return and reset profiling data."""
-        self._wait_and_transfer_profile_events()
-
-        result = {
-            kernel_name: MultiCallKernelProfile(len(times), sum(times))
-            for kernel_name, times in self._profile_results.items()
-        }
+        self._profile_events.append(ProfileEvent(start, stop, t_unit_name))
 
+    def _reset_profiling_data(self) -> None:
+        """Reset profiling data."""
         self._profile_results = {}
 
-        return result
-
     # }}}
 
     @property

arraycontext/impl/pytato/utils.py

@@ -4,6 +4,12 @@
 __doc__ = """
 .. autofunction:: transfer_from_numpy
 .. autofunction:: transfer_to_numpy
+
+
+Profiling-related functions:
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. autofunction:: tabulate_profiling_data
 """
 
 
@@ -52,7 +58,7 @@
 
 from arraycontext import ArrayContext
 from arraycontext.impl.pyopencl.taggable_cl_array import Axis as ClAxis
-from arraycontext.impl.pytato import MultiCallKernelProfile
+from arraycontext.impl.pytato import PytatoPyOpenCLArrayContext
 
 
 if TYPE_CHECKING:
@@ -224,9 +230,12 @@ def transfer_to_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
 # }}}
 
 
-def tabulate_profiling_data(profile_results: dict[str, MultiCallKernelProfile])\
-       -> pytools.Table:
+# {{{ Profiling
+
+def tabulate_profiling_data(actx: PytatoPyOpenCLArrayContext) -> pytools.Table:
     """Return a :class:`pytools.Table` with the profiling results."""
+    actx._wait_and_transfer_profile_events()
+
     tbl = pytools.Table()
 
     # Table header
@@ -238,23 +247,27 @@ def tabulate_profiling_data(profile_results: dict[str, MultiCallKernelProfile])\
     total_calls = 0
     total_time = 0.0
 
-    for kernel_name, mckp in profile_results.items():
-        total_calls += mckp.num_calls
+    for kernel_name, times in actx._profile_results.items():
+        num_calls = len(times)
+        total_calls += num_calls
 
-        t_sum = mckp.time
-        t_avg = mckp.time / mckp.num_calls
+        t_sum = sum(times)
+        t_avg = t_sum / num_calls
         if t_sum is not None:
             total_time += t_sum
 
         time_sum = f"{t_sum:{g}}"
         time_avg = f"{t_avg:{g}}"
 
-        tbl.add_row((kernel_name, mckp.num_calls, time_sum,
-                time_avg,))
+        tbl.add_row((kernel_name, num_calls, time_sum, time_avg))
 
     tbl.add_row(("", "", "", ""))
     tbl.add_row(("Total", total_calls, f"{total_time:{g}}", "--"))
 
+    actx._reset_profiling_data()
+
     return tbl
 
+# }}}
+
 # vim: foldmethod=marker

test/test_pytato_arraycontext.py

@@ -303,7 +303,7 @@ def twice(x):
 
     from arraycontext.impl.pytato.utils import tabulate_profiling_data
 
-    print(tabulate_profiling_data(actx.get_and_reset_profiling_data()))
+    print(tabulate_profiling_data(actx))
     assert len(actx._profile_results) == 0
 
     # }}}
@@ -321,7 +321,7 @@ def twice(x):
     assert len(actx._profile_results) == 1
     assert len(actx._profile_results["frozen_result"]) == 10
 
-    print(tabulate_profiling_data(actx.get_and_reset_profiling_data()))
+    print(tabulate_profiling_data(actx))
 
     assert len(actx._profile_results) == 0