def knl(): return make_loopy_program( [ "{[idof_init]: 0 <= idof_init < nnodes_tgt}", "{[jdof_init]: 0 <= jdof_init < nnodes_src}", "{[iel]: 0 <= iel < nelements}", "{[idof]: 0 <= idof < n_to_nodes}", "{[jdof]: 0 <= jdof < n_from_nodes}" ], """ result[idof_init, jdof_init] = 0 {id=init} ... gbarrier {id=barrier, dep=init} result[itgt_base + to_element_indices[iel]*n_to_nodes + idof, \ isrc_base + from_element_indices[iel]*n_from_nodes + jdof] \ = resample_mat[idof, jdof] {dep=barrier} """, [ lp.GlobalArg("result", None, shape="nnodes_tgt, nnodes_src", offset=lp.auto), lp.ValueArg("itgt_base, isrc_base", np.int32), lp.ValueArg("nnodes_tgt, nnodes_src", np.int32), "...", ], name="oversample_mat")
def kproj(): return make_loopy_program( [ "{[iel_init]: 0 <= iel_init < n_to_elements}", "{[idof_init]: 0 <= idof_init < n_to_nodes}", "{[iel]: 0 <= iel < nelements}", "{[i_quad]: 0 <= i_quad < n_to_nodes}", "{[ibasis]: 0 <= ibasis < n_to_nodes}" ], """ result[iel_init, idof_init] = 0 {id=init} ... gbarrier {id=barrier, dep=init} result[to_element_indices[iel], ibasis] = \ result[to_element_indices[iel], ibasis] + \ sum(i_quad, ary[from_element_indices[iel], i_quad] \ * basis_tabulation[ibasis, i_quad] \ * weights[i_quad]) {dep=barrier} """, [ lp.GlobalArg("ary", None, shape=("n_from_elements", "n_from_nodes")), lp.GlobalArg( "result", None, shape=("n_to_elements", "n_to_nodes")), lp.GlobalArg("basis_tabulation", None, shape=("n_to_nodes", "n_to_nodes")), lp.GlobalArg("weights", None, shape="n_from_nodes"), lp.ValueArg("n_from_elements", np.int32), lp.ValueArg("n_from_nodes", np.int32), lp.ValueArg("n_to_elements", np.int32), lp.ValueArg("n_to_nodes", np.int32), "..." ], name="conn_projection_knl")
def knl(sym_then, sym_else): return make_loopy_program( "{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}", [ lp.Assignment( var("out")[iel, idof], p.If(var("crit")[iel, idof], sym_then, sym_else)) ])
def prg(): return make_loopy_program( "{[iel, idof, jdof]: 0<=iel<nelements and 0<=idof, jdof<ndofs}", """ result[iel, idof] = %s(jdof, operand[iel, jdof]) """ % op_name, name="grudge_elementwise_%s" % op_name)
def knl(): return make_loopy_program( """{[iel,idof,f,j]: 0<=iel<nelements and 0<=f<nfaces and 0<=idof<nvol_nodes and 0<=j<nface_nodes}""", "result[iel,idof] = " "sum(f, sum(j, mat[idof, f, j] * vec[f, iel, j]))", name="face_mass")
def prg(): return make_loopy_program("""{[iel,idof,j]: 0<=iel<nelements and 0<=idof<ndiscr_nodes and 0<=j<nmesh_nodes}""", """ result[iel, idof] = \ sum(j, resampling_mat[idof, j] * nodes[iel, j]) """, name="nodes")
def prg(): result = make_loopy_program( """{[iel, idof, j]: 0<=iel<nelements and 0<=idof<ndiscr_nodes_out and 0<=j<ndiscr_nodes_in}""", "result[iel, idof] = sum(j, mat[idof, j] * vec[iel, j])", name="elwise_linear") result = lp.tag_array_axes(result, "mat", "stride:auto,stride:auto") return result
def resample_by_mat_prg(): from arraycontext import make_loopy_program return make_loopy_program( """ {[iel, idof, j]: 0 <= iel < nelements and 0 <= idof < nmesh_nodes and 0 <= j < ndiscr_nodes} """, """ result[iel, idof] = sum(j, resampling_mat[idof, j] * nodes[iel, j]) """, name="resample_by_mat_prg")
def node_knl(): t_unit = make_loopy_program("""{[iel, idof, jdof]: 0<=iel<nelements and 0<=idof, jdof<ndofs}""", """ result[iel, idof] = %s(jdof, operand[iel, jdof]) """ % reduction_name, name="nodewise_reduce") return lp.tag_inames( t_unit, { "iel": ConcurrentElementInameTag(), "idof": ConcurrentDOFInameTag(), })
def element_knl(): # FIXME: This computes the reduction value redundantly for each # output DOF. t_unit = make_loopy_program("""{[iel, jdof]: 0<=iel<nelements and 0<=jdof<ndofs} """, """ result[iel, 0] = %s(jdof, operand[iel, jdof]) """ % reduction_name, name="elementwise_reduce") return lp.tag_inames(t_unit, { "iel": ConcurrentElementInameTag(), })
def prg(): return make_loopy_program( [ "{[iel]: 0 <= iel < nelements}", "{[idof]: 0 <= idof < ndofs_per_element}" ], """ result[iel * ndofs_per_element + idof] = grp_ary[iel, idof] """, [ lp.GlobalArg( "result", None, shape="nelements * ndofs_per_element"), lp.GlobalArg( "grp_ary", None, shape=("nelements", "ndofs_per_element")), lp.ValueArg("nelements", np.int32), lp.ValueArg("ndofs_per_element", np.int32), "..." ], name="flatten_grp_ary")
def elementwise_prg(): # FIXME: This computes the reduction value redundantly for each # output DOF. t_unit = make_loopy_program([ "{[iel]: 0 <= iel < nelements}", "{[idof, jdof]: 0 <= idof, jdof < ndofs}" ], """ result[iel, idof] = %s(jdof, operand[iel, jdof]) """ % op_name, name="grudge_elementwise_%s_knl" % op_name) import loopy as lp from meshmode.transform_metadata import (ConcurrentElementInameTag, ConcurrentDOFInameTag) return lp.tag_inames(t_unit, { "iel": ConcurrentElementInameTag(), "idof": ConcurrentDOFInameTag() })
def prg(): t_unit = make_loopy_program([ "{[iel]: 0 <= iel < nelements}", "{[f]: 0 <= f < nfaces}", "{[idof]: 0 <= idof < nvol_nodes}", "{[jdof]: 0 <= jdof < nface_nodes}" ], """ result[iel, idof] = sum(f, sum(jdof, mat[idof, f, jdof] * jac_surf[f, iel, jdof] * vec[f, iel, jdof])) """, name="face_mass") import loopy as lp from meshmode.transform_metadata import (ConcurrentElementInameTag, ConcurrentDOFInameTag) return lp.tag_inames(t_unit, { "iel": ConcurrentElementInameTag(), "idof": ConcurrentDOFInameTag() })
def prg(): from arraycontext import make_loopy_program t_unit = make_loopy_program( "{[iel, idof]: 0 <= iel < nelements and 0 <= idof < nunit_dofs}", """ result[iel, 2*idof] = ary1[iel, idof] result[iel, 2*idof + 1] = ary2[iel, idof] """, [ lp.GlobalArg("ary1", shape="(nelements, nunit_dofs)"), lp.GlobalArg("ary2", shape="(nelements, nunit_dofs)"), lp.GlobalArg("result", shape="(nelements, 2*nunit_dofs)"), ... ], name="interleave") from meshmode.transform_metadata import (ConcurrentElementInameTag, ConcurrentDOFInameTag) return lp.tag_inames( t_unit, { "iel": ConcurrentElementInameTag(), "idof": ConcurrentDOFInameTag() })
def pick_knl(): knl = make_loopy_program( """{[iel, idof]: 0<=iel<nelements and 0<=idof<n_to_nodes}""", "result[to_element_indices[iel], idof] \ = ary[from_element_indices[iel], pick_list[idof]]", [ lp.GlobalArg("result", None, shape="nelements_result, n_to_nodes", offset=lp.auto), lp.GlobalArg("ary", None, shape="nelements_vec, n_from_nodes", offset=lp.auto), lp.ValueArg("nelements_result", np.int32), lp.ValueArg("nelements_vec", np.int32), lp.ValueArg("n_from_nodes", np.int32), "...", ], name="resample_by_picking") return knl
def mat_knl(): knl = make_loopy_program( """{[iel, idof, j]: 0<=iel<nelements and 0<=idof<n_to_nodes and 0<=j<n_from_nodes}""", "result[to_element_indices[iel], idof] \ = sum(j, resample_mat[idof, j] \ * ary[from_element_indices[iel], j])", [ lp.GlobalArg("result", None, shape="nelements_result, n_to_nodes", offset=lp.auto), lp.GlobalArg("ary", None, shape="nelements_vec, n_from_nodes", offset=lp.auto), lp.ValueArg("nelements_result", np.int32), lp.ValueArg("nelements_vec", np.int32), "...", ], name="resample_by_mat") return knl
def prg(): return make_loopy_program( "{[iel,idof,j]: 0 <= iel < nelements and 0 <= idof, j < nunit_dofs}", "result[iel,idof] = sum(j, diff_mat[idof, j] * vec[iel, j])", name="diff")
def prg(): return make_loopy_program( "{[iel,idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}", "result[iel,idof] = weights[idof]", name="quad_weights")
def prg(): return make_loopy_program( "{[iel,idof]: 0<=iel<nelements and 0<=idof<ndofs_per_element}", "result[iel, idof] = ary[grp_start + iel*ndofs_per_element + idof]", name="unflatten")