Merge pull request #240 from JuliaGPU/tb/painful_changes

Fixes to broadcast and indexing

Author: maleadt

src/host/broadcast.jl

@@ -1,23 +1,21 @@
 # broadcasting operations
 
+export AbstractGPUArrayStyle
+
 using Base.Broadcast
 
-import Base.Broadcast: BroadcastStyle, Broadcasted, ArrayStyle
+import Base.Broadcast: BroadcastStyle, Broadcasted, AbstractArrayStyle
 
-# we define a generic `BroadcastStyle` here that should be sufficient for most cases.
-# dependent packages like `CuArrays` can define their own `BroadcastStyle` allowing
-# them to further change or optimize broadcasting.
-#
-# TODO: investigate if we should define out own `GPUArrayStyle{N} <: AbstractArrayStyle{N}`
-#
-# NOTE: this uses the specific `T` that was used e.g. `JLArray` or `CLArray` for ArrayStyle,
-#       instead of using `ArrayStyle{AbstractGPUArray}`, due to the fact how `similar` works.
-BroadcastStyle(::Type{T}) where {T<:AbstractGPUArray} = ArrayStyle{T}()
+"""
+Abstract supertype for GPU array styles. The `N` parameter is the dimensionality.
+
+Downstream implementations should provide a concrete array style type that inherits from
+this supertype.
+"""
+abstract type AbstractGPUArrayStyle{N} <: AbstractArrayStyle{N} end
 
 # Wrapper types otherwise forget that they are GPU compatible
-#
-# NOTE: Don't directly use ArrayStyle{AbstractGPUArray} here since that would mean that `CuArrays`
-#       customization no longer take effect.
+# NOTE: don't directly use GPUArrayStyle here not to lose downstream customizations.
 for (W, ctor) in Adapt.wrappers
   @eval begin
     BroadcastStyle(::Type{<:$W}) where {AT<:AbstractGPUArray} = BroadcastStyle(AT)
@@ -28,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:
@@ -49,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/host/indexing.jl

@@ -1,6 +1,6 @@
 # host-level indexing
 
-export allowscalar, @allowscalar, assertscalar
+export allowscalar, @allowscalar, @disallowscalar, assertscalar
 
 
 # mechanism to disallow scalar operations
@@ -82,26 +82,18 @@ end
 
 Base.IndexStyle(::Type{<:AbstractGPUArray}) = Base.IndexLinear()
 
-function _getindex(xs::AbstractGPUArray{T}, i::Integer) where T
+function Base.getindex(xs::AbstractGPUArray{T}, i::Integer) where T
+    ndims(xs) > 0 && assertscalar("scalar getindex")
     x = Array{T}(undef, 1)
     copyto!(x, 1, xs, i, 1)
     return x[1]
 end
 
-function Base.getindex(xs::AbstractGPUArray{T}, i::Integer) where T
-    ndims(xs) > 0 && assertscalar("scalar getindex")
-    _getindex(xs, i)
-end
-
-function _setindex!(xs::AbstractGPUArray{T}, v::T, i::Integer) where T
-    x = T[v]
-    copyto!(xs, i, x, 1, 1)
-    return v
-end
-
 function Base.setindex!(xs::AbstractGPUArray{T}, v::T, i::Integer) where T
     assertscalar("scalar setindex!")
-    _setindex!(xs, v, i)
+    x = T[v]
+    copyto!(xs, i, x, 1, 1)
+    return xs
 end
 
 Base.setindex!(xs::AbstractGPUArray, v, i::Integer) = xs[i] = convert(eltype(xs), v)

src/host/mapreduce.jl

@@ -57,8 +57,8 @@ gpu_promote_type(::typeof(min), ::Type{T}) where {T<: WidenReduceResult} = T
 gpu_promote_type(::typeof(abs), ::Type{Complex{T}}) where {T} = T
 gpu_promote_type(::typeof(abs2), ::Type{Complex{T}}) where {T} = T
 
-import Base.Broadcast: Broadcasted, ArrayStyle
-const GPUSrcArray = Union{Broadcasted{ArrayStyle{AT}}, AbstractGPUArray{T, N}} where {T, N, AT<:AbstractGPUArray}
+import Base.Broadcast: Broadcasted
+const GPUSrcArray = Union{Broadcasted{<:AbstractGPUArrayStyle}, <:AbstractGPUArray}
 
 function Base.mapreduce(f::Function, op::Function, A::GPUSrcArray; dims = :, init...)
     mapreduce_impl(f, op, init.data, A, dims)

src/host/quirks.jl

@@ -22,6 +22,6 @@ if VERSION >= v"1.3.0-alpha.107"
     @inline combine_axes(A, B...) = broadcast_shape(axes(A), combine_axes(B...))
     combine_axes(A) = axes(A)
 
