Fix mixed-buffer/mixed-shape broadcasts. (#2290)

Author: maleadt

src/broadcast.jl

@@ -12,9 +12,9 @@ BroadcastStyle(W::Type{<:AnyCuArray{T,N}}) where {T,N} =
 
 # when we are dealing with different buffer styles, we cannot know
 # which one is better, so use unified memory
-BroadcastStyle(::CUDA.CuArrayStyle{N, B1},
-               ::CUDA.CuArrayStyle{N, B2}) where {N,B1,B2} =
-    CuArrayStyle{N, Mem.Unified}()
+BroadcastStyle(::CUDA.CuArrayStyle{N1, B1},
+               ::CUDA.CuArrayStyle{N2, B2}) where {N1,N2,B1,B2} =
+    CuArrayStyle{max(N1,N2), Mem.Unified}()
 
 # allocation of output arrays
 Base.similar(bc::Broadcasted{CuArrayStyle{N,B}}, ::Type{T}, dims) where {T,N,B} =

test/base/broadcast.jl

@@ -62,4 +62,10 @@ end
   d = cu([1]; host=true)
   e = c .+ d
   @test is_unified(e)
+
+  # this should also work with differently-sized inputs
+  f = cu([1]; device=true)
+  g = cu([1 2]; host=true)
+  h = f .+ g
+  @test is_unified(h)
 end