def _remove_inames_for_shared_hw_axes(self, cond_inames):
"""
See if cond_inames contains references to two (or more) inames that
boil down to the same tag. If so, exclude them. (We shouldn't be writing
conditionals for such inames because we would be implicitly restricting
the other inames as well.)
"""
tag_key_uses = defaultdict(list)
from loopy.kernel.data import HardwareConcurrentTag
for iname in cond_inames:
tags = self.iname_tags_of_type(iname, HardwareConcurrentTag, max_num=1)
if tags:
tag, = tags
tag_key_uses[tag.key].append(iname)
multi_use_keys = set(
key for key, user_inames in six.iteritems(tag_key_uses)
if len(user_inames) > 1)
multi_use_inames = set()
for iname in cond_inames:
tags = self.iname_tags_of_type(iname, HardwareConcurrentTag)
if tags:
tag, = filter_iname_tags_by_type(tags, HardwareConcurrentTag, 1)
if tag.key in multi_use_keys:
multi_use_inames.add(iname)
return frozenset(cond_inames - multi_use_inames)
def _remove_inames_for_shared_hw_axes(self, cond_inames):
"""
See if cond_inames contains references to two (or more) inames that
boil down to the same tag. If so, exclude them. (We shouldn't be writing
conditionals for such inames because we would be implicitly restricting
the other inames as well.)
"""
tag_key_uses = defaultdict(list)
from loopy.kernel.data import HardwareConcurrentTag
for iname in cond_inames:
tags = self.iname_tags_of_type(iname, HardwareConcurrentTag, max_num=1)
if tags:
tag, = tags
tag_key_uses[tag.key].append(iname)
multi_use_keys = set(
key for key, user_inames in six.iteritems(tag_key_uses)
if len(user_inames) > 1)
multi_use_inames = set()
for iname in cond_inames:
tags = self.iname_tags_of_type(iname, HardwareConcurrentTag)
if tags:
tag, = filter_iname_tags_by_type(tags, HardwareConcurrentTag, 1)
if tag.key in multi_use_keys:
multi_use_inames.add(iname)
return frozenset(cond_inames - multi_use_inames)
def check_multiple_tags_allowed(kernel):
from loopy.kernel.data import (GroupIndexTag, LocalIndexTag, VectorizeTag,
UnrollTag, ForceSequentialTag, IlpBaseTag, filter_iname_tags_by_type)
illegal_combinations = [
(GroupIndexTag, LocalIndexTag, VectorizeTag, UnrollTag, ForceSequentialTag),
(IlpBaseTag, ForceSequentialTag)
]
for iname, tags in six.iteritems(kernel.iname_to_tags):
for comb in illegal_combinations:
if len(filter_iname_tags_by_type(tags, comb)) > 1:
raise LoopyError("iname {0} has illegal combination of "
"tags: {1}".format(iname, tags))
def check_multiple_tags_allowed(kernel):
from loopy.kernel.data import (GroupIndexTag, LocalIndexTag, VectorizeTag,
UnrollTag, ForceSequentialTag, IlpBaseTag, filter_iname_tags_by_type)
illegal_combinations = [
(GroupIndexTag, LocalIndexTag, VectorizeTag, UnrollTag, ForceSequentialTag),
(IlpBaseTag, ForceSequentialTag)
]
for iname, tags in six.iteritems(kernel.iname_to_tags):
for comb in illegal_combinations:
if len(filter_iname_tags_by_type(tags, comb)) > 1:
raise LoopyError("iname {0} has illegal combination of "
"tags: {1}".format(iname, tags))
def get_admissible_conditional_inames_for(codegen_state, sched_index):
"""This function disallows conditionals on local-idx tagged
inames if there is a barrier nested somewhere within.
"""
kernel = codegen_state.kernel
from loopy.kernel.data import (LocalIndexTag, HardwareConcurrentTag,
filter_iname_tags_by_type)
from loopy.schedule import find_active_inames_at, has_barrier_within
result = find_active_inames_at(kernel, sched_index)
has_barrier = has_barrier_within(kernel, sched_index)
for iname, tags in six.iteritems(kernel.iname_to_tags):
if (filter_iname_tags_by_type(tags, HardwareConcurrentTag)
and codegen_state.is_generating_device_code):
if not has_barrier or not filter_iname_tags_by_type(tags, LocalIndexTag):
result.add(iname)
return frozenset(result)
def check_multiple_tags_allowed(kernel):
"""
Checks if a multiple tags of an iname are compatible.
"""
from loopy.kernel.data import (GroupIndexTag, LocalIndexTag, VectorizeTag,
UnrollTag, ForceSequentialTag, IlpBaseTag,
filter_iname_tags_by_type)
illegal_combinations = [(GroupIndexTag, LocalIndexTag, VectorizeTag,
UnrollTag, ForceSequentialTag),
(IlpBaseTag, ForceSequentialTag)]
for iname in kernel.inames.values():
for comb in illegal_combinations:
if len(filter_iname_tags_by_type(iname.tags, comb)) > 1:
raise LoopyError("iname {} has illegal combination of "
"tags: {}".format(iname.name, iname.tags))
def iname_tags_of_type(self, iname, tag_type_or_types,
max_num=None, min_num=None):
"""Return a subset of *tags* that matches type *tag_type*. Raises exception
if the number of tags found were greater than *max_num* or less than
*min_num*.
:arg tags: An iterable of tags.
:arg tag_type_or_types: a subclass of :class:`loopy.kernel.data.IndexTag`.
:arg max_num: the maximum number of tags expected to be found.
:arg min_num: the minimum number of tags expected to be found.
"""
from loopy.kernel.data import filter_iname_tags_by_type
return filter_iname_tags_by_type(
self.iname_to_tags.get(iname, frozenset()),
tag_type_or_types, max_num=max_num, min_num=min_num)
def iname_tags_of_type(self, iname, tag_type_or_types,
max_num=None, min_num=None):
"""Return a subset of *tags* that matches type *tag_type*. Raises exception
if the number of tags found were greater than *max_num* or less than
*min_num*.
:arg tags: An iterable of tags.
:arg tag_type_or_types: a subclass of :class:`loopy.kernel.data.IndexTag`.
:arg max_num: the maximum number of tags expected to be found.
:arg min_num: the minimum number of tags expected to be found.
"""
from loopy.kernel.data import filter_iname_tags_by_type
return filter_iname_tags_by_type(
self.iname_to_tags.get(iname, frozenset()),
tag_type_or_types, max_num=max_num, min_num=min_num)
def get_hw_axis_sizes_and_tags_for_save_slot(self, temporary):
"""
This is used for determining the amount of global storage needed for saving
and restoring the temporary across kernel calls, due to hardware
parallel inames (the inferred axes get prefixed to the number of
dimensions in the temporary).
In the case of local temporaries, inames that are tagged
hw-local do not contribute to the global storage shape.
"""
accessor_insn_ids = frozenset(
self.kernel.reader_map()[temporary.name]
| self.kernel.writer_map()[temporary.name])
group_tags = None
local_tags = None
def _sortedtags(tags):
return sorted(tags, key=lambda tag: tag.axis)
for insn_id in accessor_insn_ids:
insn = self.kernel.id_to_insn[insn_id]
my_group_tags = []
my_local_tags = []
for iname in insn.within_inames:
tags = self.kernel.iname_tags(iname)
if not tags:
continue
from loopy.kernel.data import (GroupInameTag, LocalInameTag,
ConcurrentTag,
filter_iname_tags_by_type)
if filter_iname_tags_by_type(tags, GroupInameTag):
tag, = filter_iname_tags_by_type(tags, GroupInameTag, 1)
my_group_tags.append(tag)
elif filter_iname_tags_by_type(tags, LocalInameTag):
tag, = filter_iname_tags_by_type(tags, LocalInameTag, 1)
my_local_tags.append(tag)
elif filter_iname_tags_by_type(tags, ConcurrentTag):
raise LoopyError("iname '%s' is tagged with '%s' - only "
"group and local tags are supported for "
"auto save/reload of temporaries" %
(iname, tags))
if group_tags is None:
group_tags = _sortedtags(my_group_tags)
local_tags = _sortedtags(my_local_tags)
group_tags_originating_insn_id = insn_id
if (group_tags != _sortedtags(my_group_tags)
or local_tags != _sortedtags(my_local_tags)):
raise LoopyError(
"inconsistent parallel tags across instructions that access "
"'%s' (specifically, instruction '%s' has tags '%s' but "
"instruction '%s' has tags '%s')" %
(temporary.name, group_tags_originating_insn_id, group_tags
+ local_tags, insn_id, my_group_tags + my_local_tags))
if group_tags is None:
assert local_tags is None
return (), ()
group_sizes, local_sizes = (
self.kernel.get_grid_sizes_for_insn_ids_as_exprs(
accessor_insn_ids, self.callables_table))
if temporary.address_space == lp.AddressSpace.LOCAL:
# Elide local axes in the save slot for local temporaries.
del local_tags[:]
local_sizes = ()
# We set hw_dims to be arranged according to the order:
# g.0 < g.1 < ... < l.0 < l.1 < ...
return (group_sizes + local_sizes), tuple(group_tags + local_tags)
def emit_assignment(self, codegen_state, insn):
kernel = codegen_state.kernel
ecm = codegen_state.expression_to_code_mapper
assignee_var_name, = insn.assignee_var_names()
lhs_var = codegen_state.kernel.get_var_descriptor(assignee_var_name)
lhs_dtype = lhs_var.dtype
if insn.atomicity:
raise NotImplementedError("atomic ops in ISPC")
from loopy.expression import dtype_to_type_context
from pymbolic.mapper.stringifier import PREC_NONE
rhs_type_context = dtype_to_type_context(kernel.target, lhs_dtype)
rhs_code = ecm(insn.expression,
prec=PREC_NONE,
type_context=rhs_type_context,
needed_dtype=lhs_dtype)
lhs = insn.assignee
# {{{ handle streaming stores
if "!streaming_store" in insn.tags:
ary = ecm.find_array(lhs)
from loopy.kernel.array import get_access_info
from pymbolic import evaluate
from loopy.symbolic import simplify_using_aff
index_tuple = tuple(
simplify_using_aff(kernel, idx) for idx in lhs.index_tuple)
access_info = get_access_info(
kernel.target, ary, index_tuple,
lambda expr: evaluate(expr, self.codegen_state.var_subst_map),
codegen_state.vectorization_info)
from loopy.kernel.data import ArrayArg, TemporaryVariable
if not isinstance(ary, (ArrayArg, TemporaryVariable)):
raise LoopyError("array type not supported in ISPC: %s" %
type(ary).__name)
if len(access_info.subscripts) != 1:
raise LoopyError("streaming stores must have a subscript")
subscript, = access_info.subscripts
from pymbolic.primitives import Sum, flattened_sum, Variable
if isinstance(subscript, Sum):
terms = subscript.children
else:
terms = (subscript.children, )
new_terms = []
from loopy.kernel.data import LocalIndexTag, filter_iname_tags_by_type
from loopy.symbolic import get_dependencies
saw_l0 = False
for term in terms:
if (isinstance(term, Variable) and kernel.iname_tags_of_type(
term.name, LocalIndexTag)):
tag, = kernel.iname_tags_of_type(term.name,
LocalIndexTag,
min_num=1,
max_num=1)
if tag.axis == 0:
if saw_l0:
raise LoopyError(
"streaming store must have stride 1 in "
"local index, got: %s" % subscript)
saw_l0 = True
continue
else:
for dep in get_dependencies(term):
if filter_iname_tags_by_type(
kernel.iname_to_tags.get(dep, []),
LocalIndexTag):
tag, = filter_iname_tags_by_type(
kernel.iname_to_tags.get(dep, []),
LocalIndexTag, 1)
if tag.axis == 0:
raise LoopyError(
"streaming store must have stride 1 in "
"local index, got: %s" % subscript)
new_terms.append(term)
if not saw_l0:
raise LoopyError("streaming store must have stride 1 in "
"local index, got: %s" % subscript)
if access_info.vector_index is not None:
raise LoopyError("streaming store may not use a short-vector "
"data type")
rhs_has_programindex = any(
isinstance(tag, LocalIndexTag) and tag.axis == 0
for tag in kernel.iname_tags(dep)
for dep in get_dependencies(insn.expression))
if not rhs_has_programindex:
rhs_code = "broadcast(%s, 0)" % rhs_code
from cgen import Statement
return Statement(
"streaming_store(%s + %s, %s)" %
(access_info.array_name,
ecm(flattened_sum(new_terms), PREC_NONE, 'i'), rhs_code))
# }}}
from cgen import Assign
return Assign(ecm(lhs, prec=PREC_NONE, type_context=None), rhs_code)
def generate_code_for_sched_index(codegen_state, sched_index):
kernel = codegen_state.kernel
sched_item = kernel.linearization[sched_index]
if isinstance(sched_item, CallKernel):
assert not codegen_state.is_generating_device_code
from loopy.schedule import (gather_schedule_block,
get_insn_ids_for_block_at)
_, past_end_i = gather_schedule_block(kernel.linearization,
sched_index)
assert past_end_i <= codegen_state.schedule_index_end
extra_args = synthesize_idis_for_extra_args(kernel, sched_index)
new_codegen_state = codegen_state.copy(
is_generating_device_code=True,
gen_program_name=sched_item.kernel_name,
schedule_index_end=past_end_i - 1,
implemented_data_info=(codegen_state.implemented_data_info +
extra_args))
from loopy.codegen.result import generate_host_or_device_program
codegen_result = generate_host_or_device_program(
new_codegen_state, sched_index)
if codegen_state.is_entrypoint:
glob_grid, loc_grid = kernel.get_grid_sizes_for_insn_ids_as_exprs(
get_insn_ids_for_block_at(kernel.linearization, sched_index),
codegen_state.callables_table)
return merge_codegen_results(codegen_state, [
codegen_result,
codegen_state.ast_builder.get_kernel_call(
codegen_state, sched_item.kernel_name, glob_grid, loc_grid,
extra_args),
])
else:
# do not generate host code for non-entrypoint kernels
return codegen_result
elif isinstance(sched_item, EnterLoop):
from loopy.kernel.data import (UnrolledIlpTag, UnrollTag,
ForceSequentialTag, LoopedIlpTag,
VectorizeTag, InameImplementationTag,
InOrderSequentialSequentialTag,
filter_iname_tags_by_type)
tags = kernel.iname_tags_of_type(sched_item.iname,
InameImplementationTag)
tags = tuple(tag for tag in tags if tag)
from loopy.codegen.loop import (generate_unroll_loop,
generate_vectorize_loop,
generate_sequential_loop_dim_code)
if filter_iname_tags_by_type(tags, (UnrollTag, UnrolledIlpTag)):
func = generate_unroll_loop
elif filter_iname_tags_by_type(tags, VectorizeTag):
func = generate_vectorize_loop
elif not tags or filter_iname_tags_by_type(
tags, (LoopedIlpTag, ForceSequentialTag,
InOrderSequentialSequentialTag)):
func = generate_sequential_loop_dim_code
else:
raise RuntimeError(
"encountered (invalid) EnterLoop "
"for '%s', tagged '%s'" %
(sched_item.iname, ", ".join(str(tag) for tag in tags)))
return func(codegen_state, sched_index)
elif isinstance(sched_item, Barrier):
# {{{ emit barrier code
from loopy.codegen.result import CodeGenerationResult
if codegen_state.is_generating_device_code:
barrier_ast = codegen_state.ast_builder.emit_barrier(
sched_item.synchronization_kind, sched_item.mem_kind,
sched_item.comment)
if sched_item.originating_insn_id:
return CodeGenerationResult.new(
codegen_state, sched_item.originating_insn_id, barrier_ast,
codegen_state.implemented_domain)
else:
return barrier_ast
else:
# host code
if sched_item.synchronization_kind in ["global", "local"]:
# host code is assumed globally and locally synchronous
return CodeGenerationResult(
host_program=None,
device_programs=[],
implemented_domains={},
implemented_data_info=codegen_state.implemented_data_info)
else:
raise LoopyError("do not know how to emit code for barrier "
"synchronization kind '%s'"
"in host code" %
sched_item.synchronization_kind)
# }}}
elif isinstance(sched_item, RunInstruction):
insn = kernel.id_to_insn[sched_item.insn_id]
from loopy.codegen.instruction import generate_instruction_code
return codegen_state.try_vectorized(
"instruction %s" % insn.id,
lambda inner_cgs: generate_instruction_code(inner_cgs, insn))
else:
raise RuntimeError("unexpected schedule item type: %s" %
type(sched_item))
def emit_assignment(self, codegen_state, insn):
kernel = codegen_state.kernel
ecm = codegen_state.expression_to_code_mapper
assignee_var_name, = insn.assignee_var_names()
lhs_var = codegen_state.kernel.get_var_descriptor(assignee_var_name)
lhs_dtype = lhs_var.dtype
if insn.atomicity:
raise NotImplementedError("atomic ops in ISPC")
from loopy.expression import dtype_to_type_context
from pymbolic.mapper.stringifier import PREC_NONE
rhs_type_context = dtype_to_type_context(kernel.target, lhs_dtype)
rhs_code = ecm(insn.expression, prec=PREC_NONE,
type_context=rhs_type_context,
needed_dtype=lhs_dtype)
lhs = insn.assignee
# {{{ handle streaming stores
if "!streaming_store" in insn.tags:
ary = ecm.find_array(lhs)
from loopy.kernel.array import get_access_info
from pymbolic import evaluate
from loopy.symbolic import simplify_using_aff
index_tuple = tuple(
simplify_using_aff(kernel, idx) for idx in lhs.index_tuple)
access_info = get_access_info(kernel.target, ary, index_tuple,
lambda expr: evaluate(expr, codegen_state.var_subst_map),
codegen_state.vectorization_info)
from loopy.kernel.data import ArrayArg, TemporaryVariable
if not isinstance(ary, (ArrayArg, TemporaryVariable)):
raise LoopyError("array type not supported in ISPC: %s"
% type(ary).__name)
if len(access_info.subscripts) != 1:
raise LoopyError("streaming stores must have a subscript")
subscript, = access_info.subscripts
from pymbolic.primitives import Sum, flattened_sum, Variable
if isinstance(subscript, Sum):
terms = subscript.children
else:
terms = (subscript.children,)
new_terms = []
from loopy.kernel.data import LocalIndexTag, filter_iname_tags_by_type
from loopy.symbolic import get_dependencies
saw_l0 = False
for term in terms:
if (isinstance(term, Variable)
and kernel.iname_tags_of_type(term.name, LocalIndexTag)):
tag, = kernel.iname_tags_of_type(
term.name, LocalIndexTag, min_num=1, max_num=1)
if tag.axis == 0:
if saw_l0:
raise LoopyError(
"streaming store must have stride 1 in "
"local index, got: %s" % subscript)
saw_l0 = True
continue
else:
for dep in get_dependencies(term):
if filter_iname_tags_by_type(
kernel.iname_to_tags.get(dep, []), LocalIndexTag):
tag, = filter_iname_tags_by_type(
kernel.iname_to_tags.get(dep, []), LocalIndexTag, 1)
if tag.axis == 0:
raise LoopyError(
"streaming store must have stride 1 in "
"local index, got: %s" % subscript)
new_terms.append(term)
if not saw_l0:
raise LoopyError("streaming store must have stride 1 in "
"local index, got: %s" % subscript)
if access_info.vector_index is not None:
raise LoopyError("streaming store may not use a short-vector "
"data type")
rhs_has_programindex = any(
isinstance(tag, LocalIndexTag) and tag.axis == 0
for tag in kernel.iname_tags(dep)
for dep in get_dependencies(insn.expression))
if not rhs_has_programindex:
rhs_code = "broadcast(%s, 0)" % rhs_code
from cgen import Statement
return Statement(
"streaming_store(%s + %s, %s)"
% (
access_info.array_name,
ecm(flattened_sum(new_terms), PREC_NONE, 'i'),
rhs_code))
# }}}
from cgen import Assign
return Assign(ecm(lhs, prec=PREC_NONE, type_context=None), rhs_code)
def get_hw_axis_sizes_and_tags_for_save_slot(self, temporary):
"""
This is used for determining the amount of global storage needed for saving
and restoring the temporary across kernel calls, due to hardware
parallel inames (the inferred axes get prefixed to the number of
dimensions in the temporary).
In the case of local temporaries, inames that are tagged
hw-local do not contribute to the global storage shape.
"""
accessor_insn_ids = frozenset(
self.kernel.reader_map()[temporary.name]
| self.kernel.writer_map()[temporary.name])
group_tags = None
local_tags = None
def _sortedtags(tags):
return sorted(tags, key=lambda tag: tag.axis)
for insn_id in accessor_insn_ids:
insn = self.kernel.id_to_insn[insn_id]
my_group_tags = []
my_local_tags = []
for iname in insn.within_inames:
tags = self.kernel.iname_tags(iname)
if not tags:
continue
from loopy.kernel.data import (GroupIndexTag, LocalIndexTag,
ConcurrentTag, filter_iname_tags_by_type)
if filter_iname_tags_by_type(tags, GroupIndexTag):
tag, = filter_iname_tags_by_type(tags, GroupIndexTag, 1)
my_group_tags.append(tag)
elif filter_iname_tags_by_type(tags, LocalIndexTag):
tag, = filter_iname_tags_by_type(tags, LocalIndexTag, 1)
my_local_tags.append(tag)
elif filter_iname_tags_by_type(tags, ConcurrentTag):
raise LoopyError(
"iname '%s' is tagged with '%s' - only "
"group and local tags are supported for "
"auto save/reload of temporaries" %
(iname, tags))
if group_tags is None:
group_tags = _sortedtags(my_group_tags)
local_tags = _sortedtags(my_local_tags)
group_tags_originating_insn_id = insn_id
if (
group_tags != _sortedtags(my_group_tags)
or local_tags != _sortedtags(my_local_tags)):
raise LoopyError(
"inconsistent parallel tags across instructions that access "
"'%s' (specifically, instruction '%s' has tags '%s' but "
"instruction '%s' has tags '%s')"
% (temporary.name,
group_tags_originating_insn_id, group_tags + local_tags,
insn_id, my_group_tags + my_local_tags))
if group_tags is None:
assert local_tags is None
return (), ()
group_sizes, local_sizes = (
self.kernel.get_grid_sizes_for_insn_ids_as_exprs(accessor_insn_ids))
if temporary.address_space == lp.AddressSpace.LOCAL:
# Elide local axes in the save slot for local temporaries.
del local_tags[:]
local_sizes = ()
# We set hw_dims to be arranged according to the order:
# g.0 < g.1 < ... < l.0 < l.1 < ...
return (group_sizes + local_sizes), tuple(group_tags + local_tags)
def precompute(kernel, subst_use, sweep_inames=[], within=None,
storage_axes=None, temporary_name=None, precompute_inames=None,
precompute_outer_inames=None,
storage_axis_to_tag={},
# "None" is a valid value here, distinct from the default.
default_tag=_not_provided,
dtype=None,
fetch_bounding_box=False,
temporary_address_space=None,
compute_insn_id=None,
**kwargs):
"""Precompute the expression described in the substitution rule determined by
*subst_use* and store it in a temporary array. A precomputation needs two
things to operate, a list of *sweep_inames* (order irrelevant) and an
ordered list of *storage_axes* (whose order will describe the axis ordering
of the temporary array).
:arg subst_use: Describes what to prefetch.
The following objects may be given for *subst_use*:
* The name of the substitution rule.
* The tagged name ("name$tag") of the substitution rule.
* A list of invocations of the substitution rule.
This list of invocations, when swept across *sweep_inames*, then serves
to define the footprint of the precomputation.
Invocations may be tagged ("name$tag") to filter out a subset of the
usage sites of the substitution rule. (Namely those usage sites that
use the same tagged name.)
Invocations may be given as a string or as a
:class:`pymbolic.primitives.Expression` object.
If only one invocation is to be given, then the only entry of the list
may be given directly.
If the list of invocations generating the footprint is not given,
all (tag-matching, if desired) usage sites of the substitution rule
are used to determine the footprint.
The following cases can arise for each sweep axis:
* The axis is an iname that occurs within arguments specified at
usage sites of the substitution rule. This case is assumed covered
by the storage axes provided for the argument.
* The axis is an iname that occurs within the *value* of the rule, but not
within its arguments. A new, dedicated storage axis is allocated for
such an axis.
:arg sweep_inames: A :class:`list` of inames to be swept.
May also equivalently be a comma-separated string.
:arg within: a stack match as understood by
:func:`loopy.match.parse_stack_match`.
:arg storage_axes: A :class:`list` of inames and/or rule argument
names/indices to be used as storage axes.
May also equivalently be a comma-separated string.
:arg temporary_name:
The temporary variable name to use for storing the precomputed data.
If it does not exist, it will be created. If it does exist, its properties
(such as size, type) are checked (and updated, if possible) to match
its use.
:arg precompute_inames:
A tuple of inames to be used to carry out the precomputation.
If the specified inames do not already exist, they will be
created. If they do already exist, their loop domain is verified
against the one required for this precomputation. This tuple may
be shorter than the (provided or automatically found) *storage_axes*
tuple, in which case names will be automatically created.
May also equivalently be a comma-separated string.
:arg precompute_outer_inames: A :class:`frozenset` of inames within which
the compute instruction is nested. If *None*, make an educated guess.
May also be specified as a comma-separated string.
:arg default_tag: The :ref:`iname tag <iname-tags>` to be applied to the
inames created to perform the precomputation. The current default will
make them local axes and automatically split them to fit the work
group size, but this default will disappear in favor of simply leaving them
untagged in 2019. For 2018, a warning will be issued if no *default_tag* is
specified.
:arg compute_insn_id: The ID of the instruction generated to perform the
precomputation.
If `storage_axes` is not specified, it defaults to the arrangement
`<direct sweep axes><arguments>` with the direct sweep axes being the
slower-varying indices.
Trivial storage axes (i.e. axes of length 1 with respect to the sweep) are
eliminated.
"""
# {{{ unify temporary_address_space / temporary_scope
temporary_scope = kwargs.pop("temporary_scope", None)
from loopy.kernel.data import AddressSpace
if temporary_scope is not None:
from warnings import warn
warn("temporary_scope is deprecated. Use temporary_address_space instead",
DeprecationWarning, stacklevel=2)
if temporary_address_space is not None:
raise LoopyError("may not specify both temporary_address_space and "
"temporary_scope")
temporary_address_space = temporary_scope
del temporary_scope
# }}}
if kwargs:
raise TypeError("unrecognized keyword arguments: %s"
% ", ".join(kwargs.keys()))
# {{{ check, standardize arguments
if isinstance(sweep_inames, str):
sweep_inames = [iname.strip() for iname in sweep_inames.split(",")]
for iname in sweep_inames:
if iname not in kernel.all_inames():
raise RuntimeError("sweep iname '%s' is not a known iname"
% iname)
sweep_inames = list(sweep_inames)
sweep_inames_set = frozenset(sweep_inames)
if isinstance(storage_axes, str):
storage_axes = [ax.strip() for ax in storage_axes.split(",")]
if isinstance(precompute_inames, str):
precompute_inames = [iname.strip() for iname in precompute_inames.split(",")]
if isinstance(precompute_outer_inames, str):
precompute_outer_inames = frozenset(
iname.strip() for iname in precompute_outer_inames.split(","))
if isinstance(subst_use, str):
subst_use = [subst_use]
footprint_generators = None
subst_name = None
subst_tag = None
from pymbolic.primitives import Variable, Call
from loopy.symbolic import parse, TaggedVariable
for use in subst_use:
if isinstance(use, str):
use = parse(use)
if isinstance(use, Call):
if footprint_generators is None:
footprint_generators = []
footprint_generators.append(use)
subst_name_as_expr = use.function
else:
subst_name_as_expr = use
if isinstance(subst_name_as_expr, TaggedVariable):
new_subst_name = subst_name_as_expr.name
new_subst_tag = subst_name_as_expr.tag
elif isinstance(subst_name_as_expr, Variable):
new_subst_name = subst_name_as_expr.name
new_subst_tag = None
else:
raise ValueError("unexpected type of subst_name")
if (subst_name, subst_tag) == (None, None):
subst_name, subst_tag = new_subst_name, new_subst_tag
else:
if (subst_name, subst_tag) != (new_subst_name, new_subst_tag):
raise ValueError("not all uses in subst_use agree "
"on rule name and tag")
from loopy.match import parse_stack_match
within = parse_stack_match(within)
try:
subst = kernel.substitutions[subst_name]
except KeyError:
raise LoopyError("substitution rule '%s' not found"
% subst_name)
c_subst_name = subst_name.replace(".", "_")
# {{{ handle default_tag
from loopy.transform.data import _not_provided \
as transform_data_not_provided
if default_tag is _not_provided or default_tag is transform_data_not_provided:
# no need to warn for scalar precomputes
if sweep_inames:
from warnings import warn
warn(
"Not specifying default_tag is deprecated, and default_tag "
"will become mandatory in 2019.x. "
"Pass 'default_tag=\"l.auto\" to match the current default, "
"or Pass 'default_tag=None to leave the loops untagged, which "
"is the recommended behavior.",
DeprecationWarning, stacklevel=(
# In this case, we came here through add_prefetch. Increase
# the stacklevel.
3 if default_tag is transform_data_not_provided
else 2))
default_tag = "l.auto"
from loopy.kernel.data import parse_tag
default_tag = parse_tag(default_tag)
# }}}
# }}}
# {{{ process invocations in footprint generators, start access_descriptors
if footprint_generators:
from pymbolic.primitives import Variable, Call
access_descriptors = []
for fpg in footprint_generators:
if isinstance(fpg, Variable):
args = ()
elif isinstance(fpg, Call):
args = fpg.parameters
else:
raise ValueError("footprint generator must "
"be substitution rule invocation")
access_descriptors.append(
RuleAccessDescriptor(
identifier=access_descriptor_id(args, None),
args=args
))
# }}}
# {{{ gather up invocations in kernel code, finish access_descriptors
if not footprint_generators:
rule_mapping_context = SubstitutionRuleMappingContext(
kernel.substitutions, kernel.get_var_name_generator())
invg = RuleInvocationGatherer(
rule_mapping_context, kernel, subst_name, subst_tag, within)
del rule_mapping_context
import loopy as lp
for insn in kernel.instructions:
if isinstance(insn, lp.MultiAssignmentBase):
for assignee in insn.assignees:
invg(assignee, kernel, insn)
invg(insn.expression, kernel, insn)
access_descriptors = invg.access_descriptors
if not access_descriptors:
raise RuntimeError("no invocations of '%s' found" % subst_name)
# }}}
# {{{ find inames used in arguments
expanding_usage_arg_deps = set()
for accdesc in access_descriptors:
for arg in accdesc.args:
expanding_usage_arg_deps.update(
get_dependencies(arg) & kernel.all_inames())
# }}}
var_name_gen = kernel.get_var_name_generator()
# {{{ use given / find new storage_axes
# extra axes made necessary because they don't occur in the arguments
extra_storage_axes = set(sweep_inames_set - expanding_usage_arg_deps)
from loopy.symbolic import SubstitutionRuleExpander
submap = SubstitutionRuleExpander(kernel.substitutions)
value_inames = (
get_dependencies(submap(subst.expression))
- frozenset(subst.arguments)
) & kernel.all_inames()
if value_inames - expanding_usage_arg_deps < extra_storage_axes:
raise RuntimeError("unreferenced sweep inames specified: "
+ ", ".join(extra_storage_axes
- value_inames - expanding_usage_arg_deps))
new_iname_to_tag = {}
if storage_axes is None:
storage_axes = []
# Add sweep_inames (in given--rather than arbitrary--order) to
# storage_axes *if* they are part of extra_storage_axes.
for iname in sweep_inames:
if iname in extra_storage_axes:
extra_storage_axes.remove(iname)
storage_axes.append(iname)
if extra_storage_axes:
if (precompute_inames is not None
and len(storage_axes) < len(precompute_inames)):
raise LoopyError("must specify a sufficient number of "
"storage_axes to uniquely determine the meaning "
"of the given precompute_inames. (%d storage_axes "
"needed)" % len(precompute_inames))
storage_axes.extend(sorted(extra_storage_axes))
storage_axes.extend(range(len(subst.arguments)))
del extra_storage_axes
prior_storage_axis_name_dict = {}
storage_axis_names = []
storage_axis_sources = [] # number for arg#, or iname
# {{{ check for pre-existing precompute_inames
if precompute_inames is not None:
preexisting_precompute_inames = (
set(precompute_inames) & kernel.all_inames())
else:
preexisting_precompute_inames = set()
# }}}
for i, saxis in enumerate(storage_axes):
tag_lookup_saxis = saxis
if saxis in subst.arguments:
saxis = subst.arguments.index(saxis)
storage_axis_sources.append(saxis)
if isinstance(saxis, int):
# argument index
name = old_name = subst.arguments[saxis]
else:
old_name = saxis
name = "%s_%s" % (c_subst_name, old_name)
if (precompute_inames is not None
and i < len(precompute_inames)
and precompute_inames[i]):
name = precompute_inames[i]
tag_lookup_saxis = name
if (name not in preexisting_precompute_inames
and var_name_gen.is_name_conflicting(name)):
raise RuntimeError("new storage axis name '%s' "
"conflicts with existing name" % name)
else:
name = var_name_gen(name)
storage_axis_names.append(name)
if name not in preexisting_precompute_inames:
new_iname_to_tag[name] = storage_axis_to_tag.get(
tag_lookup_saxis, default_tag)
prior_storage_axis_name_dict[name] = old_name
del storage_axis_to_tag
del storage_axes
del precompute_inames
# }}}
# {{{ fill out access_descriptors[...].storage_axis_exprs
access_descriptors = [
accdesc.copy(
storage_axis_exprs=storage_axis_exprs(
storage_axis_sources, accdesc.args))
for accdesc in access_descriptors]
# }}}
expanding_inames = sweep_inames_set | frozenset(expanding_usage_arg_deps)
assert expanding_inames <= kernel.all_inames()
if storage_axis_names:
# {{{ find domain to be changed
change_inames = expanding_inames | preexisting_precompute_inames
from loopy.kernel.tools import DomainChanger
domch = DomainChanger(kernel, change_inames)
if domch.leaf_domain_index is not None:
# If the sweep inames are at home in parent domains, then we'll add
# fetches with loops over copies of these parent inames that will end
# up being scheduled *within* loops over these parents.
for iname in sweep_inames_set:
if kernel.get_home_domain_index(iname) != domch.leaf_domain_index:
raise RuntimeError("sweep iname '%s' is not 'at home' in the "
"sweep's leaf domain" % iname)
# }}}
abm = ArrayToBufferMap(kernel, domch.domain, sweep_inames,
access_descriptors, len(storage_axis_names))
non1_storage_axis_names = []
for i, saxis in enumerate(storage_axis_names):
if abm.non1_storage_axis_flags[i]:
non1_storage_axis_names.append(saxis)
else:
del new_iname_to_tag[saxis]
if saxis in preexisting_precompute_inames:
raise LoopyError("precompute axis %d (1-based) was "
"eliminated as "
"having length 1 but also mapped to existing "
"iname '%s'" % (i+1, saxis))
mod_domain = domch.domain
# {{{ modify the domain, taking into account preexisting inames
# inames may already exist in mod_domain, add them primed to start
primed_non1_saxis_names = [
iname+"'" for iname in non1_storage_axis_names]
mod_domain = abm.augment_domain_with_sweep(
domch.domain, primed_non1_saxis_names,
boxify_sweep=fetch_bounding_box)
check_domain = mod_domain
for i, saxis in enumerate(non1_storage_axis_names):
var_dict = mod_domain.get_var_dict(isl.dim_type.set)
if saxis in preexisting_precompute_inames:
# add equality constraint between existing and new variable
dt, dim_idx = var_dict[saxis]
saxis_aff = isl.Aff.var_on_domain(mod_domain.space, dt, dim_idx)
dt, dim_idx = var_dict[primed_non1_saxis_names[i]]
new_var_aff = isl.Aff.var_on_domain(mod_domain.space, dt, dim_idx)
mod_domain = mod_domain.add_constraint(
isl.Constraint.equality_from_aff(new_var_aff - saxis_aff))
# project out the new one
mod_domain = mod_domain.project_out(dt, dim_idx, 1)
else:
# remove the prime from the new variable
dt, dim_idx = var_dict[primed_non1_saxis_names[i]]
mod_domain = mod_domain.set_dim_name(dt, dim_idx, saxis)
def add_assumptions(d):
assumption_non_param = isl.BasicSet.from_params(kernel.assumptions)
assumptions, domain = isl.align_two(assumption_non_param, d)
return assumptions & domain
# {{{ check that we got the desired domain
check_domain = add_assumptions(
check_domain.project_out_except(
primed_non1_saxis_names, [isl.dim_type.set]))
mod_check_domain = add_assumptions(mod_domain)
# re-add the prime from the new variable
var_dict = mod_check_domain.get_var_dict(isl.dim_type.set)
for saxis in non1_storage_axis_names:
dt, dim_idx = var_dict[saxis]
mod_check_domain = mod_check_domain.set_dim_name(dt, dim_idx, saxis+"'")
mod_check_domain = mod_check_domain.project_out_except(
primed_non1_saxis_names, [isl.dim_type.set])
mod_check_domain, check_domain = isl.align_two(
mod_check_domain, check_domain)
# The modified domain can't get bigger by adding constraints
assert mod_check_domain <= check_domain
if not check_domain <= mod_check_domain:
print(check_domain)
print(mod_check_domain)
raise LoopyError("domain of preexisting inames does not match "
"domain needed for precompute")
# }}}
# {{{ check that we didn't shrink the original domain
# project out the new names from the modified domain
orig_domain_inames = list(domch.domain.get_var_dict(isl.dim_type.set))
mod_check_domain = add_assumptions(
mod_domain.project_out_except(
orig_domain_inames, [isl.dim_type.set]))
check_domain = add_assumptions(domch.domain)
mod_check_domain, check_domain = isl.align_two(
mod_check_domain, check_domain)
# The modified domain can't get bigger by adding constraints
assert mod_check_domain <= check_domain
if not check_domain <= mod_check_domain:
print(check_domain)
print(mod_check_domain)
raise LoopyError("original domain got shrunk by applying the precompute")
# }}}
# }}}
new_kernel_domains = domch.get_domains_with(mod_domain)
else:
# leave kernel domains unchanged
new_kernel_domains = kernel.domains
non1_storage_axis_names = []
abm = NoOpArrayToBufferMap()
kernel = kernel.copy(domains=new_kernel_domains)
# {{{ set up compute insn
if temporary_name is None:
temporary_name = var_name_gen(based_on=c_subst_name)
assignee = var(temporary_name)
if non1_storage_axis_names:
assignee = assignee[
tuple(var(iname) for iname in non1_storage_axis_names)]
# {{{ process substitutions on compute instruction
storage_axis_subst_dict = {}
for arg_name, bi in zip(storage_axis_names, abm.storage_base_indices):
if arg_name in non1_storage_axis_names:
arg = var(arg_name)
else:
arg = 0
storage_axis_subst_dict[
prior_storage_axis_name_dict.get(arg_name, arg_name)] = arg+bi
rule_mapping_context = SubstitutionRuleMappingContext(
kernel.substitutions, kernel.get_var_name_generator())
from loopy.match import parse_stack_match
expr_subst_map = RuleAwareSubstitutionMapper(
rule_mapping_context,
make_subst_func(storage_axis_subst_dict),
within=parse_stack_match(None))
compute_expression = expr_subst_map(subst.expression, kernel, None)
# }}}
from loopy.kernel.data import Assignment
if compute_insn_id is None:
compute_insn_id = kernel.make_unique_instruction_id(based_on=c_subst_name)
compute_insn = Assignment(
id=compute_insn_id,
assignee=assignee,
expression=compute_expression,
# within_inames determined below
)
compute_dep_id = compute_insn_id
added_compute_insns = [compute_insn]
if temporary_address_space == AddressSpace.GLOBAL:
barrier_insn_id = kernel.make_unique_instruction_id(
based_on=c_subst_name+"_barrier")
from loopy.kernel.instruction import BarrierInstruction
barrier_insn = BarrierInstruction(
id=barrier_insn_id,
depends_on=frozenset([compute_insn_id]),
synchronization_kind="global",
mem_kind="global")
compute_dep_id = barrier_insn_id
added_compute_insns.append(barrier_insn)
# }}}
# {{{ substitute rule into expressions in kernel (if within footprint)
from loopy.symbolic import SubstitutionRuleExpander
expander = SubstitutionRuleExpander(kernel.substitutions)
invr = RuleInvocationReplacer(rule_mapping_context,
subst_name, subst_tag, within,
access_descriptors, abm,
storage_axis_names, storage_axis_sources,
non1_storage_axis_names,
temporary_name, compute_insn_id, compute_dep_id,
compute_read_variables=get_dependencies(expander(compute_expression)))
kernel = invr.map_kernel(kernel)
kernel = kernel.copy(
instructions=added_compute_insns + kernel.instructions)
kernel = rule_mapping_context.finish_kernel(kernel)
# }}}
# {{{ add dependencies to compute insn
kernel = kernel.copy(
instructions=[
insn.copy(depends_on=frozenset(invr.compute_insn_depends_on))
if insn.id == compute_insn_id
else insn
for insn in kernel.instructions])
# }}}
# {{{ propagate storage iname subst to dependencies of compute instructions
from loopy.kernel.tools import find_recursive_dependencies
compute_deps = find_recursive_dependencies(
kernel, frozenset([compute_insn_id]))
# FIXME: Need to verify that there are no outside dependencies
# on compute_deps
prior_storage_axis_names = frozenset(storage_axis_subst_dict)
new_insns = []
for insn in kernel.instructions:
if (insn.id in compute_deps
and insn.within_inames & prior_storage_axis_names):
insn = (insn
.with_transformed_expressions(
lambda expr: expr_subst_map(expr, kernel, insn))
.copy(within_inames=frozenset(
storage_axis_subst_dict.get(iname, var(iname)).name
for iname in insn.within_inames)))
new_insns.append(insn)
else:
new_insns.append(insn)
kernel = kernel.copy(instructions=new_insns)
# }}}
# {{{ determine inames for compute insn
if precompute_outer_inames is None:
from loopy.kernel.tools import guess_iname_deps_based_on_var_use
precompute_outer_inames = (
frozenset(non1_storage_axis_names)
| frozenset(
(expanding_usage_arg_deps | value_inames)
- sweep_inames_set)
| guess_iname_deps_based_on_var_use(kernel, compute_insn))
else:
if not isinstance(precompute_outer_inames, frozenset):
raise TypeError("precompute_outer_inames must be a frozenset")
precompute_outer_inames = precompute_outer_inames \
| frozenset(non1_storage_axis_names)
kernel = kernel.copy(
instructions=[
insn.copy(within_inames=precompute_outer_inames)
if insn.id == compute_insn_id
else insn
for insn in kernel.instructions])
# }}}
# {{{ set up temp variable
import loopy as lp
if dtype is None:
dtype = lp.auto
else:
dtype = np.dtype(dtype)
import loopy as lp
if temporary_address_space is None:
temporary_address_space = lp.auto
new_temp_shape = tuple(abm.non1_storage_shape)
new_temporary_variables = kernel.temporary_variables.copy()
if temporary_name not in new_temporary_variables:
temp_var = lp.TemporaryVariable(
name=temporary_name,
dtype=dtype,
base_indices=(0,)*len(new_temp_shape),
shape=tuple(abm.non1_storage_shape),
address_space=temporary_address_space,
dim_names=tuple(non1_storage_axis_names))
else:
temp_var = new_temporary_variables[temporary_name]
# {{{ check and adapt existing temporary
if temp_var.dtype is lp.auto:
pass
elif temp_var.dtype is not lp.auto and dtype is lp.auto:
dtype = temp_var.dtype
elif temp_var.dtype is not lp.auto and dtype is not lp.auto:
if temp_var.dtype != dtype:
raise LoopyError("Existing and new dtype of temporary '%s' "
"do not match (existing: %s, new: %s)"
% (temporary_name, temp_var.dtype, dtype))
temp_var = temp_var.copy(dtype=dtype)
if len(temp_var.shape) != len(new_temp_shape):
raise LoopyError("Existing and new temporary '%s' do not "
"have matching number of dimensions ('%d' vs. '%d') "
% (temporary_name,
len(temp_var.shape), len(new_temp_shape)))
if temp_var.base_indices != (0,) * len(new_temp_shape):
raise LoopyError("Existing and new temporary '%s' do not "
"have matching number of dimensions ('%d' vs. '%d') "
% (temporary_name,
len(temp_var.shape), len(new_temp_shape)))
new_temp_shape = tuple(
max(i, ex_i)
for i, ex_i in zip(new_temp_shape, temp_var.shape))
temp_var = temp_var.copy(shape=new_temp_shape)
if temporary_address_space == temp_var.address_space:
pass
elif temporary_address_space is lp.auto:
temporary_address_space = temp_var.address_space
elif temp_var.address_space is lp.auto:
pass
else:
raise LoopyError("Existing and new temporary '%s' do not "
"have matching scopes (existing: %s, new: %s)"
% (temporary_name,
AddressSpace.stringify(temp_var.address_space),
AddressSpace.stringify(temporary_address_space)))
temp_var = temp_var.copy(address_space=temporary_address_space)
# }}}
new_temporary_variables[temporary_name] = temp_var
kernel = kernel.copy(
temporary_variables=new_temporary_variables)
# }}}
from loopy import tag_inames
kernel = tag_inames(kernel, new_iname_to_tag)
from loopy.kernel.data import AutoFitLocalIndexTag, filter_iname_tags_by_type
if filter_iname_tags_by_type(new_iname_to_tag.values(), AutoFitLocalIndexTag):
from loopy.kernel.tools import assign_automatic_axes
kernel = assign_automatic_axes(kernel)
return kernel
def precompute(
kernel,
subst_use,
sweep_inames=[],
within=None,
storage_axes=None,
temporary_name=None,
precompute_inames=None,
precompute_outer_inames=None,
storage_axis_to_tag={},
# "None" is a valid value here, distinct from the default.
default_tag=_not_provided,
dtype=None,
fetch_bounding_box=False,
temporary_address_space=None,
compute_insn_id=None,
**kwargs):
"""Precompute the expression described in the substitution rule determined by
*subst_use* and store it in a temporary array. A precomputation needs two
things to operate, a list of *sweep_inames* (order irrelevant) and an
ordered list of *storage_axes* (whose order will describe the axis ordering
of the temporary array).
:arg subst_use: Describes what to prefetch.
The following objects may be given for *subst_use*:
* The name of the substitution rule.
* The tagged name ("name$tag") of the substitution rule.
* A list of invocations of the substitution rule.
This list of invocations, when swept across *sweep_inames*, then serves
to define the footprint of the precomputation.
Invocations may be tagged ("name$tag") to filter out a subset of the
usage sites of the substitution rule. (Namely those usage sites that
use the same tagged name.)
Invocations may be given as a string or as a
:class:`pymbolic.primitives.Expression` object.
If only one invocation is to be given, then the only entry of the list
may be given directly.
If the list of invocations generating the footprint is not given,
all (tag-matching, if desired) usage sites of the substitution rule
are used to determine the footprint.
The following cases can arise for each sweep axis:
* The axis is an iname that occurs within arguments specified at
usage sites of the substitution rule. This case is assumed covered
by the storage axes provided for the argument.
* The axis is an iname that occurs within the *value* of the rule, but not
within its arguments. A new, dedicated storage axis is allocated for
such an axis.
:arg sweep_inames: A :class:`list` of inames to be swept.
May also equivalently be a comma-separated string.
:arg within: a stack match as understood by
:func:`loopy.match.parse_stack_match`.
:arg storage_axes: A :class:`list` of inames and/or rule argument
names/indices to be used as storage axes.
May also equivalently be a comma-separated string.
:arg temporary_name:
The temporary variable name to use for storing the precomputed data.
If it does not exist, it will be created. If it does exist, its properties
(such as size, type) are checked (and updated, if possible) to match
its use.
:arg precompute_inames:
A tuple of inames to be used to carry out the precomputation.
If the specified inames do not already exist, they will be
created. If they do already exist, their loop domain is verified
against the one required for this precomputation. This tuple may
be shorter than the (provided or automatically found) *storage_axes*
tuple, in which case names will be automatically created.
May also equivalently be a comma-separated string.
:arg precompute_outer_inames: A :class:`frozenset` of inames within which
the compute instruction is nested. If *None*, make an educated guess.
May also be specified as a comma-separated string.
:arg default_tag: The :ref:`iname tag <iname-tags>` to be applied to the
inames created to perform the precomputation. The current default will
make them local axes and automatically split them to fit the work
group size, but this default will disappear in favor of simply leaving them
untagged in 2019. For 2018, a warning will be issued if no *default_tag* is
specified.
:arg compute_insn_id: The ID of the instruction generated to perform the
precomputation.
If `storage_axes` is not specified, it defaults to the arrangement
`<direct sweep axes><arguments>` with the direct sweep axes being the
slower-varying indices.
Trivial storage axes (i.e. axes of length 1 with respect to the sweep) are
eliminated.
"""
# {{{ unify temporary_address_space / temporary_scope
temporary_scope = kwargs.pop("temporary_scope", None)
from loopy.kernel.data import AddressSpace
if temporary_scope is not None:
from warnings import warn
warn(
"temporary_scope is deprecated. Use temporary_address_space instead",
DeprecationWarning,
stacklevel=2)
if temporary_address_space is not None:
raise LoopyError(
"may not specify both temporary_address_space and "
"temporary_scope")
temporary_address_space = temporary_scope
del temporary_scope
# }}}
if kwargs:
raise TypeError("unrecognized keyword arguments: %s" %
", ".join(kwargs.keys()))
# {{{ check, standardize arguments
if isinstance(sweep_inames, str):
sweep_inames = [iname.strip() for iname in sweep_inames.split(",")]
for iname in sweep_inames:
if iname not in kernel.all_inames():
raise RuntimeError("sweep iname '%s' is not a known iname" % iname)
sweep_inames = list(sweep_inames)
sweep_inames_set = frozenset(sweep_inames)
if isinstance(storage_axes, str):
storage_axes = [ax.strip() for ax in storage_axes.split(",")]
if isinstance(precompute_inames, str):
precompute_inames = [
iname.strip() for iname in precompute_inames.split(",")
]
if isinstance(precompute_outer_inames, str):
precompute_outer_inames = frozenset(
iname.strip() for iname in precompute_outer_inames.split(","))
if isinstance(subst_use, str):
subst_use = [subst_use]
footprint_generators = None
subst_name = None
subst_tag = None
from pymbolic.primitives import Variable, Call
from loopy.symbolic import parse, TaggedVariable
for use in subst_use:
if isinstance(use, str):
use = parse(use)
if isinstance(use, Call):
if footprint_generators is None:
footprint_generators = []
footprint_generators.append(use)
subst_name_as_expr = use.function
else:
subst_name_as_expr = use
if isinstance(subst_name_as_expr, TaggedVariable):
new_subst_name = subst_name_as_expr.name
new_subst_tag = subst_name_as_expr.tag
elif isinstance(subst_name_as_expr, Variable):
new_subst_name = subst_name_as_expr.name
new_subst_tag = None
else:
raise ValueError("unexpected type of subst_name")
if (subst_name, subst_tag) == (None, None):
subst_name, subst_tag = new_subst_name, new_subst_tag
else:
if (subst_name, subst_tag) != (new_subst_name, new_subst_tag):
raise ValueError("not all uses in subst_use agree "
"on rule name and tag")
from loopy.match import parse_stack_match
within = parse_stack_match(within)
try:
subst = kernel.substitutions[subst_name]
except KeyError:
raise LoopyError("substitution rule '%s' not found" % subst_name)
c_subst_name = subst_name.replace(".", "_")
# {{{ handle default_tag
from loopy.transform.data import _not_provided \
as transform_data_not_provided
if default_tag is _not_provided or default_tag is transform_data_not_provided:
# no need to warn for scalar precomputes
if sweep_inames:
from warnings import warn
warn(
"Not specifying default_tag is deprecated, and default_tag "
"will become mandatory in 2019.x. "
"Pass 'default_tag=\"l.auto\" to match the current default, "
"or Pass 'default_tag=None to leave the loops untagged, which "
"is the recommended behavior.",
DeprecationWarning,
stacklevel=(
# In this case, we came here through add_prefetch. Increase
# the stacklevel.
3 if default_tag is transform_data_not_provided else 2))
default_tag = "l.auto"
from loopy.kernel.data import parse_tag
default_tag = parse_tag(default_tag)
# }}}
# }}}
# {{{ process invocations in footprint generators, start access_descriptors
if footprint_generators:
from pymbolic.primitives import Variable, Call
access_descriptors = []
for fpg in footprint_generators:
if isinstance(fpg, Variable):
args = ()
elif isinstance(fpg, Call):
args = fpg.parameters
else:
raise ValueError("footprint generator must "
"be substitution rule invocation")
access_descriptors.append(
RuleAccessDescriptor(identifier=access_descriptor_id(
args, None),
args=args))
# }}}
# {{{ gather up invocations in kernel code, finish access_descriptors
if not footprint_generators:
rule_mapping_context = SubstitutionRuleMappingContext(
kernel.substitutions, kernel.get_var_name_generator())
invg = RuleInvocationGatherer(rule_mapping_context, kernel, subst_name,
subst_tag, within)
del rule_mapping_context
import loopy as lp
for insn in kernel.instructions:
if isinstance(insn, lp.MultiAssignmentBase):
for assignee in insn.assignees:
invg(assignee, kernel, insn)
invg(insn.expression, kernel, insn)
access_descriptors = invg.access_descriptors
if not access_descriptors:
raise RuntimeError("no invocations of '%s' found" % subst_name)
# }}}
# {{{ find inames used in arguments
expanding_usage_arg_deps = set()
for accdesc in access_descriptors:
for arg in accdesc.args:
expanding_usage_arg_deps.update(
get_dependencies(arg) & kernel.all_inames())
# }}}
var_name_gen = kernel.get_var_name_generator()
# {{{ use given / find new storage_axes
# extra axes made necessary because they don't occur in the arguments
extra_storage_axes = set(sweep_inames_set - expanding_usage_arg_deps)
from loopy.symbolic import SubstitutionRuleExpander
submap = SubstitutionRuleExpander(kernel.substitutions)
value_inames = (get_dependencies(submap(subst.expression)) -
frozenset(subst.arguments)) & kernel.all_inames()
if value_inames - expanding_usage_arg_deps < extra_storage_axes:
raise RuntimeError("unreferenced sweep inames specified: " +
", ".join(extra_storage_axes - value_inames -
expanding_usage_arg_deps))
new_iname_to_tag = {}
if storage_axes is None:
storage_axes = []
# Add sweep_inames (in given--rather than arbitrary--order) to
# storage_axes *if* they are part of extra_storage_axes.
for iname in sweep_inames:
if iname in extra_storage_axes:
extra_storage_axes.remove(iname)
storage_axes.append(iname)
if extra_storage_axes:
if (precompute_inames is not None
and len(storage_axes) < len(precompute_inames)):
raise LoopyError(
"must specify a sufficient number of "
"storage_axes to uniquely determine the meaning "
"of the given precompute_inames. (%d storage_axes "
"needed)" % len(precompute_inames))
storage_axes.extend(sorted(extra_storage_axes))
storage_axes.extend(range(len(subst.arguments)))
del extra_storage_axes
prior_storage_axis_name_dict = {}
storage_axis_names = []
storage_axis_sources = [] # number for arg#, or iname
# {{{ check for pre-existing precompute_inames
if precompute_inames is not None:
preexisting_precompute_inames = (set(precompute_inames)
& kernel.all_inames())
else:
preexisting_precompute_inames = set()
# }}}
for i, saxis in enumerate(storage_axes):
tag_lookup_saxis = saxis
if saxis in subst.arguments:
saxis = subst.arguments.index(saxis)
storage_axis_sources.append(saxis)
if isinstance(saxis, int):
# argument index
name = old_name = subst.arguments[saxis]
else:
old_name = saxis
name = "%s_%s" % (c_subst_name, old_name)
if (precompute_inames is not None and i < len(precompute_inames)
and precompute_inames[i]):
name = precompute_inames[i]
tag_lookup_saxis = name
if (name not in preexisting_precompute_inames
and var_name_gen.is_name_conflicting(name)):
raise RuntimeError("new storage axis name '%s' "
"conflicts with existing name" % name)
else:
name = var_name_gen(name)
storage_axis_names.append(name)
if name not in preexisting_precompute_inames:
new_iname_to_tag[name] = storage_axis_to_tag.get(
tag_lookup_saxis, default_tag)
prior_storage_axis_name_dict[name] = old_name
del storage_axis_to_tag
del storage_axes
del precompute_inames
# }}}
# {{{ fill out access_descriptors[...].storage_axis_exprs
access_descriptors = [
accdesc.copy(storage_axis_exprs=storage_axis_exprs(
storage_axis_sources, accdesc.args))
for accdesc in access_descriptors
]
# }}}
expanding_inames = sweep_inames_set | frozenset(expanding_usage_arg_deps)
assert expanding_inames <= kernel.all_inames()
if storage_axis_names:
# {{{ find domain to be changed
change_inames = expanding_inames | preexisting_precompute_inames
from loopy.kernel.tools import DomainChanger
domch = DomainChanger(kernel, change_inames)
if domch.leaf_domain_index is not None:
# If the sweep inames are at home in parent domains, then we'll add
# fetches with loops over copies of these parent inames that will end
# up being scheduled *within* loops over these parents.
for iname in sweep_inames_set:
if kernel.get_home_domain_index(
iname) != domch.leaf_domain_index:
raise RuntimeError(
"sweep iname '%s' is not 'at home' in the "
"sweep's leaf domain" % iname)
# }}}
abm = ArrayToBufferMap(kernel, domch.domain, sweep_inames,
access_descriptors, len(storage_axis_names))
non1_storage_axis_names = []
for i, saxis in enumerate(storage_axis_names):
if abm.non1_storage_axis_flags[i]:
non1_storage_axis_names.append(saxis)
else:
del new_iname_to_tag[saxis]
if saxis in preexisting_precompute_inames:
raise LoopyError(
"precompute axis %d (1-based) was "
"eliminated as "
"having length 1 but also mapped to existing "
"iname '%s'" % (i + 1, saxis))
mod_domain = domch.domain
# {{{ modify the domain, taking into account preexisting inames
# inames may already exist in mod_domain, add them primed to start
primed_non1_saxis_names = [
iname + "'" for iname in non1_storage_axis_names
]
mod_domain = abm.augment_domain_with_sweep(
domch.domain,
primed_non1_saxis_names,
boxify_sweep=fetch_bounding_box)
check_domain = mod_domain
for i, saxis in enumerate(non1_storage_axis_names):
var_dict = mod_domain.get_var_dict(isl.dim_type.set)
if saxis in preexisting_precompute_inames:
# add equality constraint between existing and new variable
dt, dim_idx = var_dict[saxis]
saxis_aff = isl.Aff.var_on_domain(mod_domain.space, dt,
dim_idx)
dt, dim_idx = var_dict[primed_non1_saxis_names[i]]
new_var_aff = isl.Aff.var_on_domain(mod_domain.space, dt,
dim_idx)
mod_domain = mod_domain.add_constraint(
isl.Constraint.equality_from_aff(new_var_aff - saxis_aff))
# project out the new one
mod_domain = mod_domain.project_out(dt, dim_idx, 1)
else:
# remove the prime from the new variable
dt, dim_idx = var_dict[primed_non1_saxis_names[i]]
mod_domain = mod_domain.set_dim_name(dt, dim_idx, saxis)
def add_assumptions(d):
assumption_non_param = isl.BasicSet.from_params(kernel.assumptions)
assumptions, domain = isl.align_two(assumption_non_param, d)
return assumptions & domain
# {{{ check that we got the desired domain
check_domain = add_assumptions(
check_domain.project_out_except(primed_non1_saxis_names,
[isl.dim_type.set]))
mod_check_domain = add_assumptions(mod_domain)
# re-add the prime from the new variable
var_dict = mod_check_domain.get_var_dict(isl.dim_type.set)
for saxis in non1_storage_axis_names:
dt, dim_idx = var_dict[saxis]
mod_check_domain = mod_check_domain.set_dim_name(
dt, dim_idx, saxis + "'")
mod_check_domain = mod_check_domain.project_out_except(
primed_non1_saxis_names, [isl.dim_type.set])
mod_check_domain, check_domain = isl.align_two(mod_check_domain,
check_domain)
# The modified domain can't get bigger by adding constraints
assert mod_check_domain <= check_domain
if not check_domain <= mod_check_domain:
print(check_domain)
print(mod_check_domain)
raise LoopyError("domain of preexisting inames does not match "
"domain needed for precompute")
# }}}
# {{{ check that we didn't shrink the original domain
# project out the new names from the modified domain
orig_domain_inames = list(domch.domain.get_var_dict(isl.dim_type.set))
mod_check_domain = add_assumptions(
mod_domain.project_out_except(orig_domain_inames,
[isl.dim_type.set]))
check_domain = add_assumptions(domch.domain)
mod_check_domain, check_domain = isl.align_two(mod_check_domain,
check_domain)
# The modified domain can't get bigger by adding constraints
assert mod_check_domain <= check_domain
if not check_domain <= mod_check_domain:
print(check_domain)
print(mod_check_domain)
raise LoopyError(
"original domain got shrunk by applying the precompute")
# }}}
# }}}
new_kernel_domains = domch.get_domains_with(mod_domain)
else:
# leave kernel domains unchanged
new_kernel_domains = kernel.domains
non1_storage_axis_names = []
abm = NoOpArrayToBufferMap()
kernel = kernel.copy(domains=new_kernel_domains)
# {{{ set up compute insn
if temporary_name is None:
temporary_name = var_name_gen(based_on=c_subst_name)
assignee = var(temporary_name)
if non1_storage_axis_names:
assignee = assignee[tuple(
var(iname) for iname in non1_storage_axis_names)]
# {{{ process substitutions on compute instruction
storage_axis_subst_dict = {}
for arg_name, bi in zip(storage_axis_names, abm.storage_base_indices):
if arg_name in non1_storage_axis_names:
arg = var(arg_name)
else:
arg = 0
storage_axis_subst_dict[prior_storage_axis_name_dict.get(
arg_name, arg_name)] = arg + bi
rule_mapping_context = SubstitutionRuleMappingContext(
kernel.substitutions, kernel.get_var_name_generator())
from loopy.match import parse_stack_match
expr_subst_map = RuleAwareSubstitutionMapper(
rule_mapping_context,
make_subst_func(storage_axis_subst_dict),
within=parse_stack_match(None))
compute_expression = expr_subst_map(subst.expression, kernel, None)
# }}}
from loopy.kernel.data import Assignment
if compute_insn_id is None:
compute_insn_id = kernel.make_unique_instruction_id(
based_on=c_subst_name)
compute_insn = Assignment(
id=compute_insn_id,
assignee=assignee,
expression=compute_expression,
# within_inames determined below
)
compute_dep_id = compute_insn_id
added_compute_insns = [compute_insn]
if temporary_address_space == AddressSpace.GLOBAL:
barrier_insn_id = kernel.make_unique_instruction_id(
based_on=c_subst_name + "_barrier")
from loopy.kernel.instruction import BarrierInstruction
barrier_insn = BarrierInstruction(id=barrier_insn_id,
depends_on=frozenset(
[compute_insn_id]),
synchronization_kind="global",
mem_kind="global")
compute_dep_id = barrier_insn_id
added_compute_insns.append(barrier_insn)
# }}}
# {{{ substitute rule into expressions in kernel (if within footprint)
from loopy.symbolic import SubstitutionRuleExpander
expander = SubstitutionRuleExpander(kernel.substitutions)
invr = RuleInvocationReplacer(rule_mapping_context,
subst_name,
subst_tag,
within,
access_descriptors,
abm,
storage_axis_names,
storage_axis_sources,
non1_storage_axis_names,
temporary_name,
compute_insn_id,
compute_dep_id,
compute_read_variables=get_dependencies(
expander(compute_expression)))
kernel = invr.map_kernel(kernel)
kernel = kernel.copy(instructions=added_compute_insns +
kernel.instructions)
kernel = rule_mapping_context.finish_kernel(kernel)
# }}}
# {{{ add dependencies to compute insn
kernel = kernel.copy(instructions=[
insn.copy(depends_on=frozenset(invr.compute_insn_depends_on)) if insn.
id == compute_insn_id else insn for insn in kernel.instructions
])
# }}}
# {{{ propagate storage iname subst to dependencies of compute instructions
from loopy.kernel.tools import find_recursive_dependencies
compute_deps = find_recursive_dependencies(kernel,
frozenset([compute_insn_id]))
# FIXME: Need to verify that there are no outside dependencies
# on compute_deps
prior_storage_axis_names = frozenset(storage_axis_subst_dict)
new_insns = []
for insn in kernel.instructions:
if (insn.id in compute_deps
and insn.within_inames & prior_storage_axis_names):
insn = (insn.with_transformed_expressions(
lambda expr: expr_subst_map(expr, kernel, insn)).copy(
within_inames=frozenset(
storage_axis_subst_dict.get(iname, var(iname)).name
for iname in insn.within_inames)))
new_insns.append(insn)
else:
new_insns.append(insn)
kernel = kernel.copy(instructions=new_insns)
# }}}
# {{{ determine inames for compute insn
if precompute_outer_inames is None:
from loopy.kernel.tools import guess_iname_deps_based_on_var_use
precompute_outer_inames = (
frozenset(non1_storage_axis_names)
| frozenset((expanding_usage_arg_deps | value_inames) -
sweep_inames_set)
| guess_iname_deps_based_on_var_use(kernel, compute_insn))
else:
if not isinstance(precompute_outer_inames, frozenset):
raise TypeError("precompute_outer_inames must be a frozenset")
precompute_outer_inames = precompute_outer_inames \
| frozenset(non1_storage_axis_names)
kernel = kernel.copy(instructions=[
insn.copy(within_inames=precompute_outer_inames) if insn.id ==
compute_insn_id else insn for insn in kernel.instructions
])
# }}}
# {{{ set up temp variable
import loopy as lp
if dtype is not None:
dtype = np.dtype(dtype)
if temporary_address_space is None:
temporary_address_space = lp.auto
new_temp_shape = tuple(abm.non1_storage_shape)
new_temporary_variables = kernel.temporary_variables.copy()
if temporary_name not in new_temporary_variables:
temp_var = lp.TemporaryVariable(
name=temporary_name,
dtype=dtype,
base_indices=(0, ) * len(new_temp_shape),
shape=tuple(abm.non1_storage_shape),
address_space=temporary_address_space,
dim_names=tuple(non1_storage_axis_names))
else:
temp_var = new_temporary_variables[temporary_name]
# {{{ check and adapt existing temporary
if temp_var.dtype is lp.auto:
pass
elif temp_var.dtype is not lp.auto and dtype is lp.auto:
dtype = temp_var.dtype
elif temp_var.dtype is not lp.auto and dtype is not lp.auto:
if temp_var.dtype != dtype:
raise LoopyError("Existing and new dtype of temporary '%s' "
"do not match (existing: %s, new: %s)" %
(temporary_name, temp_var.dtype, dtype))
temp_var = temp_var.copy(dtype=dtype)
if len(temp_var.shape) != len(new_temp_shape):
raise LoopyError(
"Existing and new temporary '%s' do not "
"have matching number of dimensions ('%d' vs. '%d') " %
(temporary_name, len(temp_var.shape), len(new_temp_shape)))
if temp_var.base_indices != (0, ) * len(new_temp_shape):
raise LoopyError(
"Existing and new temporary '%s' do not "
"have matching number of dimensions ('%d' vs. '%d') " %
(temporary_name, len(temp_var.shape), len(new_temp_shape)))
new_temp_shape = tuple(
max(i, ex_i) for i, ex_i in zip(new_temp_shape, temp_var.shape))
temp_var = temp_var.copy(shape=new_temp_shape)
if temporary_address_space == temp_var.address_space:
pass
elif temporary_address_space is lp.auto:
temporary_address_space = temp_var.address_space
elif temp_var.address_space is lp.auto:
pass
else:
raise LoopyError("Existing and new temporary '%s' do not "
"have matching scopes (existing: %s, new: %s)" %
(temporary_name,
AddressSpace.stringify(temp_var.address_space),
AddressSpace.stringify(temporary_address_space)))
temp_var = temp_var.copy(address_space=temporary_address_space)
# }}}
new_temporary_variables[temporary_name] = temp_var
kernel = kernel.copy(temporary_variables=new_temporary_variables)
# }}}
from loopy import tag_inames
kernel = tag_inames(kernel, new_iname_to_tag)
from loopy.kernel.data import AutoFitLocalIndexTag, filter_iname_tags_by_type
if filter_iname_tags_by_type(new_iname_to_tag.values(),
AutoFitLocalIndexTag):
from loopy.kernel.tools import assign_automatic_axes
kernel = assign_automatic_axes(kernel)
return kernel
def generate_code_for_sched_index(codegen_state, sched_index):
kernel = codegen_state.kernel
sched_item = kernel.schedule[sched_index]
if isinstance(sched_item, CallKernel):
assert not codegen_state.is_generating_device_code
from loopy.schedule import (gather_schedule_block, get_insn_ids_for_block_at)
_, past_end_i = gather_schedule_block(kernel.schedule, sched_index)
assert past_end_i <= codegen_state.schedule_index_end
extra_args = synthesize_idis_for_extra_args(kernel, sched_index)
new_codegen_state = codegen_state.copy(
is_generating_device_code=True,
gen_program_name=sched_item.kernel_name,
schedule_index_end=past_end_i-1,
implemented_data_info=(codegen_state.implemented_data_info
+ extra_args))
from loopy.codegen.result import generate_host_or_device_program
codegen_result = generate_host_or_device_program(
new_codegen_state, sched_index)
glob_grid, loc_grid = kernel.get_grid_sizes_for_insn_ids_as_exprs(
get_insn_ids_for_block_at(kernel.schedule, sched_index))
return merge_codegen_results(codegen_state, [
codegen_result,
codegen_state.ast_builder.get_kernel_call(
codegen_state,
sched_item.kernel_name,
glob_grid, loc_grid,
extra_args),
])
elif isinstance(sched_item, EnterLoop):
tags = kernel.iname_tags(sched_item.iname)
tags = tuple(tag for tag in tags if tag)
from loopy.codegen.loop import (
generate_unroll_loop,
generate_vectorize_loop,
generate_sequential_loop_dim_code)
from loopy.kernel.data import (UnrolledIlpTag, UnrollTag,
ForceSequentialTag, LoopedIlpTag, VectorizeTag,
InOrderSequentialSequentialTag, filter_iname_tags_by_type)
if filter_iname_tags_by_type(tags, (UnrollTag, UnrolledIlpTag)):
func = generate_unroll_loop
elif filter_iname_tags_by_type(tags, VectorizeTag):
func = generate_vectorize_loop
elif not tags or filter_iname_tags_by_type(tags, (LoopedIlpTag,
ForceSequentialTag, InOrderSequentialSequentialTag)):
func = generate_sequential_loop_dim_code
else:
raise RuntimeError("encountered (invalid) EnterLoop "
"for '%s', tagged '%s'"
% (sched_item.iname, ", ".join(str(tag) for tag in tags)))
return func(codegen_state, sched_index)
elif isinstance(sched_item, Barrier):
# {{{ emit barrier code
from loopy.codegen.result import CodeGenerationResult
if codegen_state.is_generating_device_code:
barrier_ast = codegen_state.ast_builder.emit_barrier(
sched_item.synchronization_kind, sched_item.mem_kind,
sched_item.comment)
if sched_item.originating_insn_id:
return CodeGenerationResult.new(
codegen_state,
sched_item.originating_insn_id,
barrier_ast,
codegen_state.implemented_domain)
else:
return barrier_ast
else:
# host code
if sched_item.synchronization_kind in ["global", "local"]:
# host code is assumed globally and locally synchronous
return CodeGenerationResult(
host_program=None,
device_programs=[],
implemented_domains={},
implemented_data_info=codegen_state.implemented_data_info)
else:
raise LoopyError("do not know how to emit code for barrier "
"synchronization kind '%s'" "in host code"
% sched_item.synchronization_kind)
# }}}
elif isinstance(sched_item, RunInstruction):
insn = kernel.id_to_insn[sched_item.insn_id]
from loopy.codegen.instruction import generate_instruction_code
return codegen_state.try_vectorized(
"instruction %s" % insn.id,
lambda inner_cgs: generate_instruction_code(inner_cgs, insn))
else:
raise RuntimeError("unexpected schedule item type: %s"
% type(sched_item))
def iname_tags_of_type(self, iname, tag_type_or_types):
from loopy.kernel.data import filter_iname_tags_by_type
return filter_iname_tags_by_type(self.inames[iname].tags,
tag_type_or_types)
