Optimize indices through @generated reduce.

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.14.10"
+version = "2.14.11"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/ranges.jl

@@ -443,21 +443,159 @@ specified, then the indices for visiting each index of `x` are returned.
 end
 @inline indices(x::AbstractUnitRange{<:Integer}) = Base.Slice(OptionallyStaticUnitRange(x))
 
+"""
+  reduce_tup(f::F, inds::Tuple{Vararg{Any,N}}) where {F,N}
+
+An optimized `reduce` for tuples. `Base.reduce`'s `afoldl` will often not inline.
+Additionally, `reduce_tup` attempts to order the reduction in an optimal manner.
+
+```julia
+julia> using StaticArrays, ArrayInterface, BenchmarkTools
+
+julia> rsum(v::SVector) = ArrayInterface.reduce_tup(+, v.data)
+rsum (generic function with 2 methods)
+
+julia> for n ∈ 2:16
+           @show n
+           v = @SVector rand(n)
+           s1 = @btime  sum(\$(Ref(v))[])
+           s2 = @btime rsum(\$(Ref(v))[])
+       end
+n = 2
+  0.863 ns (0 allocations: 0 bytes)
+  0.864 ns (0 allocations: 0 bytes)
+n = 3
+  0.863 ns (0 allocations: 0 bytes)
+  0.863 ns (0 allocations: 0 bytes)
+n = 4
+  0.864 ns (0 allocations: 0 bytes)
+  1.075 ns (0 allocations: 0 bytes)
+n = 5
+  0.864 ns (0 allocations: 0 bytes)
+  1.076 ns (0 allocations: 0 bytes)
+n = 6
+  0.865 ns (0 allocations: 0 bytes)
+  1.077 ns (0 allocations: 0 bytes)
+n = 7
+  1.075 ns (0 allocations: 0 bytes)
+  0.866 ns (0 allocations: 0 bytes)
+n = 8
+  0.974 ns (0 allocations: 0 bytes)
+  1.076 ns (0 allocations: 0 bytes)
+n = 9
+  1.081 ns (0 allocations: 0 bytes)
+  1.077 ns (0 allocations: 0 bytes)
+n = 10
+  1.203 ns (0 allocations: 0 bytes)
+  1.077 ns (0 allocations: 0 bytes)
+n = 11
+  1.355 ns (0 allocations: 0 bytes)
+  1.292 ns (0 allocations: 0 bytes)
+n = 12
+  1.539 ns (0 allocations: 0 bytes)
+  1.079 ns (0 allocations: 0 bytes)
+n = 13
+  1.704 ns (0 allocations: 0 bytes)
+  1.290 ns (0 allocations: 0 bytes)
+n = 14
+  1.916 ns (0 allocations: 0 bytes)
+  1.185 ns (0 allocations: 0 bytes)
+n = 15
+  2.072 ns (0 allocations: 0 bytes)
+  1.292 ns (0 allocations: 0 bytes)
+n = 16
+  2.273 ns (0 allocations: 0 bytes)
+  1.076 ns (0 allocations: 0 bytes)
+```
+
+More importantly, `reduce_tup(_pick_range, inds)` often performs better than `reduce(_pick_range, inds)`.
+```julia
+julia> using ArrayInterface, BenchmarkTools
+
+julia> inds = (Base.OneTo(100), 1:100, 1:ArrayInterface.StaticInt(100))
+(Base.OneTo(100), 1:100, 1:Static(100))
+
+julia> @btime reduce(ArrayInterface._pick_range, \$(Ref(inds))[])
+  6.000 ns (0 allocations: 0 bytes)
+Base.Slice(Static(1):Static(100))
+
+julia> @btime ArrayInterface.reduce_tup(ArrayInterface._pick_range, \$(Ref(inds))[])
+  2.578 ns (0 allocations: 0 bytes)
+Base.Slice(Static(1):Static(100))
+
+julia> inds = (Base.OneTo(100), 1:100, 1:UInt(100))
+(Base.OneTo(100), 1:100, 0x0000000000000001:0x0000000000000064)
+
+julia> @btime reduce(ArrayInterface._pick_range, \$(Ref(inds))[])
+  6.191 ns (0 allocations: 0 bytes)
+Base.Slice(Static(1):100)
+
+julia> @btime ArrayInterface.reduce_tup(ArrayInterface._pick_range, \$(Ref(inds))[])
+  2.591 ns (0 allocations: 0 bytes)
+Base.Slice(Static(1):100)
+
+julia> inds = (Base.OneTo(100), 1:100, 1:UInt(100), Int32(1):Int32(100))
+(Base.OneTo(100), 1:100, 0x0000000000000001:0x0000000000000064, 1:100)
+
+julia> @btime reduce(ArrayInterface._pick_range, $(Ref(inds))[])
+  9.268 ns (0 allocations: 0 bytes)
+Base.Slice(Static(1):100)
+
+julia> @btime ArrayInterface.reduce_tup(ArrayInterface._pick_range, $(Ref(inds))[])
+  2.570 ns (0 allocations: 0 bytes)
+Base.Slice(Static(1):100)
+```
+"""
+@generated function reduce_tup(f::F, inds::Tuple{Vararg{Any,N}}) where {F,N}
+    q = Expr(:block, Expr(:meta, :inline))
+    if N == 1
+        push!(q.args, :(inds[1]))
+        return q
+    end
+    splits = 0
+    _N = N
+    while _N > 1
+        _Nhalf = _N >> 1
+        for n ∈ 1:_Nhalf
+            assign = Symbol(:r_,n,:_,splits)
+            call = if splits == 0
+                Expr(:call, :f, Expr(:ref, :inds, n), Expr(:ref, :inds, n + _Nhalf))
+            else
+                Expr(:call, :f, Symbol(:r_,n,:_,splits-1), Symbol(:r_,n + _Nhalf,:_,splits-1))
+            end
+            push!(q.args, Expr(:(=), assign, call))
+        end
+        for (i,n) ∈ enumerate((_Nhalf<<1)+1:_N)
+            assign = Symbol(:r_,i,:_,splits)
+            call = if _N == N
+                Expr(:call, :f, assign, Expr(:ref, :inds, n))
+            else
+                Expr(:call, :f, assign, Symbol(:r_, n, :_, splits-1))
+            end
+            push!(q.args, Expr(:(=), assign, call))
+        end
+        splits += 1
+        _N = _Nhalf
+    end
+    push!(q.args, Symbol(:r_,1,:_,splits - 1))
+    q
+end
+
 function indices(x::Tuple)
     inds = map(eachindex, x)
-    return reduce(_pick_range, inds)
+    return reduce_tup(_pick_range, inds)
 end
 
 @inline indices(x, d) = indices(axes(x, d))
 
 @inline function indices(x::Tuple{Vararg{Any,N}}, dim) where {N}
     inds = map(x_i -> indices(x_i, dim), x)
-    return reduce(_pick_range, inds)
+    return reduce_tup(_pick_range, inds)
 end
 
 @inline function indices(x::Tuple{Vararg{Any,N}}, dim::Tuple{Vararg{Any,N}}) where {N}
     inds = map(indices, x, dim)
-    return reduce(_pick_range, inds)
+    return reduce_tup(_pick_range, inds)
 end
 
 @inline function _pick_range(x, y)