Handle broadcasting Ref.

Author: maleadt

src/host/broadcast.jl

@@ -26,7 +26,22 @@ end
 # This Union is a hack. Ideally Base would have a Transpose <: WrappedArray <: AbstractArray
 # and we could define our methods in terms of Union{AbstractGPUArray, WrappedArray{<:Any, <:AbstractGPUArray}}
 @eval const GPUDestArray =
-  Union{AbstractGPUArray, $((:($W where {AT <: AbstractGPUArray}) for (W, _) in Adapt.wrappers)...)}
+  Union{AbstractGPUArray,
+        $((:($W where {AT <: AbstractGPUArray}) for (W, _) in Adapt.wrappers)...),
+        Base.RefValue{<:AbstractGPUArray} }
+
+# Ref is special: it's not a real wrapper, so not part of Adapt,
+# but it is commonly used to bypass broadcasting of an argument
+# so we need to preserve its dimensionless properties.
+BroadcastStyle(::Type{Base.RefValue{AT}}) where {AT<:AbstractGPUArray} = typeof(BroadcastStyle(AT))(Val(0))
+backend(::Type{Base.RefValue{AT}}) where {AT<:AbstractGPUArray} = backend(AT)
+# but make sure we don't dispatch to the optimized copy method that directly indexes
+function Broadcast.copy(bc::Broadcasted{<:AbstractGPUArrayStyle{0}})
+    ElType = Broadcast.combine_eltypes(bc.f, bc.args)
+    isbitstype(ElType) || error("Cannot broadcast function returning non-isbits $ElType.")
+    dest = copyto!(similar(bc, ElType), bc)
+    return @allowscalar dest[CartesianIndex()]  # 0D broadcast needs to unwrap results
+end
 
 # We purposefully only specialize `copyto!`, dependent packages need to make sure that they
 # can handle:
@@ -47,7 +62,15 @@ end
     bc′ = Broadcast.preprocess(dest, bc)
     gpu_call(dest, bc′) do ctx, dest, bc′
         let I = CartesianIndex(@cartesianidx(dest))
-            @inbounds dest[I] = bc′[I]
+            #@inbounds dest[I] = bc′[I]
+            @inbounds let
+                val = bc′[I]
+                if val !== nothing
+                  # FIXME: CuArrays.jl crashes on assigning Nothing (this happens with
+                  #        broadcasts that don't return anything but assign anyway)
+                  dest[I] = val
+                end
+            end
         end
         return
     end

src/reference.jl

@@ -147,6 +147,7 @@ struct JLArray{T, N} <: AbstractGPUArray{T, N}
     dims::Dims{N}
 
     function JLArray{T,N}(data::Array{T, N}, dims::Dims{N}) where {T,N}
+        @assert isbitstype(T) "JLArray only supports bits types"
         new(data, dims)
     end
 end

test/testsuite/broadcasting.jl

@@ -126,6 +126,21 @@ function broadcasting(AT)
         x /= 2
         @test collect(x)[] == 0.5
     end
+
+    @testset "Ref" begin
+        # as first arg, 0d broadcast
+        @test compare(x->getindex.(Ref(x),1), AT, [0])
+
+        void_setindex!(args...) = (setindex!(args...); return)
+        @test compare(x->(void_setindex!.(Ref(x),1); x), AT, [0])
+
+        # regular broadcast
+        a = AT(rand(10))
+        b = AT(rand(10))
+        cpy(i,a,b) = (a[i] = b[i]; return)
+        cpy.(1:10, Ref(a), Ref(b))
+        @test Array(a) == Array(b)
+    end
 end
 
 function vec3(AT)