Adapt to broadcast changes in GPUArrays.

Author: maleadt

Manifest.toml

@@ -40,9 +40,11 @@ version = "5.0.1"
 
 [[CUDAnative]]
 deps = ["Adapt", "CEnum", "CUDAapi", "CUDAdrv", "DataStructures", "InteractiveUtils", "LLVM", "Libdl", "Printf", "TimerOutputs"]
-git-tree-sha1 = "e0c2805c9a7d338823c0d8f574242e284410fa61"
+git-tree-sha1 = "f38cf81a0c6c08cb532b7324eb85f3e2d438bed7"
+repo-rev = "0f65bc6ebc9c490c068d65436dc42b93c7650fce"
+repo-url = "https://github.com/JuliaGPU/CUDAnative.jl.git"
 uuid = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17"
-version = "2.9.1"
+version = "2.10.0"
 
 [[DataStructures]]
 deps = ["InteractiveUtils", "OrderedCollections"]
@@ -60,8 +62,8 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
 [[GPUArrays]]
 deps = ["AbstractFFTs", "Adapt", "LinearAlgebra", "Printf", "Random", "Serialization"]
-git-tree-sha1 = "64bf1d9f7634b35e1383955213cafbda3a9ac184"
-repo-rev = "a89fa3c84e97488f88449036d6a860db689669f2"
+git-tree-sha1 = "1bc994886f5c6fa82e60da92debfc9be79400993"
+repo-rev = "cc7bb3030ed18c8e545539c540b675171ff1b73e"
 repo-url = "https://github.com/JuliaGPU/GPUArrays.jl.git"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 version = "2.0.1"

src/broadcast.jl

@@ -1,25 +1,29 @@
-import Base.Broadcast: Broadcasted, Extruded, BroadcastStyle, ArrayStyle
+# broadcasting
 
-BroadcastStyle(::Type{<:CuArray}) = ArrayStyle{CuArray}()
+using Base.Broadcast: BroadcastStyle, Broadcasted
 
-function Base.similar(bc::Broadcasted{ArrayStyle{CuArray}}, ::Type{T}) where T
+struct CuArrayStyle{N} <: AbstractGPUArrayStyle{N} end
+CuArrayStyle(::Val{N}) where N = CuArrayStyle{N}()
+CuArrayStyle{M}(::Val{N}) where {N,M} = CuArrayStyle{N}()
+
+BroadcastStyle(::Type{<:CuArray{T,N}}) where {T,N} = CuArrayStyle{N}()
+
+Base.similar(bc::Broadcasted{CuArrayStyle{N}}, ::Type{T}) where {N,T} =
     similar(CuArray{T}, axes(bc))
-end
 
-function Base.similar(bc::Broadcasted{ArrayStyle{CuArray}}, ::Type{T}, dims...) where {T}
-    similar(CuArray{T}, dims...)
-end
+Base.similar(bc::Broadcasted{CuArrayStyle{N}}, ::Type{T}, dims...) where {N,T} =
+    CuArray{T}(undef, dims...)
+
 
-# replace base functions with libdevice alternatives
-# TODO: do this with Cassette.jl
+## replace base functions with libdevice alternatives
 
 cufunc(f) = f
 cufunc(::Type{T}) where T = (x...) -> T(x...) # broadcasting type ctors isn't GPU compatible
 
-Broadcast.broadcasted(::ArrayStyle{CuArray}, f, args...) =
-  Broadcasted{ArrayStyle{CuArray}}(cufunc(f), args, nothing)
+Broadcast.broadcasted(::CuArrayStyle{N}, f, args...) where {N} =
+  Broadcasted{CuArrayStyle{N}}(cufunc(f), args, nothing)
 
-libdevice = :[
+const libdevice = :[
   cos, cospi, sin, sinpi, tan, acos, asin, atan,
   cosh, sinh, tanh, acosh, asinh, atanh,
   log, log10, log1p, log2, logb, ilogb,
@@ -40,7 +44,8 @@ for f in libdevice
   @eval cufunc(::typeof(Base.$f)) = CUDAnative.$f
 end
 
-#broadcast ^
+# broadcast ^
+
 culiteral_pow(::typeof(^), x::T, ::Val{0}) where {T<:Real} = one(x)
 culiteral_pow(::typeof(^), x::T, ::Val{1}) where {T<:Real} = x
 culiteral_pow(::typeof(^), x::T, ::Val{2}) where {T<:Real} = x * x

test/base.jl

@@ -170,16 +170,6 @@ end
   @test testf(x -> reduce(|, x, init=false), map(t -> t > 0.5, cu(rand(2, 3))))
 end
 
-@testset "0D" begin
-  x = CuArray{Float64}(undef)
-  x .= 1
-  @test collect(x)[] == 1
-  # broken test that throws
-  # https://github.com/JuliaGPU/GPUArrays.jl/issues/204
-  @test_throws ErrorException x /= 2
-  #@test collect(x)[] == 0.5
-end
-
 @testset "SubArray" begin
   @test testf(rand(5)) do x
     y = x[2:4]