Skip to content

Use style dispatch in broadcast(!) #393

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.

Sign up for GitHub

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

Merged
merged 3 commits into from
Mar 15, 2022
Merged

Use style dispatch in broadcast(!) #393

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b53148e
Use style to dispatch Broadcast
N5N3 Feb 20, 2022
bee0d4f
Make sure dispatch with `dest`'s style work.
N5N3 Feb 21, 2022
6bf6326
Test fix and add more comments.
N5N3 Feb 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 10 additions & 8 deletions src/host/broadcast.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ export AbstractGPUArrayStyle

using Base.Broadcast

import Base.Broadcast: BroadcastStyle, Broadcasted, AbstractArrayStyle
import Base.Broadcast: BroadcastStyle, Broadcasted, AbstractArrayStyle, instantiate

const BroadcastGPUArray{T} = Union{AnyGPUArray{T},
Base.RefValue{<:AbstractGPUArray{T}}}
Expand Down Expand Up @@ -47,7 +47,15 @@ end
copyto!(similar(bc, ElType), bc)
end

@inline function Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{Nothing})
@inline function Base.materialize!(::Style, dest, bc::Broadcasted) where {Style<:AbstractGPUArrayStyle}
return _copyto!(dest, instantiate(Broadcasted{Style}(bc.f, bc.args, axes(dest))))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this, and the original Base implementation at https://github.com/JuliaLang/julia/blob/17775537232de5b12a38b257e402756b6df55386/base/broadcast.jl#L871, creates a Broadcasted{<:AbstractGPUArrayStyle}, why doesn't the copyto!(::AbstractArray, Broadcasted{<:AbstractGPUArrayStyle}) implementation below suffice?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Base won't create a Broadcasted{<:AbstractGPUArrayStyle}, bc's stype is kept there.
So we need to create Broadcasted{<:AbstractGPUArrayStyle} ourself.

end

@inline Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{Nothing}) = _copyto!(dest, bc) # Keep it for ArrayConflict

@inline Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractGPUArrayStyle}) = _copyto!(dest, bc)

@inline function _copyto!(dest::AbstractArray, bc::Broadcasted)
axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
isempty(dest) && return dest
bc′ = Broadcast.preprocess(dest, bc)
Expand All @@ -72,12 +80,6 @@ end
return dest
end

# Base defines this method as a performance optimization, but we don't know how to do
# `fill!` in general for all `BroadcastGPUArray` so we just go straight to the fallback
@inline Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{<:Broadcast.AbstractArrayStyle{0}}) =
copyto!(dest, convert(Broadcasted{Nothing}, bc))


## map

allequal(x) = true
Expand Down
32 changes: 32 additions & 0 deletions test/testsuite/broadcasting.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
@testsuite "broadcasting" (AT, eltypes)->begin
broadcasting(AT, eltypes)
vec3(AT, eltypes)
unknown_wrapper(AT, eltypes)

@testset "type instabilities" begin
f(x) = x ? 1.0 : 0
Expand Down Expand Up @@ -205,3 +206,34 @@ function vec3(AT, eltypes)
@test all(map((a,b)-> all((1,2,3) .≈ (1,2,3)), Array(res2), res2c))
end
end

# A help struct to test style-based broadcast dispatch with unknown array wrapper.
# `WrapArray(A)` behaves like `A` during broadcast. But its not a `BroadcastGPUArray`.
struct WrapArray{T,N,P<:AbstractArray{T,N}} <: AbstractArray{T,N}
data::P
end
Base.@propagate_inbounds Base.getindex(A::WrapArray, i::Integer...) = A.data[i...]
Base.@propagate_inbounds Base.setindex!(A::WrapArray, v::Any, i::Integer...) = setindex!(A.data, v, i...)
Base.size(A::WrapArray) = size(A.data)
# For kernal support
Adapt.adapt_structure(to, s::WrapArray) = WrapArray(Adapt.adapt(to, s.data))
# For broadcast support
GPUArrays.backend(::Type{WrapArray{T,N,P}}) where {T,N,P} = GPUArrays.backend(P)
Broadcast.BroadcastStyle(::Type{WrapArray{T,N,P}}) where {T,N,P} = Broadcast.BroadcastStyle(P)

function unknown_wrapper(AT, eltypes)
for ET in eltypes
@views @testset "unknown wrapper $ET" begin
A = AT(rand(ET, 10, 10))
WA = WrapArray(A)
# test for dispatch with src's BroadcastStyle.
@test Array(WA .+ ET(1)) == Array(A .+ ET(1))
@test Array(WA .+ WA) == Array(WA .+ A) == Array(A .+ A)
@test Array(WA .+ A[:,1]) == Array(A .+ A[:,1])
@test Array(WA .+ A[1,:]) == Array(A .+ A[1,:])
# test for dispatch with dest's BroadcastStyle.
WA .= ET(1)
@test all(isequal(ET(1)), Array(A))
end
end
end