def get_translation_loopy_insns(self):
from sumpy.symbolic import make_sympy_vector
dvec = make_sympy_vector("d", self.dim)
src_coeff_exprs = [sp.Symbol("src_coeff%d" % i)
for i in range(len(self.src_expansion))]
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, dvec))]
sac.run_global_cse()
from sumpy.symbolic import kill_trivial_assignments
assignments = kill_trivial_assignments([
(name, expr)
for name, expr in six.iteritems(sac.assignments)],
retain_names=tgt_coeff_names)
from sumpy.codegen import to_loopy_insns
return to_loopy_insns(
assignments,
vector_names=set(["d"]),
pymbolic_expr_maps=[self.tgt_expansion.get_code_transformer()],
complex_dtype=np.complex128 # FIXME
)
def coefficients_from_source(self, avec, bvec):
from sumpy.kernel import DirectionalSourceDerivative
kernel = self.kernel
from sumpy.tools import mi_power, mi_factorial
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_sympy_vector
dir_vec = make_sympy_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)
else:
result = [
mi_power(avec, mi) / mi_factorial(mi)
for mi in self.get_full_coefficient_identifiers()]
return self.full_to_stored(result)
def get_loopy_instructions(self):
from sumpy.symbolic import make_sympy_vector
avec = make_sympy_vector("a", self.dim)
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))]
sac.run_global_cse()
from sumpy.symbolic import kill_trivial_assignments
assignments = kill_trivial_assignments([
(name, expr)
for name, expr in six.iteritems(sac.assignments)],
retain_names=coeff_names)
from sumpy.codegen import to_loopy_insns
return to_loopy_insns(
assignments,
vector_names=set(["a"]),
pymbolic_expr_maps=[self.expansion.get_code_transformer()],
complex_dtype=np.complex128 # FIXME
)
def get_loopy_insns_and_result_names(self):
from sumpy.symbolic import make_sympy_vector
dvec = make_sympy_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],
complex_dtype=np.complex128 # FIXME
)
return loopy_insns, result_names
def postprocess_at_source(self, expr, avec):
expr = self.inner_kernel.postprocess_at_source(expr, avec)
dimensions = len(avec)
assert dimensions == self.dim
from sumpy.symbolic import make_sym_vector as make_sympy_vector
dir_vec = make_sympy_vector(self.dir_vec_name, dimensions)
# avec = center-src -> minus sign from chain rule
return sum(-dir_vec[axis] * expr.diff(avec[axis])
for axis in range(dimensions))
def postprocess_at_target(self, expr, bvec):
expr = self.inner_kernel.postprocess_at_target(expr, bvec)
dim = len(bvec)
assert dim == self.dim
from sumpy.symbolic import make_sym_vector as make_sympy_vector
dir_vec = make_sympy_vector(self.dir_vec_name, dim)
# bvec = tgt-center
return sum(dir_vec[axis] * expr.diff(bvec[axis])
for axis in range(dim))
def postprocess_at_source(self, expr, avec):
expr = self.inner_kernel.postprocess_at_source(expr, avec)
dimensions = len(avec)
assert dimensions == self.dim
from sumpy.symbolic import make_sympy_vector
dir_vec = make_sympy_vector(self.dir_vec_name, dimensions)
# avec = center-src -> minus sign from chain rule
return sum(-dir_vec[axis]*expr.diff(avec[axis])
for axis in range(dimensions))
def postprocess_at_target(self, expr, bvec):
expr = self.inner_kernel.postprocess_at_target(expr, bvec)
dim = len(bvec)
assert dim == self.dim
from sumpy.symbolic import make_sympy_vector
dir_vec = make_sympy_vector(self.dir_vec_name, dim)
# bvec = tgt-center
return sum(dir_vec[axis]*expr.diff(bvec[axis])
for axis in range(dim))
def get_loopy_insns_and_result_names(self):
from sumpy.symbolic import make_sympy_vector
bvec = make_sympy_vector("b", self.dim)
from sumpy.assignment_collection import SymbolicAssignmentCollection
sac = SymbolicAssignmentCollection()
coeff_exprs = [sp.Symbol("coeff%d" % i)
for i in range(len(self.expansion.get_coefficient_identifiers()))]
value = self.expansion.evaluate(coeff_exprs, bvec)
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.symbolic import kill_trivial_assignments
assignments = kill_trivial_assignments([
(name, expr)
for name, expr in six.iteritems(sac.assignments)],
retain_names=result_names)
from sumpy.codegen import to_loopy_insns
loopy_insns = to_loopy_insns(assignments,
vector_names=set(["b"]),
pymbolic_expr_maps=[self.expansion.get_code_transformer()],
complex_dtype=np.complex128 # FIXME
)
from pymbolic.interop.sympy import SympyToPymbolicMapper
sympy_conv = SympyToPymbolicMapper()
loopy_insns.append(
lp.Assignment(id=None,
assignee="kernel_scaling",
expression=sympy_conv(self.expansion.kernel.get_scaling()),
temp_var_type=lp.auto))
return loopy_insns, result_names
def get_kernel(self):
from sumpy.symbolic import make_sympy_vector
avec = make_sympy_vector("a", self.dim)
bvec = make_sympy_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.symbolic import kill_trivial_assignments
assignments = kill_trivial_assignments([
(name, expr.subs("tau", 0))
for name, expr in six.iteritems(sac.assignments)],
retain_names=result_names)
from sumpy.codegen import to_loopy_insns
loopy_insns = to_loopy_insns(assignments,
vector_names=set(["a", "b"]),
pymbolic_expr_maps=[
expn.kernel.get_code_transformer() for expn in self.expansions],
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*)"
)
loopy_knl = lp.fix_parameters(loopy_knl, 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
