add add_profiling_event

Author: matthiasdiener

arraycontext/impl/pytato/__init__.py

@@ -374,6 +374,11 @@ def _wait_and_transfer_profile_events(self) -> None:
 
         self.profile_events = []
 
+    def add_profiling_event(self, evt: cl._cl.Event, translation_unit: Any) -> None:
+        """Add a profiling event to the list of profiling events."""
+        if self.profile_kernels:
+            self.profile_events.append(ProfileEvent(evt, translation_unit))
+
     def get_profiling_data_for_kernel(self, kernel_name: str) \
           -> MultiCallKernelProfile:
         """Return profiling data for kernel *kernel_name*."""
@@ -660,8 +665,7 @@ def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
                 allocator=self.allocator,
                 **bound_arguments)
 
-        if self.profile_kernels:
-            self.profile_events.append(ProfileEvent(evt, pt_prg))
+        self.add_profiling_event(evt, pt_prg)
 
         assert len(set(out_dict) & set(key_to_frozen_subary)) == 0
 

arraycontext/impl/pytato/compile.py

@@ -640,9 +640,7 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
                                             allocator=self.actx.allocator,
                                             **input_kwargs_for_loopy)
 
-        if self.actx.profile_kernels:
-            from arraycontext.impl.pytato import ProfileEvent
-            self.actx.profile_events.append(ProfileEvent(evt, self.pytato_program))
+        self.actx.add_profiling_event(evt, self.pytato_program)
 
         def to_output_template(keys, _):
             name_in_program = self.output_id_to_name_in_program[keys]
@@ -683,9 +681,7 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
                                             allocator=self.actx.allocator,
                                             **input_kwargs_for_loopy)
 
-        if self.actx.profile_kernels:
-            from arraycontext.impl.pytato import ProfileEvent
-            self.actx.profile_events.append(ProfileEvent(evt, self.pytato_program))
+        self.actx.add_profiling_event(evt, self.pytato_program)
 
         return self.actx.thaw(to_tagged_cl_array(out_dict[self.output_name],
                                                  axes=get_cl_axes_from_pt_axes(