Handle broadcasting when storage types are different (#605)

Co-authored-by: Christian Guinard <[email protected]>

Author: limarta

src/broadcast.jl

@@ -12,9 +12,14 @@ BroadcastStyle(W::Type{<:WrappedMtlArray{T,N}}) where {T,N} =
 
 # when we are dealing with different buffer styles, we cannot know
 # which one is better, so use shared memory
-BroadcastStyle(::MtlArrayStyle{N, S1},
-               ::MtlArrayStyle{N, S2}) where {N,S1,S2} =
-    MtlArrayStyle{N, SharedStorage}()
+BroadcastStyle(::MtlArrayStyle{N1, S1},
+               ::MtlArrayStyle{N2, S2}) where {N1,N2,S1,S2} =
+    MtlArrayStyle{max(N1, N2), SharedStorage}()
+
+# resolve ambiguity: different N, same memory type
+BroadcastStyle(::MtlArrayStyle{N1, S},
+               ::MtlArrayStyle{N2, S}) where {N1,N2,S} =
+    MtlArrayStyle{max(N1, N2), S}()
 
 # allocation of output arrays
 Base.similar(::Broadcasted{MtlArrayStyle{N, S}}, ::Type{T}, dims) where {T, N, S} =

test/array.jl

@@ -557,6 +557,24 @@ end
     @test testf(x->min.(x, one(Float32)), randn(Float32, 1000))
     @test testf(x->min.(max.(x, zero(Float32)), one(Float32)), randn(Float32, 1000))
     @test testf(x->max.(min.(x, one(Float32)), zero(Float32)), randn(Float32, 1000))
+
+    # preserving buffer types
+    let x = Metal.zeros(Float32, 1; storage=Metal.SharedStorage)
+        y = x .+ 1
+        @test is_shared(y)
+    end
+
+    # when storages are different, choose shared
+    let x = Metal.zeros(Float32, 1; storage=Metal.SharedStorage), y = Metal.zeros(Float32, 1; storage=Metal.PrivateStorage)
+        z = x .+ y
+        @test is_shared(z)
+    end
+
+    let x = Metal.zeros(Float32, 2, 2; storage=Metal.SharedStorage), y = Metal.zeros(Float32, 2; storage=Metal.PrivateStorage)
+        z = x .+ y
+        @test is_shared(z)
+    end
 end
 
+
 end