factor out profile enable/disable

Author: matthiasdiener

arraycontext/impl/pytato/__init__.py

@@ -284,22 +284,6 @@ def __init__(
         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 "
-                 "cl.command_queue_properties.PROFILING_ENABLE.")
-
-            # List of ProfileEvents that haven't been transferred to profiled
-            # results yet
-            self._profile_events: list[ProfileEvent] = []
-
-            # Dict of kernel name -> list of kernel execution times
-            self._profile_results: dict[str, list[int]] = {}
-
         self.using_svm = None
 
         if allocator is None:
@@ -348,8 +332,29 @@ def __init__(
 
         self._force_svm_arg_limit = _force_svm_arg_limit
 
+        self._enable_profiling(profile_kernels)
+
     # {{{ Profiling functionality
 
+    def _enable_profiling(self, enable: bool) -> None:
+        # List of ProfileEvents that haven't been transferred to profiled
+        # results yet
+        self._profile_events: list[ProfileEvent] = []
+
+        # Dict of kernel name -> list of kernel execution times
+        self._profile_results: dict[str, list[int]] = {}
+
+        if enable:
+            import pyopencl as cl
+            if not self.queue.properties & cl.command_queue_properties.PROFILING_ENABLE:
+                raise RuntimeError("Profiling was not enabled in the command queue. "
+                    "Please create the queue with "
+                    "cl.command_queue_properties.PROFILING_ENABLE.")
+            self.profile_kernels = True
+
+        else:
+            self.profile_kernels = False
+
     def _wait_and_transfer_profile_events(self) -> None:
         """Wait for all profiling events to finish and transfer the results
         to *self._profile_results*."""

test/test_pytato_arraycontext.py

@@ -329,7 +329,7 @@ def twice(x):
 
     # {{{ test disabling profiling
 
-    actx.profile_kernels = False
+    actx._enable_profiling(False)
 
     assert len(actx._profile_events) == 0
 
@@ -341,6 +341,32 @@ def twice(x):
 
     # }}}
 
+    # {{{ test enabling profiling
+
+    actx._enable_profiling(True)
+
+    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
+
+    # }}}
+
+    queue2 = cl.CommandQueue(cl_ctx)
+
+    with pytest.raises(RuntimeError):
+        PytatoPyOpenCLArrayContext(queue2, profile_kernels=True)
+
+    actx2 = PytatoPyOpenCLArrayContext(queue2)
+
+    with pytest.raises(RuntimeError):
+        actx2._enable_profiling(True)
+
 
 if __name__ == "__main__":
     import sys