Adapt Type call broadcasting to a function (#2000)

This is a more generic solution to the existing `broadcasted` definition that fixes #1761 (as suggested in #1761 (comment)).
JuliaGPU · Jul 17, 2023 · 583329d · 583329d
1 parent 0341a35
commit 583329d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/src/broadcast.jl b/src/broadcast.jl
@@ -13,7 +13,3 @@ Base.similar(bc::Broadcasted{CuArrayStyle{N}}, ::Type{T}) where {N,T} =
 
 Base.similar(bc::Broadcasted{CuArrayStyle{N}}, ::Type{T}, dims) where {N,T} =
     CuArray{T}(undef, dims)
-
-# broadcasting type ctors isn't GPU compatible
-Broadcast.broadcasted(::CuArrayStyle{N}, f::Type{T}, args...) where {N, T} =
-    Broadcasted{CuArrayStyle{N}}((x...) -> T(x...), args, nothing)
diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
@@ -140,6 +140,10 @@ 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[]}()
 
+# case where type is the function being broadcasted
+Adapt.adapt_structure(to::Adaptor, bc::Base.Broadcast.Broadcasted{Style, <:Any, Type{T}}) where {Style, T} = 
+    Base.Broadcast.Broadcasted{Style}((x...) -> T(x...), adapt(to, bc.args), bc.axes)
+
 Adapt.adapt_storage(::Adaptor, xs::CuArray{T,N}) where {T,N} =
   Base.unsafe_convert(CuDeviceArray{T,N,AS.Global}, xs)
 

diff --git a/test/base/broadcast.jl b/test/base/broadcast.jl
@@ -39,3 +39,11 @@ end
   A = CuArray{ComplexF64}(undef, (2,2))
   @test eltype(convert.(ComplexF32, A)) == ComplexF32
 end
+
+# https://github.com/JuliaGPU/CUDA.jl/issues/1761
+@testset "Broadcast Type(args)" begin
+  A = CuArray{ComplexF64}(undef, (2,2))
+  @test eltype(ComplexF32.(A)) == ComplexF32
+  @test eltype(A .+ ComplexF32.(1)) == ComplexF64
+  @test eltype(ComplexF32.(A) .+ ComplexF32.(1)) == ComplexF32
+end