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)