-    Broadcast._axes(::Broadcasted{ArrayStyle{AT}}, axes::Tuple) where {AT <: AbstractGPUArray} = axes
-    @inline Broadcast._axes(bc::Broadcasted{ArrayStyle{AT}}, ::Nothing) where {AT <: AbstractGPUArray} = combine_axes(bc.args...)
+    Broadcast._axes(::Broadcasted{<:AbstractGPUArrayStyle}, axes::Tuple) = axes
+    @inline Broadcast._axes(bc::Broadcasted{<:AbstractGPUArrayStyle}, ::Nothing) = combine_axes(bc.args...)
 end

src/reference.jl

@@ -5,9 +5,11 @@
 
 module JLArrays
 
+export JLArray
+
 using GPUArrays
 
-export JLArray
+using Adapt
 
 
 #
@@ -52,14 +54,21 @@ function JLKernelContext(ctx::JLKernelContext, threadidx::Int)
     )
 end
 
-to_device(ctx, x::Tuple) = to_device.(Ref(ctx), x)
-to_device(ctx, x) = x
+struct Adaptor end
+jlconvert(arg) = adapt(Adaptor(), arg)
+
+# FIXME: add Ref to Adapt.jl (but make sure it doesn't cause ambiguities with CUDAnative's)
+struct JlRefValue{T} <: Ref{T}
+  x::T
+end
+Base.getindex(r::JlRefValue) = r.x
+Adapt.adapt_structure(to::Adaptor, r::Base.RefValue) = JlRefValue(adapt(to, r[]))
 
 function GPUArrays.gpu_call(::JLBackend, f, args...; blocks::Int, threads::Int)
     ctx = JLKernelContext(threads, blocks)
-    device_args = to_device.(Ref(ctx), args)
+    device_args = jlconvert.(args)
     tasks = Array{Task}(undef, threads)
-    @allowscalar for blockidx in 1:blocks
+    @disallowscalar for blockidx in 1:blocks
         ctx.blockidx = blockidx
         for threadidx in 1:threads
             thread_ctx = JLKernelContext(ctx, threadidx)
@@ -138,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
@@ -193,15 +203,19 @@ Base.convert(::Type{T}, x::T) where T <: JLArray = x
 
 ## broadcast
 
-using Base.Broadcast: BroadcastStyle, Broadcasted, ArrayStyle
+using Base.Broadcast: BroadcastStyle, Broadcasted
+
+struct JLArrayStyle{N} <: AbstractGPUArrayStyle{N} end
+JLArrayStyle(::Val{N}) where N = JLArrayStyle{N}()
+JLArrayStyle{M}(::Val{N}) where {N,M} = JLArrayStyle{N}()
 
-BroadcastStyle(::Type{<:JLArray}) = ArrayStyle{JLArray}()
+BroadcastStyle(::Type{JLArray{T,N}}) where {T,N} = JLArrayStyle{N}()
 
-function Base.similar(bc::Broadcasted{ArrayStyle{JLArray}}, ::Type{T}) where T
+Base.similar(bc::Broadcasted{JLArrayStyle{N}}, ::Type{T}) where {N,T} =
     similar(JLArray{T}, axes(bc))
-end
 
-Base.similar(bc::Broadcasted{ArrayStyle{JLArray}}, ::Type{T}, dims...) where {T} = JLArray{T}(undef, dims...)
+Base.similar(bc::Broadcasted{JLArrayStyle{N}}, ::Type{T}, dims...) where {N,T} =
+    JLArray{T}(undef, dims...)
 
 
 ## memory operations
@@ -263,8 +277,8 @@ GPUArrays.device(x::JLArray) = JLDevice()
 
 GPUArrays.backend(::Type{<:JLArray}) = JLBackend()
 
-to_device(ctx, x::JLArray{T,N}) where {T,N} = JLDeviceArray{T,N}(x.data, x.dims)
-to_device(ctx, x::Base.RefValue{<: JLArray}) = Base.RefValue(to_device(ctx, x[]))
+Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N} =
+  JLDeviceArray{T,N}(x.data, x.dims)
 
 GPUArrays.unsafe_reinterpret(::Type{T}, A::JLArray, size::Tuple) where T =
     reshape(reinterpret(T, A.data), size)

test/testsuite/broadcasting.jl

@@ -118,6 +118,29 @@ function broadcasting(AT)
             @test compare((A, B) -> A .* B .+ ET(10), AT, rand(ET, 40, 40), rand(ET, 40, 40))
         end
     end
+
+    @testset "0D" begin
+        x = AT{Float64}(undef)
+        x .= 1
+        @test collect(x)[] == 1
+        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)

test/testsuite/indexing.jl

@@ -62,5 +62,11 @@ function test_indexing(AT)
                 @test Array(A) == Ac
             end
         end
+
+        @testset "get/setindex!" begin
+            # literal calls to get/setindex! have differen return types
+            @test compare(x->getindex(x,1), AT, zeros(Int, 2))
+            @test compare(x->setindex!(x,1,1), AT, zeros(Int, 2))
+        end
     end
 end