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