def map_subscript(self, expr, domain, insn_id):
WalkMapper.map_subscript(self, expr, domain, insn_id)
from pymbolic.primitives import Variable
assert isinstance(expr.aggregate, Variable)
shape = None
var_name = expr.aggregate.name
if var_name in self.kernel.arg_dict:
arg = self.kernel.arg_dict[var_name]
shape = arg.shape
elif var_name in self.kernel.temporary_variables:
tv = self.kernel.temporary_variables[var_name]
shape = tv.shape
if shape is not None:
subscript = expr.index
if not isinstance(subscript, tuple):
subscript = (subscript, )
from loopy.symbolic import get_dependencies
available_vars = set(domain.get_var_dict())
shape_deps = set()
for shape_axis in shape:
if shape_axis is not None:
shape_deps.update(get_dependencies(shape_axis))
if not (get_dependencies(subscript) <= available_vars
and shape_deps <= available_vars):
return
if len(subscript) != len(shape):
raise LoopyError(
"subscript to '%s' in '%s' has the wrong "
"number of indices (got: %d, expected: %d)" %
(expr.aggregate.name, expr, len(subscript), len(shape)))
access_range = self._get_access_range(domain, subscript)
if access_range is None:
# Likely: index was non-affine, nothing we can do.
return
shape_domain = isl.BasicSet.universe(access_range.get_space())
for idim in range(len(subscript)):
shape_axis = shape[idim]
if shape_axis is not None:
slab = self._make_slab(shape_domain.get_space(),
(dim_type.in_, idim), 0, shape_axis)
shape_domain = shape_domain.intersect(slab)
if not access_range.is_subset(shape_domain):
raise LoopyError(
"'%s' in instruction '%s' "
"accesses out-of-bounds array element (could not"
" establish '%s' is a subset of '%s')." %
(expr, insn_id, access_range, shape_domain))
def map_subscript(self, expr):
WalkMapper.map_subscript(self, expr)
from pymbolic.primitives import Variable
assert isinstance(expr.aggregate, Variable)
if expr.aggregate.name == self.arg_name:
if not isinstance(expr.index, tuple):
self.index_ranks.append(1)
else:
self.index_ranks.append(len(expr.index))
def map_subscript(self, expr):
WalkMapper.map_subscript(self, expr)
from pymbolic.primitives import Variable
assert isinstance(expr.aggregate, Variable)
shape = None
var_name = expr.aggregate.name
if var_name in self.kernel.arg_dict:
arg = self.kernel.arg_dict[var_name]
shape = arg.shape
elif var_name in self.kernel.temporary_variables:
tv = self.kernel.temporary_variables[var_name]
shape = tv.shape
if shape is not None:
subscript = expr.index
if not isinstance(subscript, tuple):
subscript = (subscript,)
from loopy.symbolic import get_dependencies, get_access_range
available_vars = set(self.domain.get_var_dict())
shape_deps = set()
for shape_axis in shape:
if shape_axis is not None:
shape_deps.update(get_dependencies(shape_axis))
if not (get_dependencies(subscript) <= available_vars
and shape_deps <= available_vars):
return
if len(subscript) != len(shape):
raise LoopyError("subscript to '%s' in '%s' has the wrong "
"number of indices (got: %d, expected: %d)" % (
expr.aggregate.name, expr,
len(subscript), len(shape)))
try:
access_range = get_access_range(self.domain, subscript,
self.kernel.assumptions)
except isl.Error:
# Likely: index was non-linear, nothing we can do.
return
except TypeError:
# Likely: index was non-linear, nothing we can do.
return
shape_domain = isl.BasicSet.universe(access_range.get_space())
for idim in range(len(subscript)):
shape_axis = shape[idim]
if shape_axis is not None:
from loopy.isl_helpers import make_slab
slab = make_slab(
shape_domain.get_space(), (dim_type.in_, idim),
0, shape_axis)
shape_domain = shape_domain.intersect(slab)
if not access_range.is_subset(shape_domain):
raise LoopyError("'%s' in instruction '%s' "
"accesses out-of-bounds array element"
% (expr, self.insn_id))
def extract_subst(kernel, subst_name, template, parameters=()):
"""
:arg subst_name: The name of the substitution rule to be created.
:arg template: Unification template expression.
:arg parameters: An iterable of parameters used in
*template*, or a comma-separated string of the same.
All targeted subexpressions must match ('unify with') *template*
The template may contain '*' wildcards that will have to match exactly across all
unifications.
"""
if isinstance(template, str):
from pymbolic import parse
template = parse(template)
if isinstance(parameters, str):
parameters = tuple(s.strip() for s in parameters.split(","))
var_name_gen = kernel.get_var_name_generator()
# {{{ replace any wildcards in template with new variables
def get_unique_var_name():
based_on = subst_name + "_wc"
result = var_name_gen(based_on)
return result
from loopy.symbolic import WildcardToUniqueVariableMapper
wc_map = WildcardToUniqueVariableMapper(get_unique_var_name)
template = wc_map(template)
# }}}
# {{{ gather up expressions
expr_descriptors = []
from loopy.symbolic import UnidirectionalUnifier
unif = UnidirectionalUnifier(lhs_mapping_candidates=set(parameters))
def gather_exprs(expr, mapper):
urecs = unif(template, expr)
if urecs:
if len(urecs) > 1:
raise RuntimeError(
"ambiguous unification of '%s' with template '%s'" %
(expr, template))
urec, = urecs
expr_descriptors.append(
ExprDescriptor(insn=insn,
expr=expr,
unif_var_dict=dict(
(lhs.name, rhs)
for lhs, rhs in urec.equations)))
else:
mapper.fallback_mapper(expr)
# can't nest, don't recurse
from loopy.symbolic import (CallbackMapper, WalkMapper, IdentityMapper)
dfmapper = CallbackMapper(gather_exprs, WalkMapper())
for insn in kernel.instructions:
dfmapper(insn.assignees)
dfmapper(insn.expression)
for sr in six.itervalues(kernel.substitutions):
dfmapper(sr.expression)
# }}}
if not expr_descriptors:
raise RuntimeError("no expressions matching '%s'" % template)
# {{{ substitute rule into instructions
def replace_exprs(expr, mapper):
found = False
for exprd in expr_descriptors:
if expr is exprd.expr:
found = True
break
if not found:
return mapper.fallback_mapper(expr)
args = [exprd.unif_var_dict[arg_name] for arg_name in parameters]
result = var(subst_name)
if args:
result = result(*args)
return result
# can't nest, don't recurse
cbmapper = CallbackMapper(replace_exprs, IdentityMapper())
new_insns = []
for insn in kernel.instructions:
new_insns.append(insn.with_transformed_expressions(cbmapper))
from loopy.kernel.data import SubstitutionRule
new_substs = {
subst_name:
SubstitutionRule(
name=subst_name,
arguments=tuple(parameters),
expression=template,
)
}
for subst in six.itervalues(kernel.substitutions):
new_substs[subst.name] = subst.copy(
expression=cbmapper(subst.expression))
# }}}
return kernel.copy(instructions=new_insns, substitutions=new_substs)
def extract_subst(kernel, subst_name, template, parameters=(), within=None):
"""
:arg subst_name: The name of the substitution rule to be created.
:arg template: Unification template expression.
:arg parameters: An iterable of parameters used in
*template*, or a comma-separated string of the same.
:arg within: An instance of :class:`loopy.match.MatchExpressionBase` or
:class:`str` as understood by :func:`loopy.match.parse_match`.
All targeted subexpressions must match ('unify with') *template*
The template may contain '*' wildcards that will have to match exactly across all
unifications.
"""
if isinstance(kernel, TranslationUnit):
kernel_names = [
i for i, clbl in kernel.callables_table.items()
if isinstance(clbl, CallableKernel)
]
if len(kernel_names) != 1:
raise LoopyError()
return kernel.with_kernel(
extract_subst(kernel[kernel_names[0]], subst_name, template,
parameters))
if isinstance(template, str):
from pymbolic import parse
template = parse(template)
if isinstance(parameters, str):
parameters = tuple(s.strip() for s in parameters.split(","))
from loopy.match import parse_match
within = parse_match(within)
var_name_gen = kernel.get_var_name_generator()
# {{{ replace any wildcards in template with new variables
def get_unique_var_name():
based_on = subst_name + "_wc"
result = var_name_gen(based_on)
return result
from loopy.symbolic import WildcardToUniqueVariableMapper
wc_map = WildcardToUniqueVariableMapper(get_unique_var_name)
template = wc_map(template)
# }}}
# {{{ gather up expressions
expr_descriptors = []
from loopy.symbolic import UnidirectionalUnifier
unif = UnidirectionalUnifier(lhs_mapping_candidates=set(parameters))
def gather_exprs(expr, mapper):
urecs = unif(template, expr)
if urecs:
if len(urecs) > 1:
raise RuntimeError(
"ambiguous unification of '%s' with template '%s'" %
(expr, template))
urec, = urecs
expr_descriptors.append(
ExprDescriptor(insn=insn,
expr=expr,
unif_var_dict={
lhs.name: rhs
for lhs, rhs in urec.equations
}))
else:
mapper.fallback_mapper(expr)
# can't nest, don't recurse
from loopy.symbolic import (CallbackMapper, WalkMapper, IdentityMapper)
dfmapper = CallbackMapper(gather_exprs, WalkMapper())
from loopy.kernel.instruction import MultiAssignmentBase
for insn in kernel.instructions:
if isinstance(insn, MultiAssignmentBase) and within(kernel, insn):
dfmapper(insn.assignees)
dfmapper(insn.expression)
for sr in kernel.substitutions.values():
dfmapper(sr.expression)
# }}}
if not expr_descriptors:
raise RuntimeError("no expressions matching '%s'" % template)
# {{{ substitute rule into instructions
def replace_exprs(expr, mapper):
found = False
for exprd in expr_descriptors:
if expr is exprd.expr:
found = True
break
if not found:
return mapper.fallback_mapper(expr)
args = [exprd.unif_var_dict[arg_name] for arg_name in parameters]
result = var(subst_name)
if args:
result = result(*args)
return result
# can't nest, don't recurse
cbmapper = CallbackMapper(replace_exprs, IdentityMapper())
new_insns = []
def transform_assignee(expr):
# Assignment LHS's cannot be subst rules. Treat them
# specially.
import pymbolic.primitives as prim
if isinstance(expr, tuple):
return tuple(transform_assignee(expr_i) for expr_i in expr)
elif isinstance(expr, prim.Subscript):
return type(expr)(expr.aggregate, cbmapper(expr.index))
elif isinstance(expr, prim.Variable):
return expr
else:
raise ValueError("assignment LHS not understood")
for insn in kernel.instructions:
if within(kernel, insn):
new_insns.append(
insn.with_transformed_expressions(
cbmapper, assignee_f=transform_assignee))
else:
new_insns.append(insn)
from loopy.kernel.data import SubstitutionRule
new_substs = {
subst_name:
SubstitutionRule(
name=subst_name,
arguments=tuple(parameters),
expression=template,
)
}
for subst in kernel.substitutions.values():
new_substs[subst.name] = subst.copy(
expression=cbmapper(subst.expression))
# }}}
return kernel.copy(instructions=new_insns, substitutions=new_substs)
