From fdbb6448ff880260a0e28f7d0d340b1151d52d8a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 6 Mar 2024 13:02:54 +0100 Subject: [PATCH 1/5] Specialize broadcast to avoid integer divisions. By using hardware 2d/3d indices whenever possible. --- src/broadcast.jl | 88 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/src/broadcast.jl b/src/broadcast.jl index b9a95890a..f0dab6ced 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -19,3 +19,91 @@ BroadcastStyle(::MtlArrayStyle{N, S1}, # allocation of output arrays Base.similar(bc::Broadcasted{MtlArrayStyle{N,S}}, ::Type{T}, dims) where {T,N,S} = similar(MtlArray{T,length(dims),S}, dims) + +# specialization of the broadcast implementation to avoid expensive integer divisions +@inline function Base.materialize!(::Style, dest, bc::Broadcasted) where {Style<:MtlArrayStyle} + return _copyto!(dest, Broadcast.instantiate(Broadcasted{Style}(bc.f, bc.args, axes(dest)))) +end +@inline Base.copyto!(dest::MtlArray, bc::Broadcasted{Nothing}) = + _copyto!(dest, bc) # Keep it for ArrayConflict +@inline Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:MtlArrayStyle}) = + _copyto!(dest, bc) +@inline function _copyto!(dest::AbstractArray, bc::Broadcasted) + axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc)) + isempty(dest) && return dest + bc = Broadcast.preprocess(dest, bc) + + if ndims(dest) == 1 || + (isa(IndexStyle(dest), IndexLinear) && isa(IndexStyle(bc), IndexLinear)) + function broadcast_linear(dest, bc) + i = thread_position_in_grid_1d() + stride = threads_per_grid_1d() + while 1 <= i <= length(dest) + @inbounds dest[i] = bc[i] + i += stride + end + return + end + + kernel = @metal launch=false broadcast_linear(dest, bc) + elements = cld(length(dest), 4) + threads = min(elements, kernel.pipeline.maxTotalThreadsPerThreadgroup) + groups = cld(elements, threads) + elseif ndims(dest) == 2 + function broadcast_2d(dest, bc) + is = Tuple(thread_position_in_grid_2d()) + stride = threads_per_grid_2d() + while 1 <= is[1] <= size(dest, 1) && 1 <= is[2] <= size(dest, 2) + I = CartesianIndex(is) + @inbounds dest[I] = bc[I] + is = is .+ stride + end + return + end + + kernel = @metal launch=false broadcast_2d(dest, bc) + w = min(size(dest, 1), kernel.pipeline.threadExecutionWidth) + h = min(size(dest, 2), kernel.pipeline.maxTotalThreadsPerThreadgroup ÷ w) + threads = (w, h) + groups = cld.(size(dest), threads) + elseif ndims(dest) == 3 + function broadcast_3d(dest, bc) + is = Tuple(thread_position_in_grid_3d()) + stride = threads_per_grid_3d() + while 1 <= is[1] <= size(dest, 1) && + 1 <= is[2] <= size(dest, 2) && + 1 <= is[3] <= size(dest, 3) + I = CartesianIndex(is) + @inbounds dest[I] = bc[I] + is = is .+ stride + end + return + end + + kernel = @metal launch=false broadcast_3d(dest, bc) + w = min(size(dest, 1), kernel.pipeline.threadExecutionWidth) + h = min(size(dest, 2), kernel.pipeline.threadExecutionWidth) + d = min(size(dest, 3), kernel.pipeline.maxTotalThreadsPerThreadgroup ÷ (w*h)) + threads = (w, h, d) + groups = cld.(size(dest), threads) + else + function broadcast_cartesian(dest, bc) + i = thread_position_in_grid_1d() + stride = threads_per_grid_1d() + while 1 <= i <= length(dest) + I = @inbounds CartesianIndices(dest)[i] + @inbounds dest[I] = bc[I] + i += stride + end + return + end + + kernel = @metal launch=false broadcast_cartesian(dest, bc) + elements = cld(length(dest), 4) + threads = min(elements, kernel.pipeline.maxTotalThreadsPerThreadgroup) + groups = cld(elements, threads) + end + kernel(dest, bc; threads, groups) + + return dest +end From be5a7e8b46267ae4ffbb3654c14e5c24b8c39d9d Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 7 Mar 2024 09:48:34 +0100 Subject: [PATCH 2/5] Fully specialize the kernel when we commonly encounter a broadcast. --- src/broadcast.jl | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/src/broadcast.jl b/src/broadcast.jl index f0dab6ced..b945013d9 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -20,7 +20,25 @@ BroadcastStyle(::MtlArrayStyle{N, S1}, Base.similar(bc::Broadcasted{MtlArrayStyle{N,S}}, ::Type{T}, dims) where {T,N,S} = similar(MtlArray{T,length(dims),S}, dims) +# a static version of CartesianIndices that helps avoiding integer division +# (at the expense of additional compilation) +struct StaticCartesianIndices{N, I} end +StaticCartesianIndices(iter::CartesianIndices{N}) where {N} = + StaticCartesianIndices{N, iter.indices}() +StaticCartesianIndices(x) = StaticCartesianIndices(CartesianIndices(x)) +Base.CartesianIndices(iter::StaticCartesianIndices{N, I}) where {N, I} = + CartesianIndices{N, typeof(I)}(I) +Base.@propagate_inbounds Base.getindex(I::StaticCartesianIndices, i::Integer) = + CartesianIndices(I)[Int(i)] +Base.length(I::StaticCartesianIndices) = length(CartesianIndices(I)) +function Base.show(io::IO, I::StaticCartesianIndices) + print(io, "Static") + show(io, CartesianIndices(I)) +end + # specialization of the broadcast implementation to avoid expensive integer divisions +const _broadcast_shapes = Base.IdDict() +const BROADCAST_SPECIALIZATION_THRESHOLD = 10 @inline function Base.materialize!(::Style, dest, bc::Broadcasted) where {Style<:MtlArrayStyle} return _copyto!(dest, Broadcast.instantiate(Broadcasted{Style}(bc.f, bc.args, axes(dest)))) end @@ -33,6 +51,35 @@ end isempty(dest) && return dest bc = Broadcast.preprocess(dest, bc) + # if this is a common broadcast shape, specialize the kernel on it + Is = CartesianIndices(dest) + if !haskey(_broadcast_shapes, Is) + _broadcast_shapes[Is] = 1 + else + _broadcast_shapes[Is] += 1 + end + if _broadcast_shapes[Is] > BROADCAST_SPECIALIZATION_THRESHOLD + function broadcast_cartesian_static(dest, bc, Is) + i = thread_position_in_grid_1d() + stride = threads_per_grid_1d() + while 1 <= i <= length(dest) + I = @inbounds Is[i] + @inbounds dest[I] = bc[I] + i += stride + end + return + end + + Is = StaticCartesianIndices(Is) + kernel = @metal launch=false broadcast_cartesian_static(dest, bc, Is) + elements = cld(length(dest), 4) + threads = min(elements, kernel.pipeline.maxTotalThreadsPerThreadgroup) + groups = cld(elements, threads) + kernel(dest, bc, Is; threads, groups) + return dest + end + + # try to use the most appropriate hardware index to avoid integer division if ndims(dest) == 1 || (isa(IndexStyle(dest), IndexLinear) && isa(IndexStyle(bc), IndexLinear)) function broadcast_linear(dest, bc) From 7ba12fc7bf81bf3f4c04a1b09938933411d7b3cd Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 7 Mar 2024 09:48:45 +0100 Subject: [PATCH 3/5] Avoid unreachable argument errors. --- src/broadcast.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/broadcast.jl b/src/broadcast.jl index b945013d9..1c9701c1e 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -103,7 +103,7 @@ end while 1 <= is[1] <= size(dest, 1) && 1 <= is[2] <= size(dest, 2) I = CartesianIndex(is) @inbounds dest[I] = bc[I] - is = is .+ stride + is = (is[1] + stride[1], is[2] + stride[2]) end return end @@ -122,7 +122,7 @@ end 1 <= is[3] <= size(dest, 3) I = CartesianIndex(is) @inbounds dest[I] = bc[I] - is = is .+ stride + is = (is[1] + stride[1], is[2] + stride[2], is[3] + stride[3]) end return end From 396434abe422e288ca7eea136c92bbb475ff1d78 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 7 Mar 2024 10:33:34 +0100 Subject: [PATCH 4/5] Fix 2D launch configuration. --- src/broadcast.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/broadcast.jl b/src/broadcast.jl index 1c9701c1e..ed0eb734f 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -129,7 +129,8 @@ end kernel = @metal launch=false broadcast_3d(dest, bc) w = min(size(dest, 1), kernel.pipeline.threadExecutionWidth) - h = min(size(dest, 2), kernel.pipeline.threadExecutionWidth) + h = min(size(dest, 2), kernel.pipeline.threadExecutionWidth, + kernel.pipeline.maxTotalThreadsPerThreadgroup ÷ w) d = min(size(dest, 3), kernel.pipeline.maxTotalThreadsPerThreadgroup ÷ (w*h)) threads = (w, h, d) groups = cld.(size(dest), threads) From 0e0863e8f1f723acee7b239850d646f05800b5e1 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 8 Mar 2024 11:53:03 +0100 Subject: [PATCH 5/5] Avoid broadcast invocation in validation tests. --- test/metal.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/metal.jl b/test/metal.jl index 3ea6cffeb..ac47ab2dc 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -456,7 +456,7 @@ end buf = Base.unsafe_convert(MTL.MTLBuffer, arr) Metal.unsafe_fill!(current_device(), Metal.MtlPointer{T}(buf, 0), T(val), 4) - @test all(arr .== val) + @test all(Array(arr) .== val) end end