def disambiguate_identifiers(statements_a,
statements_b,
should_disambiguate_name=None):
if should_disambiguate_name is None:
def should_disambiguate_name(name):
return True
from pymbolic.imperative.analysis import get_all_used_identifiers
id_a = get_all_used_identifiers(statements_a)
id_b = get_all_used_identifiers(statements_b)
from pytools import UniqueNameGenerator
vng = UniqueNameGenerator(id_a | id_b)
from pymbolic import var
subst_b = {}
for clash in id_a & id_b:
if should_disambiguate_name(clash):
unclash = vng(clash)
subst_b[clash] = var(unclash)
from pymbolic.mapper.substitutor import (make_subst_func,
SubstitutionMapper)
subst_map = SubstitutionMapper(make_subst_func(subst_b))
statements_b = [stmt.map_expressions(subst_map) for stmt in statements_b]
return statements_b, subst_b
def subst_into_pwqpolynomial(new_space, poly, subst_dict):
"""
Returns an instance of :class:`islpy.PwQPolynomial` with substitutions from
*subst_dict* substituted into *poly*.
:arg poly: an instance of :class:`islpy.PwQPolynomial`
:arg subst_dict: a mapping from parameters of *poly* to
:class:`pymbolic.primitives.Expression` made up of terms comprising the
parameters of *new_space*. The expression must be affine in the param
dims of *new_space*.
"""
if not poly.get_pieces():
# pw poly is univserally zero
result = isl.PwQPolynomial.zero(
new_space.insert_dims(dim_type.out, 0, 1))
assert result.dim(dim_type.out) == 1
return result
i_begin_subst_space = poly.dim(dim_type.param)
poly, subst_domain, subst_dict = get_param_subst_domain(
new_space, poly, subst_dict)
from loopy.symbolic import qpolynomial_to_expr, qpolynomial_from_expr
new_pieces = []
for valid_set, qpoly in poly.get_pieces():
valid_set = valid_set & subst_domain
if valid_set.plain_is_empty():
continue
valid_set = valid_set.project_out(dim_type.param, 0,
i_begin_subst_space)
from pymbolic.mapper.substitutor import (SubstitutionMapper,
make_subst_func)
sub_mapper = SubstitutionMapper(make_subst_func(subst_dict))
expr = sub_mapper(qpolynomial_to_expr(qpoly))
qpoly = qpolynomial_from_expr(valid_set.space, expr)
new_pieces.append((valid_set, qpoly))
if not new_pieces:
raise ValueError(
"no pieces of PwQPolynomial survived the substitution")
valid_set, qpoly = new_pieces[0]
result = isl.PwQPolynomial.alloc(valid_set, qpoly)
for valid_set, qpoly in new_pieces[1:]:
result = result.add_disjoint(isl.PwQPolynomial.alloc(valid_set, qpoly))
assert result.dim(dim_type.out)
return result
def _update_t_by_dt_factor(factor, statements):
from dagrt.language import Assign, Nop
from pymbolic import var
from pymbolic.mapper.substitutor import make_subst_func, SubstitutionMapper
mapper = SubstitutionMapper(make_subst_func({"<dt>":
factor * var("<dt>")}))
def updater(stmt):
if factor == 0:
return Nop(id=stmt.id, depends_on=stmt.depends_on)
return stmt.map_expressions(mapper)
return [
stmt if (not isinstance(stmt, Assign) or stmt.lhs != var("<t>")) else
updater(stmt) for stmt in statements
]
def parse(expr):
"""Return a pymbolic expression constructed from the string.
Values between backticks ("`") are parsed as variable names.
Tagged identifiers ("<func>f") are also parsed as variable names.
"""
from pymbolic import var
def remove_backticks(expr):
if not isinstance(expr, var):
return expr
varname = expr.name
if varname.startswith("`") and varname.endswith("`"):
return var(varname[1:-1])
return expr
from pymbolic.mapper.substitutor import SubstitutionMapper
parser = _ExtendedParser()
substitutor = SubstitutionMapper(remove_backticks)
return substitutor(parser(expr))
def fast_evaluator(matrix, sparse=False):
"""
Generate a function to evaluate a step matrix quickly.
The input comes from StepMatrixFinder.
"""
# First, rename variables in the matrix to names that are acceptable Python
# identifiers. We make use of dagrt's KeyToUniqueNameMap.
from dagrt.codegen.utils import KeyToUniqueNameMap
name_map = KeyToUniqueNameMap(forced_prefix="matrix")
def make_identifier(symbol):
from pymbolic import var
assert isinstance(symbol, var)
return var(name_map.get_or_make_name_for_key(symbol.name))
def get_var_order_from_name_map():
order = sorted(name_map)
return (order,
[name_map.get_or_make_name_for_key(key) for key in order])
from pymbolic.mapper.substitutor import SubstitutionMapper
substitutor = SubstitutionMapper(make_identifier)
from pymbolic import compile
# functools.partial ensures the resulting object is picklable.
from functools import partial
if sparse:
data = [substitutor(entry) for entry in matrix.data]
var_order, renamed_vars = get_var_order_from_name_map()
compiled_entries = [compile(entry, renamed_vars) for entry in data]
compiled_matrix = matrix.copy(data=compiled_entries)
else:
matrix = substitutor(matrix)
var_order, renamed_vars = get_var_order_from_name_map()
compiled_matrix = compile(matrix, renamed_vars)
return partial(_eval_compiled_matrix, compiled_matrix, var_order)
def subst_into_pwaff(new_space, pwaff, subst_dict):
"""
Returns an instance of :class:`islpy.PwAff` with substitutions from
*subst_dict* substituted into *pwaff*.
:arg pwaff: an instance of :class:`islpy.PwAff`
:arg subst_dict: a mapping from parameters of *pwaff* to
:class:`pymbolic.primitives.Expression` made up of terms comprising the
parameters of *new_space*. The expression must be affine in the param
dims of *new_space*.
"""
from pymbolic.mapper.substitutor import (SubstitutionMapper,
make_subst_func)
from loopy.symbolic import aff_from_expr, aff_to_expr
from functools import reduce
i_begin_subst_space = pwaff.dim(dim_type.param)
pwaff, subst_domain, subst_dict = get_param_subst_domain(
new_space, pwaff, subst_dict)
subst_mapper = SubstitutionMapper(make_subst_func(subst_dict))
pwaffs = []
for valid_set, qpoly in pwaff.get_pieces():
valid_set = valid_set & subst_domain
if valid_set.plain_is_empty():
continue
valid_set = valid_set.project_out(dim_type.param, 0,
i_begin_subst_space)
aff = aff_from_expr(valid_set.space, subst_mapper(aff_to_expr(qpoly)))
pwaffs.append(isl.PwAff.alloc(valid_set, aff))
if not pwaffs:
raise ValueError("no pieces of PwAff survived the substitution")
return reduce(lambda pwaff1, pwaff2: pwaff1.union_add(pwaff2),
pwaffs).coalesce()
def map_operator_binding(self, expr):
from hedge.optemplate import \
FluxOperatorBase, \
BoundaryPair, \
OperatorBinding, \
IdentityMapperMixin, \
InverseMassOperator
from pymbolic.mapper.substitutor import SubstitutionMapper
class FieldIntoBdrySubstitutionMapper(SubstitutionMapper,
IdentityMapperMixin):
def map_normal(self, expr):
return expr
if isinstance(expr.op, FluxOperatorBase):
if isinstance(expr.field, BoundaryPair):
return 0
else:
# Finally, an interior flux. Rewrite it.
def subst_func(expr):
if expr == self.vol_var:
return self.bdry_val_var
else:
return None
return OperatorBinding(
expr.op,
BoundaryPair(expr.field,
SubstitutionMapper(subst_func)(expr.field),
self.bdry_tag))
elif isinstance(expr.op, InverseMassOperator):
return OperatorBinding(expr.op, self.rec(expr.field))
else:
return 0
def strang_splitting(dag1, dag2, stepping_phase):
"""Given two time advancement routines (in *dag1* and *dag2*), returns a
single second-order accurate time advancement routine representing the sum
of both of those advancements.
:arg dag1: a :class:`dagrt.language.DAGCode`
:arg dag2: a :class:`dagrt.language.DAGCode`
:arg stepping_phase: the name of the phase in *dag1* and *dag2* that carries
out time stepping to which Strang splitting is to be applied.
:returns: a :class:`dagrt.language.DAGCode`
"""
from pymbolic.mapper.substitutor import make_subst_func, SubstitutionMapper
# {{{ disambiguate
id1 = dag1.existing_var_names()
id2 = dag1.existing_var_names()
from pytools import UniqueNameGenerator
vng = UniqueNameGenerator(id1 | id2)
from pymbolic import var
subst2 = {}
for clash in id1 & id2:
if not clash.startswith("<") or clash.startswith("<p>"):
unclash = vng(clash)
subst2[clash] = var(unclash)
subst2_mapper = SubstitutionMapper(make_subst_func(subst2))
# }}}
all_phases = frozenset(dag1.phases) | frozenset(dag2.phases)
from dagrt.language import DAGCode, ExecutionPhase
new_phases = {}
for phase_name in all_phases:
phase1 = dag1.phases.get(phase_name)
phase2 = dag2.phases.get(phase_name)
substed_s2_stmts = [
stmt.map_expressions(subst2_mapper) for stmt in phase2.statements
]
if phase_name == stepping_phase:
assert phase1 is not None
assert phase2 is not None
from pymbolic import var
dt_half = SubstitutionMapper(
make_subst_func({"<dt>": var("<dt>") / 2}))
phase1_half_dt = [
stmt.map_expressions(dt_half) for stmt in phase1.statements
]
if phase1.next_phase != phase2.next_phase:
raise ValueError(
"DAGs don't agree on default "
f"phase transition out of phase '{phase_name}'")
s2_name = phase_name + "_s2"
s3_name = phase_name + "_s3"
assert s2_name not in all_phases
assert s3_name not in all_phases
"""
du/dt = A + B
Time interval is [0,1]
1. Starting with u0, solve du / dt = A from t = 0 to 1/2, get u1
2. Starting with u1, solve du / dt = B from t = 0 to 1, get u2
3. Starting with u2, solve du / dt = A from t = 1/2 to 1, get u3
4. Return u3
"""
new_phases[phase_name] = ExecutionPhase(
name=phase_name,
next_phase=s2_name,
statements=(_update_t_by_dt_factor(
0, _elide_yield_state(phase1_half_dt))))
new_phases[s2_name] = ExecutionPhase(
name=s2_name,
next_phase=s3_name,
statements=(_update_t_by_dt_factor(
1 / 2, _elide_yield_state(substed_s2_stmts))))
new_phases[s3_name] = ExecutionPhase(name=s3_name,
next_phase=phase1.next_phase,
statements=phase1_half_dt)
else:
from dagrt.transform import fuse_two_phases
new_phases[phase_name] = fuse_two_phases(
phase_name, phase1, phase2.copy(statements=substed_s2_stmts))
if dag1.initial_phase != dag2.initial_phase:
raise ValueError("DAGs don't agree on initial phase")
return DAGCode(new_phases, dag1.initial_phase)