Parallelize `materialize!(::FieldNameDict,::FieldNameDict)`

[charleskawczynski]

It's occurred to me that @. cache.b₂′ = b₂ - A₂₁ * x₁ lowers to this:

function Base.Broadcast.materialize!(
    dest::FieldNameDict,
    vector_or_matrix::FieldNameDict,
)
    !is_lazy(dest) || error("Cannot materialize into a lazy FieldNameDict")
    is_subset_that_covers_set(keys(vector_or_matrix), keys(dest)) || error(
        "Broadcast result and destination keys are incompatible: \
         $(set_string(keys(vector_or_matrix))) vs. $(set_string(keys(dest)))",
    ) # It is not always the case that keys(vector_or_matrix) == keys(dest).
    foreach(keys(vector_or_matrix)) do key
        entry = vector_or_matrix[key]
        if dest[key] isa UniformScaling
            dest[key] == entry || error("UniformScaling is immutable")
        elseif entry isa UniformScaling
            dest[key] .= (entry,)
        else
            dest[key] .= entry
        end
    end
end

So, there are even more kernel launches than I realized in run_field_matrix_solver!. We should fuse this loop into a single kernel launch, and parallelize if possible.

[charleskawczynski]

Another loop:

function Base.Broadcast.materialize(vector_or_matrix::FieldNameDict)
    entries = unrolled_map(values(vector_or_matrix)) do entry
        Base.Broadcast.materialize(entry)
    end
    return FieldNameDict(keys(vector_or_matrix), entries)
end

right above that definition. This needs to be reworked.

Actually, this is not an in-place operation, so it may only be used in constructing the cache/initialization. We can probably leave this one.

[charleskawczynski]

A close-up of ldiv! shows that most of field_matrix_solve! is composed of this loop (now contained in copyto_foreach! so that we can easily identify it with an nvtx annotation), and run_field_matrix_solver! on ::BlockDiagonalSolve, which launches the multiple_field_solve_kernels:

This snapshot is all on the device side, so the function calls line up with the cuda kernels.

Fusing these kernels could benefit us if there are reads to be hoisted, or if we can get better parallelism from the stencils (presumably the RHS of the solve).