def test_list_builder(ctx_factory): from pytest import importorskip importorskip("mako") context = ctx_factory() queue = cl.CommandQueue(context) from pyopencl.algorithm import ListOfListsBuilder builder = ListOfListsBuilder(context, [("mylist", np.int32)], """//CL// void generate(LIST_ARG_DECL USER_ARG_DECL index_type i) { int count = i % 4; for (int j = 0; j < count; ++j) { APPEND_mylist(count); } } """, arg_decls=[]) result, evt = builder(queue, 2000) inf = result["mylist"] 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 < i / 5; ++j) { APPEND_mylist1(j); APPEND_mylist2(j + 1); } } } """, arg_decls=[], eliminate_empty_output_lists=True) result, evt = builder(queue, 1000) mylist1 = result["mylist1"] assert mylist1.count == 19900 assert (mylist1.starts.get()[:5] == [0, 1, 3, 6, 10]).all() assert (mylist1.nonempty_indices.get()[:5] == [5, 10, 15, 20, 25]).all() assert (mylist1.lists.get()[:6] == [0, 0, 1, 0, 1, 2]).all() mylist2 = result["mylist2"] assert mylist2.count == 19900 assert (mylist2.lists.get()[:6] == [1, 1, 2, 1, 2, 3]).all()
def filter_target_lists_in_user_order(queue, tree, flags): """ :arg flags: an array of length :attr:`boxtree.Tree.ntargets` of :class:`numpy.int8` objects, which indicate by being zero that the corresponding target (in user target order) is not part of the filtered list, or by being nonzero that it is. :returns: A :class:`FilteredTargetListsInUserOrder` """ user_order_flags = flags del flags user_target_ids = cl.array.empty(queue, tree.ntargets, tree.sorted_target_ids.dtype) user_target_ids[tree.sorted_target_ids] = cl.array.arange( queue, tree.ntargets, user_target_ids.dtype) from pyopencl.tools import VectorArg, dtype_to_ctype from pyopencl.algorithm import ListOfListsBuilder from mako.template import Template builder = ListOfListsBuilder( queue.context, [("filt_tgt_list", tree.particle_id_dtype)], Template("""//CL// typedef ${dtype_to_ctype(particle_id_dtype)} particle_id_t; void generate(LIST_ARG_DECL USER_ARG_DECL index_type i) { particle_id_t b_t_start = box_target_starts[i]; particle_id_t b_t_count = box_target_counts_nonchild[i]; for (particle_id_t j = b_t_start; j < b_t_start+b_t_count; ++j) { particle_id_t user_target_id = user_target_ids[j]; if (user_order_flags[user_target_id]) { APPEND_filt_tgt_list(user_target_id); } } } """, strict_undefined=True).render( dtype_to_ctype=dtype_to_ctype, particle_id_dtype=tree.particle_id_dtype), arg_decls=[ VectorArg(user_order_flags.dtype, "user_order_flags"), VectorArg(tree.particle_id_dtype, "user_target_ids"), VectorArg(tree.particle_id_dtype, "box_target_starts"), VectorArg(tree.particle_id_dtype, "box_target_counts_nonchild"), ]) result, evt = builder(queue, tree.nboxes, user_order_flags.data, user_target_ids.data, tree.box_target_starts.data, tree.box_target_counts_nonchild.data) return FilteredTargetListsInUserOrder( nfiltered_targets=result["filt_tgt_list"].count, target_starts=result["filt_tgt_list"].starts, target_lists=result["filt_tgt_list"].lists, ).with_queue(None)
def test_list_builder_with_offset(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", with_offset=True)]) n = 10000 input_list = cl.array.zeros(queue, (n + 10,), float) input_list[10:] = 1 result, evt = builder(queue, n, input_list[10:]) inf = result["mylist"] assert inf.count == n assert (inf.lists.get() == 1).all()
def test_list_builder(ctx_factory): from pytest import importorskip importorskip("mako") context = ctx_factory() queue = cl.CommandQueue(context) from pyopencl.algorithm import ListOfListsBuilder builder = ListOfListsBuilder(context, [("mylist", np.int32)], """//CL// void generate(LIST_ARG_DECL USER_ARG_DECL index_type i) { int count = i % 4; for (int j = 0; j < count; ++j) { APPEND_mylist(count); } } """, arg_decls=[]) result, evt = builder(queue, 2000) inf = result["mylist"] assert inf.count == 3000 assert (inf.lists.get()[-6:] == [1, 2, 2, 3, 3, 3]).all()
def get_area_query_kernel(self, dimensions, coord_dtype, box_id_dtype, ball_id_dtype, peer_list_idx_dtype, max_levels): from pyopencl.tools import dtype_to_ctype from boxtree import box_flags_enum logger.info("start building area query kernel") from boxtree.traversal import TRAVERSAL_PREAMBLE_TEMPLATE from boxtree.tree_build import TreeBuilder template = Template(TRAVERSAL_PREAMBLE_TEMPLATE + AREA_QUERY_TEMPLATE, strict_undefined=True) render_vars = dict( dimensions=dimensions, dtype_to_ctype=dtype_to_ctype, box_id_dtype=box_id_dtype, particle_id_dtype=None, coord_dtype=coord_dtype, vec_types=cl.cltypes.vec_types, max_levels=max_levels, AXIS_NAMES=AXIS_NAMES, box_flags_enum=box_flags_enum, peer_list_idx_dtype=peer_list_idx_dtype, ball_id_dtype=ball_id_dtype, debug=False, root_extent_stretch_factor=TreeBuilder.ROOT_EXTENT_STRETCH_FACTOR, # Not used (but required by TRAVERSAL_PREAMBLE_TEMPLATE) stick_out_factor=0) from pyopencl.tools import VectorArg, ScalarArg arg_decls = [ VectorArg(coord_dtype, "box_centers"), ScalarArg(coord_dtype, "root_extent"), VectorArg(np.uint8, "box_levels"), ScalarArg(box_id_dtype, "aligned_nboxes"), VectorArg(box_id_dtype, "box_child_ids"), VectorArg(box_flags_enum.dtype, "box_flags"), VectorArg(peer_list_idx_dtype, "peer_list_starts"), VectorArg(box_id_dtype, "peer_lists"), VectorArg(coord_dtype, "ball_radii"), ] + [ ScalarArg(coord_dtype, "bbox_min_" + ax) for ax in AXIS_NAMES[:dimensions] ] + [ VectorArg(coord_dtype, "ball_" + ax) for ax in AXIS_NAMES[:dimensions] ] from pyopencl.algorithm import ListOfListsBuilder area_query_kernel = ListOfListsBuilder( self.context, [("leaves", box_id_dtype)], str(template.render(**render_vars)), arg_decls=arg_decls, name_prefix="area_query", count_sharing={}, complex_kernel=True) logger.info("done building area query kernel") return area_query_kernel
def get_filter_target_lists_in_user_order_kernel(self, particle_id_dtype, user_order_flags_dtype): from pyopencl.tools import VectorArg, dtype_to_ctype from pyopencl.algorithm import ListOfListsBuilder from mako.template import Template builder = ListOfListsBuilder(self.context, [("filt_tgt_list", particle_id_dtype)], Template("""//CL// typedef ${dtype_to_ctype(particle_id_dtype)} particle_id_t; void generate(LIST_ARG_DECL USER_ARG_DECL index_type i) { particle_id_t b_t_start = box_target_starts[i]; particle_id_t b_t_count = box_target_counts_nonchild[i]; for (particle_id_t j = b_t_start; j < b_t_start+b_t_count; ++j) { particle_id_t user_target_id = user_target_ids[j]; if (user_order_flags[user_target_id]) { APPEND_filt_tgt_list(user_target_id); } } } """, strict_undefined=True).render( dtype_to_ctype=dtype_to_ctype, particle_id_dtype=particle_id_dtype ), arg_decls=[ VectorArg(user_order_flags_dtype, "user_order_flags"), VectorArg(particle_id_dtype, "user_target_ids"), VectorArg(particle_id_dtype, "box_target_starts"), VectorArg(particle_id_dtype, "box_target_counts_nonchild"), ]) return builder
def get_balls_to_leaves_kernel(self, dimensions, coord_dtype, box_id_dtype, ball_id_dtype, max_levels, stick_out_factor): from pyopencl.tools import dtype_to_ctype from boxtree import box_flags_enum render_vars = dict( dimensions=dimensions, dtype_to_ctype=dtype_to_ctype, box_id_dtype=box_id_dtype, particle_id_dtype=None, ball_id_dtype=ball_id_dtype, coord_dtype=coord_dtype, vec_types=cl.array.vec.types, max_levels=max_levels, AXIS_NAMES=AXIS_NAMES, box_flags_enum=box_flags_enum, debug=False, stick_out_factor=stick_out_factor, ) logger.info("start building leaves-to-balls lookup kernel") from boxtree.traversal import TRAVERSAL_PREAMBLE_TEMPLATE src = Template(TRAVERSAL_PREAMBLE_TEMPLATE + BALLS_TO_LEAVES_TEMPLATE, strict_undefined=True).render(**render_vars) from pyopencl.tools import VectorArg, ScalarArg from pyopencl.algorithm import ListOfListsBuilder result = ListOfListsBuilder( self.context, [ ("ball_numbers", ball_id_dtype), ("overlapping_leaves", box_id_dtype), ], str(src), arg_decls=[ VectorArg(box_flags_enum.dtype, "box_flags"), VectorArg(coord_dtype, "box_centers"), VectorArg(box_id_dtype, "box_child_ids"), VectorArg(np.uint8, "box_levels"), ScalarArg(coord_dtype, "root_extent"), ScalarArg(box_id_dtype, "aligned_nboxes"), VectorArg(coord_dtype, "ball_radii"), ] + [ VectorArg(coord_dtype, "ball_" + ax) for ax in AXIS_NAMES[:dimensions] ], name_prefix="circles_to_balls", count_sharing={ # /!\ This makes a promise that APPEND_ball_numbers will # always occur *before* APPEND_overlapping_leaves. "overlapping_leaves": "ball_numbers" }, complex_kernel=True) logger.info("done building leaves-to-balls lookup kernel") return result
def get_peer_list_finder_kernel(self, dimensions, coord_dtype, box_id_dtype, max_levels): from pyopencl.tools import dtype_to_ctype from boxtree import box_flags_enum logger.info("start building peer list finder kernel") from boxtree.traversal import (TRAVERSAL_PREAMBLE_TEMPLATE, HELPER_FUNCTION_TEMPLATE) template = Template(TRAVERSAL_PREAMBLE_TEMPLATE + HELPER_FUNCTION_TEMPLATE + PEER_LIST_FINDER_TEMPLATE, strict_undefined=True) render_vars = dict( dimensions=dimensions, dtype_to_ctype=dtype_to_ctype, box_id_dtype=box_id_dtype, particle_id_dtype=None, coord_dtype=coord_dtype, vec_types=cl.cltypes.vec_types, max_levels=max_levels, AXIS_NAMES=AXIS_NAMES, box_flags_enum=box_flags_enum, debug=False, # Not used (but required by TRAVERSAL_PREAMBLE_TEMPLATE) stick_out_factor=0, # For calls to the helper is_adjacent_or_overlapping() targets_have_extent=False, sources_have_extent=False) from pyopencl.tools import VectorArg, ScalarArg arg_decls = [ VectorArg(coord_dtype, "box_centers"), ScalarArg(coord_dtype, "root_extent"), VectorArg(np.uint8, "box_levels"), ScalarArg(box_id_dtype, "aligned_nboxes"), VectorArg(box_id_dtype, "box_child_ids"), VectorArg(box_flags_enum.dtype, "box_flags"), ] from pyopencl.algorithm import ListOfListsBuilder peer_list_finder_kernel = ListOfListsBuilder( self.context, [("peers", box_id_dtype)], str(template.render(**render_vars)), arg_decls=arg_decls, name_prefix="find_peer_lists", count_sharing={}, complex_kernel=True) logger.info("done building peer list finder kernel") return peer_list_finder_kernel
def get_kernel_info(self, dimensions, particle_id_dtype, box_id_dtype, coord_dtype, box_level_dtype, max_levels, sources_are_targets, sources_have_extent, targets_have_extent, stick_out_factor): logging.info("building traversal build kernels") debug = False from pyopencl.tools import dtype_to_ctype from boxtree.tree import box_flags_enum render_vars = dict( dimensions=dimensions, dtype_to_ctype=dtype_to_ctype, particle_id_dtype=particle_id_dtype, box_id_dtype=box_id_dtype, box_flags_enum=box_flags_enum, coord_dtype=coord_dtype, vec_types=cl.array.vec.types, max_levels=max_levels, AXIS_NAMES=AXIS_NAMES, debug=debug, sources_are_targets=sources_are_targets, sources_have_extent=sources_have_extent, targets_have_extent=targets_have_extent, stick_out_factor=stick_out_factor, ) from pyopencl.algorithm import ListOfListsBuilder from pyopencl.tools import VectorArg, ScalarArg result = {} # {{{ source boxes, their parents, target boxes src = Template( TRAVERSAL_PREAMBLE_TEMPLATE + SOURCES_PARENTS_AND_TARGETS_TEMPLATE, strict_undefined=True).render(**render_vars) result["sources_parents_and_targets_builder"] = \ ListOfListsBuilder(self.context, [ ("source_parent_boxes", box_id_dtype), ("source_boxes", box_id_dtype), ("target_or_target_parent_boxes", box_id_dtype) ] + ( [("target_boxes", box_id_dtype)] if not sources_are_targets else []), str(src), arg_decls=[ VectorArg(box_flags_enum.dtype, "box_flags"), ], debug=debug, name_prefix="sources_parents_and_targets") result["level_start_box_nrs_extractor"] = \ LEVEL_START_BOX_NR_EXTRACTOR_TEMPLATE.build(self.context, type_aliases=( ("box_id_t", box_id_dtype), ("box_level_t", box_level_dtype), ), ) # }}} # {{{ build list N builders base_args = [ VectorArg(coord_dtype, "box_centers"), ScalarArg(coord_dtype, "root_extent"), VectorArg(np.uint8, "box_levels"), ScalarArg(box_id_dtype, "aligned_nboxes"), VectorArg(box_id_dtype, "box_child_ids"), VectorArg(box_flags_enum.dtype, "box_flags"), ] for list_name, template, extra_args, extra_lists in [ ("colleagues", COLLEAGUES_TEMPLATE, [], []), ("neighbor_source_boxes", NEIGBHOR_SOURCE_BOXES_TEMPLATE, [ VectorArg(box_id_dtype, "target_boxes"), ], []), ("sep_siblings", SEP_SIBLINGS_TEMPLATE, [ VectorArg(box_id_dtype, "target_or_target_parent_boxes"), VectorArg(box_id_dtype, "box_parent_ids"), VectorArg(box_id_dtype, "colleagues_starts"), VectorArg(box_id_dtype, "colleagues_list"), ], []), ("sep_smaller", SEP_SMALLER_TEMPLATE, [ VectorArg(box_id_dtype, "target_boxes"), VectorArg(box_id_dtype, "colleagues_starts"), VectorArg(box_id_dtype, "colleagues_list"), ], ["sep_close_smaller"] if sources_have_extent or targets_have_extent else []), ("sep_bigger", SEP_BIGGER_TEMPLATE, [ VectorArg(box_id_dtype, "target_or_target_parent_boxes"), VectorArg(box_id_dtype, "box_parent_ids"), VectorArg(box_id_dtype, "colleagues_starts"), VectorArg(box_id_dtype, "colleagues_list"), ], ["sep_close_bigger"] if sources_have_extent or targets_have_extent else []), ]: src = Template( TRAVERSAL_PREAMBLE_TEMPLATE + HELPER_FUNCTION_TEMPLATE + template, strict_undefined=True).render(**render_vars) result[list_name+"_builder"] = ListOfListsBuilder(self.context, [(list_name, box_id_dtype)] + [(extra_list_name, box_id_dtype) for extra_list_name in extra_lists], str(src), arg_decls=base_args + extra_args, debug=debug, name_prefix=list_name, complex_kernel=True) # }}} logging.info("traversal build kernels built") return _KernelInfo(**result)