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 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
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
