def get_loopy_insns_and_result_names(self):
from sumpy.symbolic import make_sym_vector
bvec = make_sym_vector("b", self.dim)
import sumpy.symbolic as sp
rscale = sp.Symbol("rscale")
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
coeff_exprs = [
sym.Symbol("coeff%d" % i)
for i in range(len(self.expansion.get_coefficient_identifiers()))
]
value = self.expansion.evaluate(coeff_exprs, bvec, rscale)
result_names = [
sac.assign_unique("result_%d_p" % i,
knl.postprocess_at_target(value, bvec))
for i, knl in enumerate(self.kernels)
]
sac.run_global_cse()
from sumpy.codegen import to_loopy_insns
loopy_insns = to_loopy_insns(
six.iteritems(sac.assignments),
vector_names=set(["b"]),
pymbolic_expr_maps=[self.expansion.get_code_transformer()],
retain_names=result_names,
complex_dtype=np.complex128 # FIXME
)
return loopy_insns, result_names
def get_translation_loopy_insns(self):
from sumpy.symbolic import make_sym_vector
dvec = make_sym_vector("d", self.dim)
src_coeff_exprs = [sym.Symbol("src_coeff%d" % i)
for i in range(len(self.src_expansion))]
src_rscale = sym.Symbol("src_rscale")
tgt_rscale = sym.Symbol("tgt_rscale")
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
tgt_coeff_names = [
sac.assign_unique("coeff%d" % i, coeff_i)
for i, coeff_i in enumerate(
self.tgt_expansion.translate_from(
self.src_expansion, src_coeff_exprs, src_rscale,
dvec, tgt_rscale, sac=sac))]
sac.run_global_cse()
from sumpy.codegen import to_loopy_insns
return to_loopy_insns(
six.iteritems(sac.assignments),
vector_names=set(["d"]),
pymbolic_expr_maps=[self.tgt_expansion.get_code_transformer()],
retain_names=tgt_coeff_names,
complex_dtype=np.complex128 # FIXME
)
def expand(expansion_nr, sac, expansion, avec, bvec):
rscale = sym.Symbol("rscale")
coefficients = expansion.coefficients_from_source(avec, bvec, rscale)
assigned_coeffs = [
sym.Symbol(
sac.assign_unique(
"expn%dcoeff%s" % (expansion_nr, stringify_expn_index(i)),
coefficients[expansion.get_storage_index(i)]))
for i in expansion.get_coefficient_identifiers()
]
return sac.assign_unique("expn%d_result" % expansion_nr,
expansion.evaluate(assigned_coeffs, bvec, rscale))
def try_get_recurrence_for_derivative(self, deriv, in_terms_of):
deriv = np.array(deriv, dtype=int)
for dim in np.where(2 <= deriv)[0]:
# Check if we can reduce this dimension in terms of the other
# dimensions.
reduced_deriv = deriv.copy()
reduced_deriv[dim] -= 2
coeffs = {}
for other_dim in range(self.dim):
if other_dim == dim:
continue
needed_deriv = reduced_deriv.copy()
needed_deriv[other_dim] += 2
needed_deriv = tuple(needed_deriv)
if needed_deriv not in in_terms_of:
break
coeffs[needed_deriv] = -1
else:
k = sym.Symbol(self.helmholtz_k_name)
coeffs[tuple(reduced_deriv)] = -k * k
return coeffs
def get_loopy_instructions(self):
from sumpy.symbolic import make_sym_vector
avec = make_sym_vector("a", self.dim)
import sumpy.symbolic as sp
rscale = sp.Symbol("rscale")
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
coeff_names = [
sac.assign_unique("coeff%d" % i, coeff_i)
for i, coeff_i in enumerate(
self.expansion.coefficients_from_source(avec, None, rscale))]
sac.run_global_cse()
from sumpy.codegen import to_loopy_insns
return to_loopy_insns(
six.iteritems(sac.assignments),
vector_names=set(["a"]),
pymbolic_expr_maps=[self.expansion.get_code_transformer()],
retain_names=coeff_names,
complex_dtype=np.complex128 # FIXME
)
def get_loopy_instructions(self):
from sumpy.symbolic import make_sym_vector
avec = make_sym_vector("a", self.dim)
import sumpy.symbolic as sp
rscale = sp.Symbol("rscale")
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
coeff_names = []
for knl_idx, kernel in enumerate(self.kernels):
for i, coeff_i in enumerate(
self.expansion.coefficients_from_source(
kernel, avec, None, rscale, sac)):
sac.add_assignment(f"coeff{i}_{knl_idx}", coeff_i)
coeff_names.append(f"coeff{i}_{knl_idx}")
sac.run_global_cse()
code_transformers = [self.expansion.get_code_transformer()] \
+ [kernel.get_code_transformer() for kernel in self.kernels]
from sumpy.codegen import to_loopy_insns
return to_loopy_insns(
sac.assignments.items(),
vector_names={"a"},
pymbolic_expr_maps=code_transformers,
retain_names=coeff_names,
complex_dtype=np.complex128 # FIXME
)
def track_m2l_op_count(self, param):
knl = self.knl(param.dim)
m_expn = self.mpole_expn_class(knl, order=param.order)
l_expn = self.local_expn_class(knl, order=param.order)
src_coeff_exprs = [
sym.Symbol("src_coeff%d" % i) for i in range(len(m_expn))
]
dvec = sym.make_sym_vector("d", knl.dim)
src_rscale = sym.Symbol("src_rscale")
tgt_rscale = sym.Symbol("tgt_rscale")
result = l_expn.translate_from(m_expn, src_coeff_exprs, src_rscale,
dvec, tgt_rscale)
sac = SymbolicAssignmentCollection()
for i, expr in enumerate(result):
sac.assign_unique("coeff%d" % i, expr)
sac.run_global_cse()
insns = to_loopy_insns(six.iteritems(sac.assignments))
counter = pymbolic.mapper.flop_counter.CSEAwareFlopCounter()
return sum([counter.rec(insn.expression) + 1 for insn in insns])
def __call__(self, base="expr"):
base = self._normalize(base)
count = self.base_to_count[base]
def make_id_str(base, count):
return "{base}{suffix}".format(base=base,
suffix="" if count == 0 else "_" +
str(count - 1))
id_str = make_id_str(base, count)
while id_str in self.taken_symbols:
count += 1
id_str = make_id_str(base, count)
self.base_to_count[base] = count + 1
return sym.Symbol(id_str)
def coefficients_from_source(self, avec, bvec, rscale):
# no point in heeding rscale here--just ignore it
if bvec is None:
raise RuntimeError(
"cannot use line-Taylor expansions in a setting "
"where the center-target vector is not known at coefficient "
"formation")
tau = sym.Symbol("tau")
avec_line = avec + tau * bvec
line_kernel = self.kernel.get_expression(avec_line)
from sumpy.symbolic import USE_SYMENGINE
if USE_SYMENGINE:
from sumpy.tools import MiDerivativeTaker, my_syntactic_subs
deriv_taker = MiDerivativeTaker(line_kernel, (tau, ))
return [
my_syntactic_subs(
self.kernel.postprocess_at_target(
self.kernel.postprocess_at_source(
deriv_taker.diff(i), avec), bvec), {tau: 0})
for i in self.get_coefficient_identifiers()
]
else:
# Workaround for sympy. The automatic distribution after
# single-variable diff makes the expressions very large
# (https://github.com/sympy/sympy/issues/4596), so avoid doing
# single variable diff.
#
# See also https://gitlab.tiker.net/inducer/pytential/merge_requests/12
return [
self.kernel.postprocess_at_target(
self.kernel.postprocess_at_source(
line_kernel.diff("tau", i), avec),
bvec).subs("tau", 0)
for i in self.get_coefficient_identifiers()
]
def get_pde_as_diff_op(self, **kwargs):
w = make_identity_diff_op(self.dim)
k = sym.Symbol(self.helmholtz_k_name)
return (laplacian(w) + k**2 * w)
def save_intermediate(expr):
if sac is None:
return expr
return sym.Symbol(sac.assign_unique("compress_temp", expr))
def save_intermediate(expr):
if sac is None:
return expr
return sym.Symbol(sac.assign_unique("projection_temp", expr))
def get_bessel_arg_scaling(self):
return sym.I * sym.Symbol(self.kernel.get_base_kernel().yukawa_lambda_name)
def get_bessel_arg_scaling(self):
return sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name)
class LinearRecurrenceBasedDerivativeWrangler(DerivativeWrangler):
_rscale_symbol = sp.Symbol("_sumpy_rscale_placeholder")
def get_coefficient_identifiers(self):
return self.stored_identifiers
def get_full_kernel_derivatives_from_stored(self,
stored_kernel_derivatives,
rscale):
coeff_matrix = self.get_coefficient_matrix(rscale)
return _spmv(coeff_matrix,
stored_kernel_derivatives,
sparse_vectors=False)
def get_stored_mpole_coefficients_from_full(self, full_mpole_coefficients,
rscale):
# = M^T x, where M = coeff matrix
coeff_matrix = self.get_coefficient_matrix(rscale)
result = [0] * len(self.stored_identifiers)
for row, coeff in enumerate(full_mpole_coefficients):
for col, val in coeff_matrix[row]:
result[col] += coeff * val
return result
@property
def stored_identifiers(self):
stored_identifiers, coeff_matrix = self._get_stored_ids_and_coeff_mat()
return stored_identifiers
@memoize_method
def get_coefficient_matrix(self, rscale):
"""
Return a matrix that expresses every derivative in terms of a
set of "stored" derivatives.
For example, for the recurrence
u_xx + u_yy + u_zz = 0
the coefficient matrix features the following entries:
... u_xx u_yy ... <= cols = only stored derivatives
==================
...| ... ... ... ...
|
u_zz| ... -1 -1 ...
^ rows = one for every derivative
"""
stored_identifiers, coeff_matrix = self._get_stored_ids_and_coeff_mat()
# substitute actual rscale for internal placeholder
return defaultdict(
lambda: [],
((irow, [(icol, coeff.xreplace({self._rscale_symbol: rscale})
if isinstance(coeff, sp.Basic) else coeff)
for icol, coeff in row])
for irow, row in six.iteritems(coeff_matrix)))
@memoize_method
def _get_stored_ids_and_coeff_mat(self):
stored_identifiers = []
identifiers_so_far = {}
import time
start_time = time.time()
logger.debug("computing recurrence for Taylor coefficients: start")
# Sparse matrix, indexed by row
coeff_matrix_transpose = defaultdict(lambda: [])
# Build up the matrix transpose by row.
from six import iteritems
for i, identifier in enumerate(
self.get_full_coefficient_identifiers()):
expr = self.try_get_recurrence_for_derivative(
identifier, identifiers_so_far, rscale=self._rscale_symbol)
if expr is None:
# Identifier should be stored
coeff_matrix_transpose[len(stored_identifiers)] = [(i, 1)]
stored_identifiers.append(identifier)
else:
expr = dict((identifiers_so_far[ident], val)
for ident, val in iteritems(expr))
result = _spmv(coeff_matrix_transpose,
expr,
sparse_vectors=True)
for j, item in iteritems(result):
coeff_matrix_transpose[j].append((i, item))
identifiers_so_far[identifier] = i
stored_identifiers = stored_identifiers
coeff_matrix = defaultdict(lambda: [])
for i, row in iteritems(coeff_matrix_transpose):
for j, val in row:
coeff_matrix[j].append((i, val))
logger.debug("computing recurrence for Taylor coefficients: "
"done after {dur:.2f} seconds".format(dur=time.time() -
start_time))
logger.debug(
"number of Taylor coefficients was reduced from {orig} to {red}".
format(orig=len(self.get_full_coefficient_identifiers()),
red=len(stored_identifiers)))
return stored_identifiers, coeff_matrix
def try_get_recurrence_for_derivative(self, deriv, in_terms_of, rscale):
"""
:arg deriv: a tuple of integers identifying a derivative for which
a recurrence is sought
:arg in_terms_of: a container supporting efficient containment testing
indicating availability of derivatives for use in recurrences
"""
raise NotImplementedError
def get_derivative_taker(self, expr, var_list):
from sumpy.tools import LinearRecurrenceBasedMiDerivativeTaker
return LinearRecurrenceBasedMiDerivativeTaker(expr, var_list, self)