def get_loopy_insns_and_result_names(self):
from sumpy.symbolic import make_sym_vector
avec = make_sym_vector("a", self.dim)
bvec = make_sym_vector("b", self.dim)
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
logger.info("compute expansion expressions: start")
result_names = [
expand(i, sac, expn, avec, bvec)
for i, expn in enumerate(self.expansions)
]
logger.info("compute expansion expressions: done")
sac.run_global_cse()
from sumpy.codegen import to_loopy_insns
loopy_insns = to_loopy_insns(
sac.assignments.items(),
vector_names={"a", "b"},
pymbolic_expr_maps=[
expn.kernel.get_code_transformer() for expn in self.expansions
],
retain_names=result_names,
complex_dtype=np.complex128 # FIXME
)
return loopy_insns, result_names
def get_loopy_insns_and_result_names(self):
from sumpy.symbolic import make_sym_vector
avec = make_sym_vector("a", self.dim)
bvec = make_sym_vector("b", self.dim)
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
logger.info("compute expansion expressions: start")
result_names = [expand(i, sac, expn, avec, bvec)
for i, expn in enumerate(self.expansions)]
logger.info("compute expansion expressions: done")
sac.run_global_cse()
from sumpy.codegen import to_loopy_insns
loopy_insns = to_loopy_insns(
six.iteritems(sac.assignments),
vector_names=set(["a", "b"]),
pymbolic_expr_maps=[
expn.kernel.get_code_transformer() for expn in self.expansions],
retain_names=result_names,
complex_dtype=np.complex128 # FIXME
)
return loopy_insns, result_names
def test_line_taylor_coeff_growth():
# Regression test for LineTaylorLocalExpansion.
# See https://gitlab.tiker.net/inducer/pytential/merge_requests/12
from sumpy.kernel import LaplaceKernel
from sumpy.expansion.local import LineTaylorLocalExpansion
from sumpy.symbolic import make_sym_vector, SympyToPymbolicMapper
import numpy as np
order = 10
expn = LineTaylorLocalExpansion(LaplaceKernel(2), order)
avec = make_sym_vector("a", 2)
bvec = make_sym_vector("b", 2)
coeffs = expn.coefficients_from_source(avec, bvec, rscale=1)
sym2pymbolic = SympyToPymbolicMapper()
coeffs_pymbolic = [sym2pymbolic(c) for c in coeffs]
from pymbolic.mapper.flop_counter import FlopCounter
flop_counter = FlopCounter()
counts = [flop_counter(c) for c in coeffs_pymbolic]
indices = np.arange(1, order + 2)
max_order = 2
assert np.polyfit(np.log(indices), np.log(counts), deg=1)[0] < max_order
def test_line_taylor_coeff_growth():
# Regression test for LineTaylorLocalExpansion.
# See https://gitlab.tiker.net/inducer/pytential/merge_requests/12
from sumpy.kernel import LaplaceKernel
from sumpy.expansion.local import LineTaylorLocalExpansion
from sumpy.symbolic import make_sym_vector, SympyToPymbolicMapper
import numpy as np
order = 10
expn = LineTaylorLocalExpansion(LaplaceKernel(2), order)
avec = make_sym_vector("a", 2)
bvec = make_sym_vector("b", 2)
coeffs = expn.coefficients_from_source(avec, bvec, rscale=1)
sym2pymbolic = SympyToPymbolicMapper()
coeffs_pymbolic = [sym2pymbolic(c) for c in coeffs]
from pymbolic.mapper.flop_counter import FlopCounter
flop_counter = FlopCounter()
counts = [flop_counter(c) for c in coeffs_pymbolic]
indices = np.arange(1, order + 2)
max_order = 2
assert np.polyfit(np.log(indices), np.log(counts), deg=1)[0] < max_order
def to_sym(self, fnames=None):
from sumpy.symbolic import make_sym_vector, Function
x = list(make_sym_vector("x", self.dim))
x += list(make_sym_vector("t", self.total_dims - self.dim))
if fnames is None:
noutputs = 0
for eq in self.eqs:
for deriv_ident in eq.keys():
noutputs = max(noutputs, deriv_ident.vec_idx)
fnames = [f"f{i}" for i in range(noutputs+1)]
funcs = [Function(fname)(*x) for fname in fnames]
res = []
for eq in self.eqs:
sym_eq = 0
for deriv_ident, coeff in eq.items():
expr = funcs[deriv_ident.vec_idx]
for i, val in enumerate(deriv_ident.mi):
for j in range(val):
expr = expr.diff(x[i])
sym_eq += expr * coeff
res.append(sym_eq)
return res
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.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 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_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_insns_and_result_names(self):
from sumpy.symbolic import make_sym_vector
dvec = make_sym_vector("d", self.dim)
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
result_names = [
sac.assign_unique(
"knl%d" % i,
knl.postprocess_at_target(
knl.postprocess_at_source(knl.get_expression(dvec), dvec),
dvec)) 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(["d"]),
pymbolic_expr_maps=[
knl.get_code_transformer() for knl in self.kernels
],
retain_names=result_names,
complex_dtype=np.complex128 # FIXME
)
return loopy_insns, result_names
def get_loopy_insns_and_result_names(self):
from sumpy.symbolic import make_sym_vector
dvec = make_sym_vector("d", self.dim)
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
result_names = [
sac.assign_unique("knl%d" % i,
knl.postprocess_at_target(
knl.postprocess_at_source(
knl.get_expression(dvec),
dvec),
dvec)
)
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(["d"]),
pymbolic_expr_maps=[
knl.get_code_transformer() for knl in self.kernels],
retain_names=result_names,
complex_dtype=np.complex128 # FIXME
)
return loopy_insns, result_names
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 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 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 coefficients_from_source(self, avec, bvec, rscale, sac=None):
from sumpy.kernel import DirectionalSourceDerivative
kernel = self.kernel
from sumpy.tools import mi_power, mi_factorial
if not self.use_rscale:
rscale = 1
if isinstance(kernel, DirectionalSourceDerivative):
from sumpy.symbolic import make_sym_vector
dir_vecs = []
tmp_kernel = kernel
while isinstance(tmp_kernel, DirectionalSourceDerivative):
dir_vecs.append(
make_sym_vector(tmp_kernel.dir_vec_name, kernel.dim))
tmp_kernel = tmp_kernel.inner_kernel
if kernel.get_base_kernel() is not tmp_kernel:
raise NotImplementedError("Unknown kernel wrapper.")
nderivs = len(dir_vecs)
coeff_identifiers = self.get_full_coefficient_identifiers()
result = [0] * len(coeff_identifiers)
for i, mi in enumerate(coeff_identifiers):
# One source derivative is the dot product of the gradient and
# directional vector.
# For multiple derivatives, gradient of gradients is taken.
# For eg: in 3D, 2 source derivatives gives 9 terms and
# cartesian_product below enumerates these 9 terms.
for deriv_terms in cartesian_product(*[range(kernel.dim)] *
nderivs):
prod = 1
derivative_mi = list(mi)
for nderivative, deriv_dim in enumerate(deriv_terms):
prod *= -derivative_mi[deriv_dim]
prod *= dir_vecs[nderivative][deriv_dim]
derivative_mi[deriv_dim] -= 1
if any(v < 0 for v in derivative_mi):
continue
result[i] += mi_power(avec, derivative_mi) * prod
for i, mi in enumerate(coeff_identifiers):
result[i] /= (mi_factorial(mi) * rscale**sum(mi))
else:
avec = [sym.UnevaluatedExpr(a * rscale**-1) for a in avec]
result = [
mi_power(avec, mi) / mi_factorial(mi)
for mi in self.get_full_coefficient_identifiers()
]
return (self.expansion_terms_wrangler.
get_stored_mpole_coefficients_from_full(result,
rscale,
sac=sac))
def coefficients_from_source(self, avec, bvec, rscale):
from sumpy.kernel import DirectionalSourceDerivative
kernel = self.kernel
from sumpy.tools import mi_power, mi_factorial
if not self.use_rscale:
rscale = 1
if isinstance(kernel, DirectionalSourceDerivative):
if kernel.get_base_kernel() is not kernel.inner_kernel:
raise NotImplementedError("more than one source derivative "
"not supported at present")
from sumpy.symbolic import make_sym_vector
dir_vec = make_sym_vector(kernel.dir_vec_name, kernel.dim)
coeff_identifiers = self.get_full_coefficient_identifiers()
result = [0] * len(coeff_identifiers)
for idim in range(kernel.dim):
for i, mi in enumerate(coeff_identifiers):
if mi[idim] == 0:
continue
derivative_mi = tuple(mi_i - 1 if iaxis == idim else mi_i
for iaxis, mi_i in enumerate(mi))
result[i] += (
- mi_power(avec, derivative_mi) * mi[idim]
* dir_vec[idim])
for i, mi in enumerate(coeff_identifiers):
result[i] /= (mi_factorial(mi) * rscale ** sum(mi))
else:
avec = [sym.UnevaluatedExpr(a * rscale**-1) for a in avec]
result = [
mi_power(avec, mi) / mi_factorial(mi)
for mi in self.get_full_coefficient_identifiers()]
return (
self.derivative_wrangler.get_stored_mpole_coefficients_from_full(
result, rscale))
def coefficients_from_source(self, avec, bvec, rscale):
from sumpy.kernel import DirectionalSourceDerivative
kernel = self.kernel
from sumpy.tools import mi_power, mi_factorial
if not self.use_rscale:
rscale = 1
if isinstance(kernel, DirectionalSourceDerivative):
if kernel.get_base_kernel() is not kernel.inner_kernel:
raise NotImplementedError("more than one source derivative "
"not supported at present")
from sumpy.symbolic import make_sym_vector
dir_vec = make_sym_vector(kernel.dir_vec_name, kernel.dim)
coeff_identifiers = self.get_full_coefficient_identifiers()
result = [0] * len(coeff_identifiers)
for idim in range(kernel.dim):
for i, mi in enumerate(coeff_identifiers):
if mi[idim] == 0:
continue
derivative_mi = tuple(mi_i - 1 if iaxis == idim else mi_i
for iaxis, mi_i in enumerate(mi))
result[i] += (
- mi_power(avec, derivative_mi) * mi[idim]
* dir_vec[idim])
for i, mi in enumerate(coeff_identifiers):
result[i] /= (mi_factorial(mi) * rscale ** sum(mi))
else:
avec = avec * rscale**-1
result = [
mi_power(avec, mi) / mi_factorial(mi)
for mi in self.get_full_coefficient_identifiers()]
return (
self.derivative_wrangler.get_stored_mpole_coefficients_from_full(
result, rscale))
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 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 get_kernel(self):
from sumpy.symbolic import make_sym_vector
avec = make_sym_vector("a", self.dim)
bvec = make_sym_vector("b", self.dim)
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
logger.info("compute expansion expressions: start")
result_names = [
expand(i, sac, expn, avec, bvec)
for i, expn in enumerate(self.expansions)
]
logger.info("compute expansion expressions: done")
sac.run_global_cse()
from sumpy.codegen import to_loopy_insns
loopy_insns = to_loopy_insns(
six.iteritems(sac.assignments),
vector_names=set(["a", "b"]),
pymbolic_expr_maps=[
expn.kernel.get_code_transformer() for expn in self.expansions
],
retain_names=result_names,
complex_dtype=np.complex128 # FIXME
)
isrc_sym = var("isrc")
exprs = [
var(name) * self.get_strength_or_not(isrc_sym, i)
for i, name in enumerate(result_names)
]
from sumpy.tools import gather_loopy_source_arguments
arguments = (self.get_src_tgt_arguments() +
self.get_input_and_output_arguments() +
gather_loopy_source_arguments(self.kernels))
loopy_knl = lp.make_kernel(
"{[isrc,itgt,idim]: 0<=itgt<ntargets and 0<=isrc<nsources "
"and 0<=idim<dim}",
self.get_kernel_scaling_assignments() + ["for itgt, isrc"] +
[self.get_compute_a_and_b_vecs()] + loopy_insns + [
lp.Assignment(id=None,
assignee="pair_result_%d" % i,
expression=expr,
temp_var_type=lp.auto)
for i, (expr,
dtype) in enumerate(zip(exprs, self.value_dtypes))
] + ["end"] + self.get_result_store_instructions(),
arguments,
name=self.name,
assumptions="nsources>=1 and ntargets>=1",
default_offset=lp.auto,
silenced_warnings="write_race(write_lpot*)",
fixed_parameters=dict(dim=self.dim))
loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
for expn in self.expansions:
loopy_knl = expn.prepare_loopy_kernel(loopy_knl)
loopy_knl = lp.tag_array_axes(loopy_knl, "center", "sep,C")
return loopy_knl
def test_m2m_and_l2l_exprs_simpler(base_knl, local_expn_class,
mpole_expn_class, order,
with_source_derivative):
from sympy.core.cache import clear_cache
clear_cache()
np.random.seed(17)
extra_kwargs = {}
if isinstance(base_knl, HelmholtzKernel):
if base_knl.allow_evanescent:
extra_kwargs["k"] = 0.2 * (0.707 + 0.707j)
else:
extra_kwargs["k"] = 0.2
if isinstance(base_knl, StokesletKernel):
extra_kwargs["mu"] = 0.2
if with_source_derivative:
knl = DirectionalSourceDerivative(base_knl, "dir_vec")
else:
knl = base_knl
mpole_expn = mpole_expn_class(knl, order=order)
local_expn = local_expn_class(knl, order=order)
from sumpy.symbolic import make_sym_vector, Symbol, USE_SYMENGINE
dvec = make_sym_vector("d", knl.dim)
src_coeff_exprs = [
Symbol("src_coeff%d" % i) for i in range(len(mpole_expn))
]
src_rscale = 3
tgt_rscale = 2
faster_m2m = mpole_expn.translate_from(mpole_expn, src_coeff_exprs,
src_rscale, dvec, tgt_rscale)
slower_m2m = mpole_expn.translate_from(mpole_expn,
src_coeff_exprs,
src_rscale,
dvec,
tgt_rscale,
_fast_version=False)
def _check_equal(expr1, expr2):
if USE_SYMENGINE:
return float((expr1 - expr2).expand()) == 0.0
else:
# with sympy we are using UnevaluatedExpr and expand doesn't expand it
# Running doit replaces UnevaluatedExpr with evaluated exprs
return float((expr1 - expr2).doit().expand()) == 0.0
for expr1, expr2 in zip(faster_m2m, slower_m2m):
assert _check_equal(expr1, expr2)
faster_l2l = local_expn.translate_from(local_expn, src_coeff_exprs,
src_rscale, dvec, tgt_rscale)
slower_l2l = local_expn.translate_from(local_expn,
src_coeff_exprs,
src_rscale,
dvec,
tgt_rscale,
_fast_version=False)
for expr1, expr2 in zip(faster_l2l, slower_l2l):
assert _check_equal(expr1, expr2)
