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