Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix inference for column broadcasting #1984

Merged
merged 2 commits into from
Sep 12, 2024
Merged

Conversation

charleskawczynski
Copy link
Member

@charleskawczynski charleskawczynski commented Sep 12, 2024

It turns out that what was missing here is that column on broadcasted objects passes nothing to the axes. So, even though we call instantiate on the broadcasted objects in the Fields broadcasting:

@inline function Base.copyto!(
    dest::Field,
    bc::Base.Broadcast.Broadcasted{<:AbstractFieldStyle},
)
    copyto!(field_values(dest), Base.Broadcast.instantiate(todata(bc)))
    return dest
end

We later call column inside the datalayouts loop:

function Base.copyto!(
    dest::VIJFH{S, Nv, Nij, Nh},
    bc::BroadcastedUnionVIJFH{S, Nv, Nij, Nh},
    ::ToCPU,
) where {S, Nv, Nij, Nh}
    # copy contiguous columns
    @inbounds for h in 1:Nh, j in 1:Nij, i in 1:Nij
        col_dest = column(dest, i, j, h)
        col_bc = column(bc, i, j, h)
        copyto!(col_dest, col_bc)
    end
    return dest
end

Which then drops the axes and passes in nothing. This results in looking up the axes inside getindex rather than column.

I'm kind of surprised that this wasn't a more wide spread issue. Also, it oddly didn't seem to have a performance benefit, but at least it fixes inference and conceptually makes more sense this way.

Closes #1981.

@charleskawczynski charleskawczynski merged commit 1f51ca3 into main Sep 12, 2024
19 of 21 checks passed
@charleskawczynski charleskawczynski deleted the ck/fix_inference branch September 12, 2024 14:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Inference failure in broadcast expression
1 participant