Do think it is possible to shomehow also provide a two-argument fill and copyto? I think it would be very nice to do so because we would follow the same style as Base

Do think it is possible to shomehow also provide a two-argument fill and copyto? I think it would be very nice to do so because we would follow the same style as Base

The top-level methods are two-argument based:

Base.fill!(dest::AbstractData, val) = Base.fill!(dest, val, backed_array(dest))
Base.copyto!(
    dest::AbstractData,
    bc::Union{AbstractData, Base.Broadcast.Broadcasted},
) = Base.copyto!(dest, bc, backed_array(dest))

but they immediately call the 3-argument methods in order to properly dispatch.

Do think it is possible to shomehow also provide a two-argument fill and copyto? I think it would be very nice to do so because we would follow the same style as Base

The top-level methods are two-argument based:

Base.fill!(dest::AbstractData, val) = Base.fill!(dest, val, backed_array(dest))
Base.copyto!(
    dest::AbstractData,
    bc::Union{AbstractData, Base.Broadcast.Broadcasted},
) = Base.copyto!(dest, bc, backed_array(dest))

but they immediately call the 3-argument methods in order to properly dispatch.

Perfect! I missed this!

Ah, some of the backed_array calls lead to MArrays. I can't think of a good name for Union{Array, MArray}, so instead I've defined new types:

abstract type AbstractDispatchToDevice end
struct ToCPU <: AbstractDispatchToDevice end
struct ToCUDA <: AbstractDispatchToDevice end

and renamed backed_array to dispatch_device:

device_dispatch(dest::AbstractData) = _device_dispatch(dest)

_device_dispatch(x::Array) = ToCPU()
_device_dispatch(x::SubArray) = _device_dispatch(parent(x))
_device_dispatch(x::Base.ReshapedArray) = _device_dispatch(parent(x))
_device_dispatch(x::AbstractData) = _device_dispatch(parent(x))
_device_dispatch(x::SArray) = ToCPU()
_device_dispatch(x::MArray) = ToCPU()

And catch CPU/GPU code paths accordingly.

This is perhaps just a symptom of not having the ClimaComms context/device not stored in the datalayouts. But, I think doing that would be pretty breaking, so let's just improve the existing design for now, and we can tackle the comms context part separately if we deem it important enough.

Ah, interesting, there is a catastrophic case. On main:

function Base.copyto!(
    dest::VIFH{S, Nv, Ni},
    bc::BroadcastedUnionVIFH{S, Nv, Ni},
) where {S, Nv, Ni}
    (_, _, _, _, Nh) = size(bc)
    # copy contiguous columns
    @inbounds for h in 1:Nh, i in 1:Ni
        col_dest = column(dest, i, h)
        col_bc = column(bc, i, h)
        copyto!(col_dest, col_bc)
    end
    return dest
end

is executed for when running on the CPU and the GPU. It works because, on the CPU, we launch Nh*Ni kernel launches (column is defined on CPU and GPU for broadcasted objects), 1 per column!

We likely never run into this but, if we do / were to, this fix will be a major performance improvement.