Add undocumented capture_as kwarg to ElementwiseKernel

inducer · Aug 24, 2015 · 671ad7f · 671ad7f
1 parent fe5b8fd
commit 671ad7f
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/pyopencl/elementwise.py b/pyopencl/elementwise.py
@@ -245,6 +245,7 @@ def __call__(self, *args, **kwargs):
 
         range_ = kwargs.pop("range", None)
         slice_ = kwargs.pop("slice", None)
+        capture_as = kwargs.pop("capture_as", None)
 
         use_range = range_ is not None or slice_ is not None
         kernel, arg_descrs = self.get_kernel(use_range)
@@ -310,6 +311,12 @@ def __call__(self, *args, **kwargs):
             invocation_args.append(repr_vec.size)
             gs, ls = repr_vec.get_sizes(queue, max_wg_size)
 
+        if capture_as is not None:
+            kernel.set_args(*invocation_args)
+            kernel.capture_call(
+                    capture_as, queue,
+                    gs, ls, *invocation_args, wait_for=wait_for)
+
         kernel.set_args(*invocation_args)
         return cl.enqueue_nd_range_kernel(queue, kernel,
                 gs, ls, wait_for=wait_for)