From 39078aa243e11c59ef13e7386b34b9d25767d68d Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 23 Jun 2020 17:25:45 -0400 Subject: [PATCH] use style dispatch in materialize --- src/host/broadcast.jl | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/src/host/broadcast.jl b/src/host/broadcast.jl index f007c583..16e728ce 100644 --- a/src/host/broadcast.jl +++ b/src/host/broadcast.jl @@ -51,6 +51,21 @@ end # # For more information see the Base documentation. @inline function Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{Nothing}) + __copyto!(dest, bc) +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)) + +if VERSION >= v"1.5.0-DEV.613" + @inline function Base.materialize!(::AbstractGPUArrayStyle, dest, bc::Broadcasted{Style}) where Style + __copyto!(dest, convert(Broadcasted{Nothing}, bc)) + end +end + +@inline function __copyto!(dest, bc) axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc)) isempty(dest) && return dest bc′ = Broadcast.preprocess(dest, bc) @@ -64,12 +79,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