Updating how we flatten DOF arrays in anticipation for lazy evaluation

[thomasgibson]

This issue is specifically addressing how to update the flatten routine for DOFArrays, found here:

meshmode/meshmode/dof_array.py

Lines 381 to 412 in bbec89a

	def _flatten_dof_array(ary: Any, strict: bool = True):
	if not isinstance(ary, DOFArray):
	if strict:
	raise TypeError(f"non-DOFArray type '{type(ary).__name__}' cannot "
	"be flattened; use 'strict=False' to allow other types")
	else:
	return ary
	actx = ary.array_context
	if actx is None:
	raise ValueError("cannot flatten frozen DOFArrays")
	@memoize_in(actx, (_flatten_dof_array, "flatten_prg"))
	def prg():
	return make_loopy_program(
	"{[iel,idof]: 0<=iel<nelements and 0<=idof<ndofs_per_element}",
	"""result[grp_start + iel*ndofs_per_element + idof] \
	= grp_ary[iel, idof]""",
	name="flatten")
	group_sizes = [grp_ary.shape[0] * grp_ary.shape[1] for grp_ary in ary]
	group_starts = np.cumsum([0] + group_sizes)
	result = actx.empty(group_starts[-1], dtype=ary.entry_dtype)
	for grp_start, grp_ary in zip(group_starts, ary):
	actx.call_loopy(prg(),
	grp_ary=grp_ary,
	result=result,
	grp_start=grp_start)
	return result

Basically, the summary of the problem is that we need to remove stateful updates of the array since this will break lazy-evaluation. The intent of #192 is to hunt down and update all instances of stateful array updates in meshmode. Specifically for the flattening routines, this bit of code is the main culprit we need to get ride of:

meshmode/meshmode/dof_array.py

Lines 404 to 410 in bbec89a

	result = actx.empty(group_starts[-1], dtype=ary.entry_dtype)
	for grp_start, grp_ary in zip(group_starts, ary):
	actx.call_loopy(prg(),
	grp_ary=grp_ary,
	result=result,
	grp_start=grp_start)

What was originally proposed in #192 (which I really liked 😞) was:

def _flatten_dof_array(ary: Any, strict: bool = True):
    if not isinstance(ary, DOFArray):
        if strict:
            raise TypeError(f"non-DOFArray type '{type(ary).__name__}' cannot "
                    "be flattened; use 'strict=False' to allow other types")
        else:
            return ary

    actx = ary.array_context
    if actx is None:
        raise ValueError("cannot flatten frozen DOFArrays")

    return actx.np.concatenate([actx.np.reshape(grp_ary, -1)
                                for grp_ary in ary])

But, turns out this isn't what we want to do since we want to maintain flexibility with regards to the memory layout of the array (x-ref: #192 (comment)).

I would appreciate some guidance on how to proceed forward now that #196 is in. cc @alexfikl, @inducer

[inducer]

The only difference is that the flatten kernel needs to stick and be used in place of the reshape. (The reshape can be used if the strides are right, as with the code in #154).