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