Add an AbstractGPUArrayStyle that preserves dimensionality.

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)

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

@@ -193,15 +193,19 @@ Base.convert(::Type{T}, x::T) where T <: JLArray = x
 
 ## broadcast
 
-using Base.Broadcast: BroadcastStyle, Broadcasted, ArrayStyle
+using Base.Broadcast: BroadcastStyle, Broadcasted
 
-BroadcastStyle(::Type{<:JLArray}) = ArrayStyle{JLArray}()
+struct JLArrayStyle{N} <: AbstractGPUArrayStyle{N} end
+JLArrayStyle(::Val{N}) where N = JLArrayStyle{N}()
+JLArrayStyle{M}(::Val{N}) where {N,M} = JLArrayStyle{N}()
 
-function Base.similar(bc::Broadcasted{ArrayStyle{JLArray}}, ::Type{T}) where T
+BroadcastStyle(::Type{JLArray{T,N}}) where {T,N} = JLArrayStyle{N}()
+
+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

test/testsuite/broadcasting.jl

@@ -118,6 +118,14 @@ 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
 end
 
 function vec3(AT)