def nodes(self):
r"""
:returns: object array of shape ``(ambient_dim,)`` containing
:class:`~meshmode.dof_array.DOFArray`\ s of node coordinates.
"""
actx = self._setup_actx
@memoize_in(actx, (Discretization, "nodes_prg"))
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")
return make_obj_array([
_DOFArray.from_list(None, [
actx.freeze(
actx.call_loopy(
prg(),
resampling_mat=actx.from_numpy(
grp.from_mesh_interp_matrix()),
nodes=actx.from_numpy(
grp.mesh_el_group.nodes[iaxis]))["result"])
for grp in self.groups
]) for iaxis in range(self.ambient_dim)
])
def num_reference_derivative(self, ref_axes, vec):
actx = vec.array_context
@memoize_in(actx, (Discretization, "reference_derivative_prg"))
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 get_mat(grp):
mat = None
for ref_axis in ref_axes:
next_mat = grp.diff_matrices()[ref_axis]
if mat is None:
mat = next_mat
else:
mat = np.dot(next_mat, mat)
return mat
return _DOFArray.from_list(actx, [
actx.call_loopy(prg(),
diff_mat=actx.from_numpy(get_mat(grp)),
vec=vec[grp.index])["result"]
for grp in self.groups
])
def _interleave_dof_arrays(self, ary1, ary2):
if not isinstance(ary1, DOFArray) or not isinstance(ary2, DOFArray):
raise TypeError("non-array passed to connection")
@memoize_in(self.array_context,
(CenterGranularityConnection, "interleave"))
def prg():
from meshmode.array_context import make_loopy_program
return make_loopy_program(
"""{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}""",
"""
dst[iel, 2*idof] = src1[iel, idof]
dst[iel, 2*idof + 1] = src2[iel, idof]
""",
[
lp.GlobalArg("src1", shape="(nelements, nunit_dofs)"),
lp.GlobalArg("src2", shape="(nelements, nunit_dofs)"),
lp.GlobalArg("dst", shape="(nelements, 2*nunit_dofs)"),
"...",
],
name="interleave")
results = []
for grp, src1, src2 in zip(self.discr.groups, ary1, ary2):
if src1.dtype != src2.dtype:
raise ValueError("dtype mismatch in inputs")
result = self.array_context.empty(
(grp.nelements, 2 * grp.nunit_dofs), dtype=src1.dtype)
self.array_context.call_loopy(
prg(), src1=src1, src2=src2, dst=result,
nelements=grp.nelements, nunit_dofs=grp.nunit_dofs)
results.append(result)
return DOFArray.from_list(self.array_context, results)
def _new_array(self, actx, creation_func, dtype=None):
if dtype is None:
dtype = self.real_dtype
elif dtype == "c":
dtype = self.complex_dtype
else:
dtype = np.dtype(dtype)
return _DOFArray.from_list(actx, [
creation_func(shape=(grp.nelements, grp.nunit_dofs), dtype=dtype)
for grp in self.groups
])
def _map_elementwise_reduction(self, reduction_name, expr):
@memoize_in(self.places, "elementwise_node_" + reduction_name)
def node_knl():
from meshmode.array_context import make_loopy_program
return 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")
@memoize_in(self.places, "elementwise_" + reduction_name)
def element_knl():
from meshmode.array_context import make_loopy_program
return 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")
discr = self.places.get_discretization(expr.dofdesc.geometry,
expr.dofdesc.discr_stage)
operand = self.rec(expr.operand)
assert operand.shape == (len(discr.groups), )
def _reduce(knl, result):
for grp in discr.groups:
self.array_context.call_loopy(knl,
operand=operand[grp.index],
result=result[grp.index])
return result
dtype = operand.entry_dtype
granularity = expr.dofdesc.granularity
if granularity is sym.GRANULARITY_NODE:
return _reduce(node_knl(),
discr.empty(self.array_context, dtype=dtype))
elif granularity is sym.GRANULARITY_ELEMENT:
result = DOFArray.from_list(self.array_context, [
self.array_context.empty((grp.nelements, 1), dtype=dtype)
for grp in discr.groups
])
return _reduce(element_knl(), result)
else:
raise ValueError('unsupported granularity: %s' % granularity)
def map_insn_loopy_kernel(self, insn, profile_data=None):
kdescr = insn.kernel_descriptor
discr = self.discrwb.discr_from_dd(kdescr.governing_dd)
dof_array_kwargs = {}
other_kwargs = {}
for name, expr in kdescr.input_mappings.items():
v = self.rec(expr)
if isinstance(v, DOFArray):
dof_array_kwargs[name] = v
else:
other_kwargs[name] = v
for name in kdescr.scalar_args():
v = other_kwargs[name]
if isinstance(v, (int, float)):
other_kwargs[name] = discr.real_dtype.type(v)
elif isinstance(v, complex):
other_kwargs[name] = discr.complex_dtype.type(v)
elif isinstance(v, np.number):
pass
else:
raise ValueError("unrecognized scalar type for variable '%s': %s"
% (name, type(v)))
result = {}
for grp in discr.groups:
kwargs = other_kwargs.copy()
kwargs["nelements"] = grp.nelements
kwargs["nunit_dofs"] = grp.nunit_dofs
for name, ary in dof_array_kwargs.items():
kwargs[name] = ary[grp.index]
knl_result = self.array_context.call_loopy(
kdescr.loopy_kernel, **kwargs)
for name, val in knl_result.items():
result.setdefault(name, []).append(val)
result = {
name: DOFArray.from_list(self.array_context, val)
for name, val in result.items()}
return list(result.items()), []
def quad_weights(self):
""":returns: A :class:`~meshmode.dof_array.DOFArray` with quadrature weights.
"""
actx = self._setup_actx
@memoize_in(actx, (Discretization, "quad_weights_prg"))
def prg():
return make_loopy_program(
"{[iel,idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}",
"result[iel,idof] = weights[idof]",
name="quad_weights")
return _DOFArray.from_list(None, [
actx.freeze(
actx.call_loopy(
prg(),
weights=actx.from_numpy(grp.weights),
nelements=grp.nelements,
)["result"]) for grp in self.groups
])
def test_node_reduction(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
actx = PyOpenCLArrayContext(queue)
# {{{ build discretization
target_order = 4
nelements = 32
mesh = make_curve_mesh(starfish, np.linspace(0.0, 1.0, nelements + 1),
target_order)
discr = Discretization(
actx, mesh, InterpolatoryQuadratureSimplexGroupFactory(target_order))
# }}}
# {{{ test
# create a shuffled [1, nelements + 1] array
ary = []
el_nr_base = 0
for grp in discr.groups:
x = 1 + np.arange(el_nr_base, grp.nelements)
np.random.shuffle(x)
ary.append(actx.freeze(actx.from_numpy(x.reshape(-1, 1))))
el_nr_base += grp.nelements
from meshmode.dof_array import DOFArray
ary = DOFArray.from_list(actx, ary)
for func, expected in [
(sym.NodeSum, nelements * (nelements + 1) // 2),
(sym.NodeMax, nelements),
(sym.NodeMin, 1),
]:
r = bind(discr, func(sym.var("x")))(actx, x=ary)
assert abs(r - expected) < 1.0e-15, r
def test_to_fd_idempotency(ctx_factory, mm_mesh, fspace_degree):
"""
Make sure mm->fd->mm and (mm->)->fd->mm->fd are identity
"""
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
actx = PyOpenCLArrayContext(queue)
# make sure degree is higher order than mesh
fspace_degree += mm_mesh.groups[0].order
# Make a function space and a function with unique values at each node
factory = InterpolatoryQuadratureSimplexGroupFactory(fspace_degree)
discr = Discretization(actx, mm_mesh, factory)
fdrake_connection = build_connection_to_firedrake(discr)
fdrake_mesh = fdrake_connection.firedrake_fspace().mesh()
dtype = fdrake_mesh.coordinates.dat.data.dtype
mm_unique = discr.zeros(actx, dtype=dtype)
unique_vals = np.arange(np.size(mm_unique[0]), dtype=dtype)
mm_unique[0].set(unique_vals.reshape(mm_unique[0].shape))
mm_unique_copy = DOFArray.from_list(actx, [mm_unique[0].copy()])
# Test for idempotency mm->fd->mm
fdrake_unique = fdrake_connection.from_meshmode(mm_unique)
fdrake_connection.from_firedrake(fdrake_unique, out=mm_unique_copy)
np.testing.assert_allclose(actx.to_numpy(mm_unique_copy[0]),
actx.to_numpy(mm_unique[0]),
atol=CLOSE_ATOL)
# Test for idempotency (mm->)fd->mm->fd
fdrake_unique_copy = fdrake_connection.from_meshmode(mm_unique_copy)
np.testing.assert_allclose(fdrake_unique_copy.dat.data,
fdrake_unique.dat.data,
atol=CLOSE_ATOL)
def __call__(self, ary):
from meshmode.dof_array import DOFArray
if not isinstance(ary, DOFArray):
raise TypeError("non-array passed to discretization connection")
actx = ary.array_context
@memoize_in(actx, (
DirectDiscretizationConnection,
"resample_by_mat_knl",
self.is_surjective,
))
def mat_knl():
if self.is_surjective:
domains = [
"""
{[iel, idof, j]:
0<=iel<nelements and
0<=idof<n_to_nodes and
0<=j<n_from_nodes}
""",
]
instructions = """
result[to_element_indices[iel], idof] \
= sum(j, resample_mat[idof, j] \
* ary[from_element_indices[iel], j])
"""
else:
domains = [
"""
{[iel_init, idof_init]:
0<=iel_init<nelements_result and
0<=idof_init<n_to_nodes}
""",
"""
{[iel, idof, j]:
0<=iel<nelements and
0<=idof<n_to_nodes and
0<=j<n_from_nodes}
""",
]
instructions = """
result[iel_init, idof_init] = 0 {id=init}
... gbarrier {id=barrier, dep=init}
result[to_element_indices[iel], idof] \
= sum(j, resample_mat[idof, j] \
* ary[from_element_indices[iel], j]) {dep=barrier}
"""
knl = make_loopy_program(
domains,
instructions, [
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_mat")
return knl
@memoize_in(actx, (DirectDiscretizationConnection,
"resample_by_picking_knl", self.is_surjective))
def pick_knl():
if self.is_surjective:
domains = [
"""{[iel, idof]:
0<=iel<nelements and
0<=idof<n_to_nodes}"""
]
instructions = """
result[to_element_indices[iel], idof] \
= ary[from_element_indices[iel], pick_list[idof]]
"""
else:
domains = [
"""
{[iel_init, idof_init]:
0<=iel_init<nelements_result and
0<=idof_init<n_to_nodes}
""", """
{[iel, idof]:
0<=iel<nelements and
0<=idof<n_to_nodes}
"""
]
instructions = """
result[iel_init, idof_init] = 0 {id=init}
... gbarrier {id=barrier, dep=init}
result[to_element_indices[iel], idof] \
= ary[from_element_indices[iel], pick_list[idof]] {dep=barrier}
"""
knl = make_loopy_program(
domains,
instructions, [
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
if ary.shape != (len(self.from_discr.groups), ):
raise ValueError("invalid shape of incoming resampling data")
group_idx_to_result = []
for i_tgrp, (tgrp,
cgrp) in enumerate(zip(self.to_discr.groups,
self.groups)):
kernels = [] # get kernels for each batch; to be fused eventually
kwargs = {} # kwargs to the fused kernel
for i_batch, batch in enumerate(cgrp.batches):
if batch.from_element_indices.size == 0:
continue
point_pick_indices = self._resample_point_pick_indices(
actx, i_tgrp, i_batch)
if point_pick_indices is None:
knl = mat_knl()
knl = lp.rename_argument(knl, "resample_mat",
f"resample_mat_{i_batch}")
kwargs[f"resample_mat_{i_batch}"] = (self._resample_matrix(
actx, i_tgrp, i_batch))
else:
knl = pick_knl()
knl = lp.rename_argument(knl, "pick_list",
f"pick_list_{i_batch}")
kwargs[f"pick_list_{i_batch}"] = point_pick_indices
# {{{ enforce different namespaces for the kernels
for iname in knl.all_inames():
knl = lp.rename_iname(knl, iname, f"{iname}_{i_batch}")
knl = lp.rename_argument(knl, "ary", f"ary_{i_batch}")
knl = lp.rename_argument(knl, "from_element_indices",
f"from_element_indices_{i_batch}")
knl = lp.rename_argument(knl, "to_element_indices",
f"to_element_indices_{i_batch}")
knl = lp.rename_argument(knl, "nelements",
f"nelements_{i_batch}")
# }}}
kwargs[f"ary_{i_batch}"] = ary[batch.from_group_index]
kwargs[f"from_element_indices_{i_batch}"] = (
batch.from_element_indices)
kwargs[f"to_element_indices_{i_batch}"] = (
batch.to_element_indices)
kernels.append(knl)
fused_knl = lp.fuse_kernels(kernels)
# order of operations doesn't matter
fused_knl = lp.add_nosync(fused_knl,
"global",
"writes:result",
"writes:result",
bidirectional=True,
force=True)
result_dict = actx.call_loopy(fused_knl,
nelements_result=tgrp.nelements,
n_to_nodes=tgrp.nunit_dofs,
**kwargs)
group_idx_to_result.append(result_dict["result"])
from meshmode.dof_array import DOFArray
return DOFArray.from_list(actx, group_idx_to_result)
def map_insn_loopy_kernel(self, insn, profile_data):
kdescr = insn.kernel_descriptor
discr = self.inner_mapper.discrwb.discr_from_dd(kdescr.governing_dd)
dof_array_kwargs = {}
other_kwargs = {}
for name, expr in kdescr.input_mappings.items():
v = self.inner_mapper.rec(expr)
if isinstance(v, DOFArray):
dof_array_kwargs[name] = v
if profile_data is not None:
size = dof_array_nbytes(v)
profile_data["bytes_read"] = (
profile_data.get("bytes_read", 0) + size)
profile_data["bytes_read_by_scalar_assignments"] = (
profile_data.get("bytes_read_by_scalar_assignments",
0) + size)
else:
assert not isinstance(v, np.ndarray)
other_kwargs[name] = v
for name in kdescr.scalar_args():
v = other_kwargs[name]
if isinstance(v, (int, float)):
other_kwargs[name] = discr.real_dtype.type(v)
elif isinstance(v, complex):
other_kwargs[name] = discr.complex_dtype.type(v)
elif isinstance(v, np.number):
pass
else:
raise ValueError(
"unrecognized scalar type for variable '%s': %s" %
(name, type(v)))
result = {}
for grp in discr.groups:
kwargs = other_kwargs.copy()
kwargs["nelements"] = grp.nelements
kwargs["nunit_dofs"] = grp.nunit_dofs
for name, ary in dof_array_kwargs.items():
kwargs[name] = ary[grp.index]
knl_result = self.inner_mapper.array_context.call_loopy(
kdescr.loopy_kernel, **kwargs)
for name, val in knl_result.items():
result.setdefault(name, []).append(val)
result = {
name: DOFArray.from_list(self.inner_mapper.array_context, val)
for name, val in result.items()
}
for val in result.values():
assert isinstance(val, DOFArray)
if profile_data is not None:
size = dof_array_nbytes(val)
profile_data["bytes_written"] = (
profile_data.get("bytes_written", 0) + size)
profile_data["bytes_written_by_scalar_assignments"] = (
profile_data.get("bytes_written_by_scalar_assignments",
0) + size)
return list(result.items()), []
def map_if(self, expr):
bool_crit = self.rec(expr.condition)
if isinstance(bool_crit, DOFArray):
# continues below
pass
elif isinstance(bool_crit, np.number):
if bool_crit:
return self.rec(expr.then)
else:
return self.rec(expr.else_)
else:
raise TypeError(
"Expected criterion to be of type np.number or DOFArray")
assert isinstance(bool_crit, DOFArray)
ngroups = len(bool_crit)
from pymbolic import var
iel = var("iel")
idof = var("idof")
subscript = (iel, idof)
then = self.rec(expr.then)
else_ = self.rec(expr.else_)
import pymbolic.primitives as p
var = p.Variable
if isinstance(then, DOFArray):
sym_then = var("a")[subscript]
def get_then(igrp):
return then[igrp]
elif isinstance(then, np.number):
sym_then = var("a")
def get_then(igrp):
return then
else:
raise TypeError(
"Expected 'then' to be of type np.number or DOFArray")
if isinstance(else_, DOFArray):
sym_else = var("b")[subscript]
def get_else(igrp):
return else_[igrp]
elif isinstance(else_, np.number):
sym_else = var("b")
def get_else(igrp):
return else_
else:
raise TypeError(
"Expected 'else' to be of type np.number or DOFArray")
@memoize_in(self.array_context, (ExecutionMapper, "map_if_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))
])
return DOFArray.from_list(self.array_context, [
self.array_context.call_loopy(
knl(sym_then, sym_else),
crit=bool_crit[igrp],
a=get_then(igrp),
b=get_else(igrp))
for igrp in range(ngroups)])
