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)
