Skip to content

Commit

Permalink
Update broadcast.jl
Browse files Browse the repository at this point in the history
  • Loading branch information
N5N3 committed Feb 20, 2022
1 parent 6e47ff5 commit 6732699
Showing 1 changed file with 9 additions and 7 deletions.
16 changes: 9 additions & 7 deletions src/host/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,15 @@ end
copyto!(similar(bc, ElType), bc)
end

@inline function Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{Nothing})
@inline function materialize!(::Style, dest, bc::Broadcasted) where {Style<:AbstractGPUArrayStyle}
return _copyto!(dest, instantiate(Broadcasted{Style}(bc.f, bc.args, axes(dest))))
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

0 comments on commit 6732699

Please sign in to comment.