Specialize Ref{<:Type} for GPU compatibility. (#1109)

Author: maleadt

src/compiler/execution.jl

@@ -123,9 +123,15 @@ Adapt.adapt_storage(to::Adaptor, p::CuPtr{T}) where {T} = reinterpret(LLVMPtr{T,
 struct CuRefValue{T} <: Ref{T}
   x::T
 end
-Base.getindex(r::CuRefValue) = r.x
+Base.getindex(r::CuRefValue{T}) where T = r.x
 Adapt.adapt_structure(to::Adaptor, r::Base.RefValue) = CuRefValue(adapt(to, r[]))
 
+# broadcast sometimes passes a ref(type), resulting in a GPU-incompatible DataType box.
+# avoid that by using a special kind of ref that knows about the boxed type.
+struct CuRefType{T} <: Ref{DataType} end
+Base.getindex(r::CuRefType{T}) where T = T
+Adapt.adapt_structure(to::Adaptor, r::Base.RefValue{<:Union{DataType,Type}}) = CuRefType{r[]}()
+
 Adapt.adapt_storage(::Adaptor, xs::CuArray{T,N}) where {T,N} =
   Base.unsafe_convert(CuDeviceArray{T,N,AS.Global}, xs)
 

test/broadcast.jl

@@ -22,6 +22,7 @@
   end
   @test Array(Whatever{Int}.(CuArray([1]))) == Whatever{Int}.([1])
 end
+
 # https://github.com/JuliaGPU/CUDA.jl/issues/223
 @testset "Ref Broadcast" begin
   foobar(idx, A) = A[idx]
@@ -32,3 +33,9 @@ end
   @test testf(x -> log.(x), rand(3,3))
   @test testf((x,xs) -> log.(x.+xs), Ref(1), rand(3,3))
 end
+
+# https://github.com/JuliaGPU/CUDA.jl/issues/261
+@testset "Broadcast Ref{<:Type}" begin
+  A = CuArray{ComplexF64}(undef, (2,2))
+  @test eltype(convert.(ComplexF32, A)) == ComplexF32
+end