diff --git a/pyopencl/algorithm.py b/pyopencl/algorithm.py index 92730750e18e1490ac65a159c3bd54ee5e2a9f96..dfa5d4bbfd0721915e3979c9cd027d62ecab63f1 100644 --- a/pyopencl/algorithm.py +++ b/pyopencl/algorithm.py @@ -1042,8 +1042,11 @@ class ListOfListsBuilder: def __call__(self, queue, n_objects, *args, **kwargs): """ :arg args: arguments corresponding to arg_decls in the constructor. - :class:`pyopencl.array.Array` are not allowed directly and should - be passed as their :attr:`pyopencl.array.Array.data` attribute instead. + Array-like arguments must be either + 1D :class:`pyopencl.array.Array` objects or + :class:`pyopencl.MemoryObject` objects, of which the latter + can be obtained from a :class:`pyopencl.array.Array` using the + :attr:`pyopencl.array.Array.data` attribute. :arg allocator: optionally, the allocator to use to allocate new arrays. :arg omit_lists: An iterable of list names that should *not* be built @@ -1111,8 +1114,30 @@ class ListOfListsBuilder: if self.eliminate_empty_output_lists: compress_kernel = self.get_compress_kernel(index_dtype) - from pyopencl.tools import expand_runtime_arg_list - args = expand_runtime_arg_list(self.arg_decls, args) + data_args = [] + for i, (arg_descr, arg_val) in enumerate(zip(self.arg_decls, args)): + from pyopencl.tools import VectorArg + if isinstance(arg_descr, VectorArg): + from pyopencl import MemoryObject + if isinstance(arg_val, MemoryObject): + data_args.append(arg_val) + if arg_descr.with_offset: + raise ValueError( + "with_offset=True specified for argument %d " + "but the argument is not an array" % i) + continue + + if arg_val.ndim != 1: + raise ValueError("argument %d is a multidimensional array" % i) + + data_args.append(arg_val.base_data) + if arg_descr.with_offset: + data_args.append(arg_val.offset) + else: + data_args.append(arg_val) + + del args + data_args = tuple(data_args) # {{{ allocate memory for counts @@ -1151,7 +1176,7 @@ class ListOfListsBuilder: gsize, lsize = splay(queue, n_objects) count_event = count_kernel(queue, gsize, lsize, - *(tuple(count_list_args) + args + (n_objects,)), + *(tuple(count_list_args) + data_args + (n_objects,)), **dict(wait_for=wait_for)) compress_events = {} @@ -1257,7 +1282,7 @@ class ListOfListsBuilder: # }}} evt = write_kernel(queue, gsize, lsize, - *(tuple(write_list_args) + args + (n_objects,)), + *(tuple(write_list_args) + data_args + (n_objects,)), **dict(wait_for=scan_events)) return result, evt diff --git a/pyopencl/scan.py b/pyopencl/scan.py index 6218d5716b409557341f4e8ce12105ad0dcf1a30..cef68c67b242ef08194da2775c79ffe564718e12 100644 --- a/pyopencl/scan.py +++ b/pyopencl/scan.py @@ -1480,8 +1480,15 @@ class GenericScanKernel(_GenericScanKernelBase): # We're done here. (But pretend to return an event.) return cl.enqueue_marker(queue, wait_for=wait_for) - from pyopencl.tools import expand_runtime_arg_list - data_args = list(expand_runtime_arg_list(self.parsed_args, args)) + data_args = [] + for arg_descr, arg_val in zip(self.parsed_args, args): + from pyopencl.tools import VectorArg + if isinstance(arg_descr, VectorArg): + data_args.append(arg_val.base_data) + if arg_descr.with_offset: + data_args.append(arg_val.offset) + else: + data_args.append(arg_val) # }}} diff --git a/pyopencl/tools.py b/pyopencl/tools.py index b1dcf9f23e84acce7c83a2f9491e2a9cb9188f6b..369fb2729cdb509f50ca07e76cba0ce58055ed4b 100644 --- a/pyopencl/tools.py +++ b/pyopencl/tools.py @@ -400,18 +400,6 @@ def get_arg_offset_adjuster_code(arg_types): return "\n".join(result) - -def expand_runtime_arg_list(args, user_args): - data_args = [] - for arg_descr, arg_val in zip(args, user_args): - if isinstance(arg_descr, VectorArg): - data_args.append(arg_val.base_data) - if arg_descr.with_offset: - data_args.append(arg_val.offset) - else: - data_args.append(arg_val) - return tuple(data_args) - # }}} diff --git a/test/test_algorithm.py b/test/test_algorithm.py index 5c09b565bac9440a862e9f98fff1bddc06ad1d7b..0360d6a348b1e1e1ab46f6bb4df0785252b2bea4 100644 --- a/test/test_algorithm.py +++ b/test/test_algorithm.py @@ -880,6 +880,31 @@ def test_list_builder(ctx_factory): assert (inf.lists.get()[-6:] == [1, 2, 2, 3, 3, 3]).all() +def test_list_builder_with_memoryobject(ctx_factory): + from pytest import importorskip + importorskip("mako") + + context = ctx_factory() + queue = cl.CommandQueue(context) + + from pyopencl.algorithm import ListOfListsBuilder + from pyopencl.tools import VectorArg + builder = ListOfListsBuilder(context, [("mylist", np.int32)], """//CL// + void generate(LIST_ARG_DECL USER_ARG_DECL index_type i) + { + APPEND_mylist(input_list[i]); + } + """, arg_decls=[VectorArg(float, "input_list")]) + + n = 10000 + input_list = cl.array.zeros(queue, (n,), float) + result, evt = builder(queue, n, input_list.data) + + inf = result["mylist"] + assert inf.count == n + assert (inf.lists.get() == 0).all() + + def test_list_builder_with_offset(ctx_factory): from pytest import importorskip importorskip("mako")