diff --git a/pyopencl/algorithm.py b/pyopencl/algorithm.py index 5083b095e2d881d50b8de45d7e17b2e8dda4c415..5f8287cf3ec530cf2d9a9ebcf21b705d19e02bbe 100644 --- a/pyopencl/algorithm.py +++ b/pyopencl/algorithm.py @@ -736,7 +736,8 @@ class ListOfListsBuilder: def __init__(self, context, list_names_and_dtypes, generate_template, arg_decls, count_sharing=None, devices=None, name_prefix="plb_build_list", options=[], preamble="", - debug=False, complex_kernel=False): + debug=False, complex_kernel=False, + eliminate_empty_output_lists=False): """ :arg context: A :class:`pyopencl.Context`. :arg list_names_and_dtypes: a list of `(name, dtype)` tuples @@ -810,6 +811,7 @@ class ListOfListsBuilder: self.debug = debug self.complex_kernel = complex_kernel + self.eliminate_empty_output_lists = eliminate_empty_output_lists # {{{ kernel generators diff --git a/test/test_algorithm.py b/test/test_algorithm.py index 2e1537dfa2db1de92c00185fe094fdce37704e0e..ab8fa6f7c39f0de27e08be5f7c38c016678592d1 100644 --- a/test/test_algorithm.py +++ b/test/test_algorithm.py @@ -847,6 +847,36 @@ def test_list_builder(ctx_factory): assert inf.count == 3000 assert (inf.lists.get()[-6:] == [1, 2, 2, 3, 3, 3]).all() + builder = ListOfListsBuilder( + context, + [("mylist1", np.int32), ("mylist2", np.int32)], + """//CL// + void generate(LIST_ARG_DECL USER_ARG_DECL index_type i) + { + if (i % 5 == 0) + { + for (int j = 0; j < 10; ++j) + { + APPEND_mylist1(j); + APPEND_mylist2(1); + } + } + } + """, + arg_decls=[], + eliminate_empty_output_lists=True) + + result, evt = builder(queue, 1000) + + mylist1 = result["mylist1"] + assert mylist1.count == 2000 + assert (mylist1.starts.get()[:5] == [0, 10, 20, 30, 40]).all() + assert (mylist1.indices.get()[:5] == [0, 5, 10, 15, 20]).all() + assert (mylist1.lists.get()[:5] == [0, 1, 2, 3, 4]).all() + mylist2 = result["mylist2"] + assert mylist2.count == 2000 + assert (mylist2.lists.get()[:5] == [1, 1, 1, 1, 1]).all() + def test_key_value_sorter(ctx_factory): from pytest import importorskip