refactor tabulate_profiling_data

Author: matthiasdiener

arraycontext/impl/pytato/__init__.py

@@ -59,7 +59,7 @@
 
 import numpy as np
 
-import pytools
+import pytools as pytools
 from pytools import memoize_method
 from pytools.tag import Tag, ToTagSetConvertible, normalize_tags
 
@@ -358,6 +358,8 @@ def __init__(
 
         self._force_svm_arg_limit = _force_svm_arg_limit
 
+    # {{{ Profiling functionality
+
     def _wait_and_transfer_profile_events(self) -> None:
         import pyopencl as cl
         # First, wait for completion of all events
@@ -403,43 +405,20 @@ 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 tabulate_profiling_data(self) -> pytools.Table:
-        """Return a :class:`pytools.Table` with the profiling results."""
+    def get_and_reset_profiling_data(self) -> dict[str, MultiCallKernelProfile]:
+        """Return and reset profiling data."""
         self._wait_and_transfer_profile_events()
 
-        tbl = pytools.Table()
-
-        # Table header
-        tbl.add_row(("Kernel", "# Calls", "Time_sum [ns]", "Time_avg [ns]"))
-
-        # Precision of results
-        g = ".4g"
-
-        total_calls = 0
-        total_time = 0.0
-
-        for knl in self.profile_results:
-            r = self.get_profiling_data_for_kernel(knl)
-
-            total_calls += r.num_calls
-
-            t_sum = r.time
-            t_avg = r.time / r.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((knl, r.num_calls, time_sum,
-                 time_avg,))
-
-        tbl.add_row(("", "", "", ""))
-        tbl.add_row(("Total", total_calls, f"{total_time:{g}}", "--"))
+        result = {
+            kernel_name: MultiCallKernelProfile(len(times), sum(times))
+            for kernel_name, times in self.profile_results.items()
+        }
 
         self.profile_results = {}
 
-        return tbl
+        return result
+
+    # }}}
 
     @property
     def _frozen_array_types(self) -> tuple[type, ...]:

arraycontext/impl/pytato/utils.py

@@ -35,6 +35,7 @@
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, cast
 
+import pytools
 from pytato.array import (
     AbstractResultWithNamedArrays,
     Array,
@@ -51,6 +52,7 @@
 
 from arraycontext import ArrayContext
 from arraycontext.impl.pyopencl.taggable_cl_array import Axis as ClAxis
+from arraycontext.impl.pytato import MultiCallKernelProfile
 
 
 if TYPE_CHECKING:
@@ -221,4 +223,38 @@ def transfer_to_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
 
 # }}}
 
+
+def tabulate_profiling_data(profile_results: dict[str, MultiCallKernelProfile])\
+       -> pytools.Table:
+    """Return a :class:`pytools.Table` with the profiling results."""
+    tbl = pytools.Table()
+
+    # Table header
+    tbl.add_row(("Kernel", "# Calls", "Time_sum [ns]", "Time_avg [ns]"))
+
+    # Precision of results
+    g = ".4g"
+
+    total_calls = 0
+    total_time = 0.0
+
+    for kernel_name, mckp in profile_results.items():
+        total_calls += mckp.num_calls
+
+        t_sum = mckp.time
+        t_avg = mckp.time / mckp.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(("", "", "", ""))
+    tbl.add_row(("Total", total_calls, f"{total_time:{g}}", "--"))
+
+    return tbl
+
 # vim: foldmethod=marker

test/test_pytato_arraycontext.py

@@ -26,6 +26,7 @@
 
 import logging
 
+import numpy as np
 import pytest
 
 from pytools.tag import Tag
@@ -289,27 +290,38 @@ def twice(x):
 
     f = actx.compile(twice)
 
+    assert len(actx.profile_events) == 0
+
     for _ in range(10):
         assert actx.to_numpy(f(99)) == 198
 
+    assert len(actx.profile_events) == 10
     actx._wait_and_transfer_profile_events()
+    assert len(actx.profile_events) == 0
     assert len(actx.profile_results) == 1
     assert len(actx.profile_results["twice"]) == 10
 
-    print(actx.tabulate_profiling_data())
+    from arraycontext.impl.pytato.utils import tabulate_profiling_data
+
+    print(tabulate_profiling_data(actx.get_and_reset_profiling_data()))
+    assert len(actx.profile_results) == 0
 
     # }}}
 
     # {{{ Uncompiled/frozen test
 
+    assert len(actx.profile_events) == 0
+
     for _ in range(10):
-        assert actx.to_numpy(twice(99)) == 198
+        assert np.all(actx.to_numpy(twice(actx.from_numpy(np.array([99, 99])))) == 198)
 
+    assert len(actx.profile_events) == 10
     actx._wait_and_transfer_profile_events()
+    assert len(actx.profile_events) == 0
     assert len(actx.profile_results) == 1
-    assert len(actx.profile_results["twice"]) == 10
+    assert len(actx.profile_results["frozen_result"]) == 10
 
-    print(actx.tabulate_profiling_data())
+    print(tabulate_profiling_data(actx.get_and_reset_profiling_data()))
 
     # }}}