Fix broadcast defaulting to Mem.Unified() (#2327)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: vpuri3

src/broadcast.jl

@@ -16,6 +16,11 @@ BroadcastStyle(::CUDA.CuArrayStyle{N1, B1},
                ::CUDA.CuArrayStyle{N2, B2}) where {N1,N2,B1,B2} =
     CuArrayStyle{max(N1,N2), Mem.Unified}()
 
+# resolve ambiguity: different N, same buffer
+BroadcastStyle(::CUDA.CuArrayStyle{N1, B},
+               ::CUDA.CuArrayStyle{N2, B}) where {N1,N2,B} =
+    CuArrayStyle{max(N1,N2), B}()
+
 # allocation of output arrays
 Base.similar(bc::Broadcasted{CuArrayStyle{N,B}}, ::Type{T}, dims) where {T,N,B} =
     similar(CuArray{T,length(dims),B}, dims)

test/base/broadcast.jl

@@ -68,4 +68,10 @@ end
   g = cu([1 2]; host=true)
   h = f .+ g
   @test is_unified(h)
+
+  # however, differences in only shape shouldn't change the buffer type
+  i = cu([1]; device=true)
+  j = cu([1 2]; device=true)
+  k = i .+ j
+  @test !is_unified(k)
 end