From 7389130bf3c9866a5335b594d2626d985bd34503 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 22 Aug 2022 22:21:32 -0500
Subject: [PATCH] capture_call: also accept file-like
---
doc/runtime_program.rst | 7 +++++--
pyopencl/__init__.py | 4 ++--
pyopencl/capture_call.py | 9 ++++++---
3 files changed, 13 insertions(+), 7 deletions(-)
diff --git a/doc/runtime_program.rst b/doc/runtime_program.rst
index 84e165f6..a60a8f62 100644
--- a/doc/runtime_program.rst
+++ b/doc/runtime_program.rst
@@ -19,7 +19,7 @@ Program
caching is performed. To retain the compiled version of a kernel
in memory, simply retain the :class:`Program` and/or :class:`Kernel`
objects.
-
+
PyOpenCL will also cache "invokers", which are short snippets of Python
that are generated to accelerate passing arguments to and enqueuing
a kernel. For now, this caching is unaffected by ``PYOPENCL_NO_CACHE``.
@@ -298,7 +298,7 @@ Kernel
Added the *allow_empty_ndrange* keyword argument.
- .. method:: capture_call(filename, queue, global_size, local_size, *args, global_offset=None, wait_for=None, g_times_l=False)
+ .. method:: capture_call(output_file, queue, global_size, local_size, *args, global_offset=None, wait_for=None, g_times_l=False)
This method supports the exact same interface as :meth:`__call__`, but
instead of invoking the kernel, it writes a self-contained PyOpenCL program
@@ -310,6 +310,9 @@ Kernel
an observed problem. It can also help separate a misbehaving kernel from
a potentially large or time-consuming outer code.
+ :arg output_file: a a filename or a file-like to which the generated
+ code is to be written.
+
To use, simply change::
evt = my_kernel(queue, gsize, lsize, arg1, arg2, ...)
diff --git a/pyopencl/__init__.py b/pyopencl/__init__.py
index 0e880bc2..eaf90963 100644
--- a/pyopencl/__init__.py
+++ b/pyopencl/__init__.py
@@ -900,10 +900,10 @@ def _add_functionality():
# kernel_set_scalar_arg_dtypes is called.
return self._enqueue(self, queue, global_size, local_size, *args, **kwargs)
- def kernel_capture_call(self, filename, queue, global_size, local_size,
+ def kernel_capture_call(self, output_file, queue, global_size, local_size,
*args, **kwargs):
from pyopencl.capture_call import capture_kernel_call
- capture_kernel_call(self, filename, queue, global_size, local_size,
+ capture_kernel_call(self, output_file, queue, global_size, local_size,
*args, **kwargs)
def kernel_get_info(self, param_name):
diff --git a/pyopencl/capture_call.py b/pyopencl/capture_call.py
index 86736531..c0648b21 100644
--- a/pyopencl/capture_call.py
+++ b/pyopencl/capture_call.py
@@ -26,7 +26,7 @@ import numpy as np
from pytools.py_codegen import PythonCodeGenerator, Indentation
-def capture_kernel_call(kernel, filename, queue, g_size, l_size, *args, **kwargs):
+def capture_kernel_call(kernel, output_file, queue, g_size, l_size, *args, **kwargs):
try:
source = kernel._source
except AttributeError:
@@ -168,5 +168,8 @@ def capture_kernel_call(kernel, filename, queue, g_size, l_size, *args, **kwargs
# }}}
- with open(filename, "w") as outf:
- outf.write(cg.get())
+ if isinstance(output_file, str):
+ with open(output_file, "w") as outf:
+ outf.write(cg.get())
+ else:
+ output_file.write(cg.get())
--
GitLab