More efficient offsets (#127)

* More efficient offsets

* Small doc fix to trigger CI

Author: Tokazama

Project.toml

@@ -12,7 +12,7 @@ Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 [compat]
 IfElse = "0.1"
 Requires = "0.5, 1.0"
-Static = "0.1"
+Static = "0.2"
 julia = "1.2"
 
 [extras]

src/ArrayInterface.jl

@@ -666,7 +666,7 @@ end
 defines_strides(::Type{<:BitArray}) = true
 
 """
-can_avx(f)
+    can_avx(f) -> Bool
 
 Returns `true` if the function `f` is guaranteed to be compatible with
 `LoopVectorization.@avx` for supported element and array types. While a return
@@ -1071,7 +1071,6 @@ function __init__()
     @require OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" begin
         size(A::OffsetArrays.OffsetArray) = size(parent(A))
         strides(A::OffsetArrays.OffsetArray) = strides(parent(A))
-        # offsets(A::OffsetArrays.OffsetArray) = map(+, A.offsets, offsets(parent(A)))
         function parent_type(
             ::Type{O},
         ) where {T,N,A<:AbstractArray{T,N},O<:OffsetArrays.OffsetArray{T,N,A}}
@@ -1084,8 +1083,16 @@ function __init__()
         function contiguous_batch_size(::Type{A}) where {A<:OffsetArrays.OffsetArray}
             return contiguous_batch_size(parent_type(A))
         end
-        stride_rank(::Type{A}) where {A<:OffsetArrays.OffsetArray} =
-            stride_rank(parent_type(A))
+
+        function _offset_axis_type(::Type{T}, dim::StaticInt{D}) where {T,D}
+            return OffsetArrays.IdOffsetRange{Int,ArrayInterface.axes_types(T, dim)}
+        end
+        function ArrayInterface.axes_types(::Type{T}) where {T<:OffsetArrays.OffsetArray}
+            return Static.eachop_tuple(_offset_axis_type, Static.nstatic(Val(ndims(T))), ArrayInterface.parent_type(T))
+        end
+        function stride_rank(::Type{A}) where {A<:OffsetArrays.OffsetArray}
+            return stride_rank(parent_type(A))
+        end
         ArrayInterface.axes(A::OffsetArrays.OffsetArray) = Base.axes(A)
         ArrayInterface.axes(A::OffsetArrays.OffsetArray, dim::Integer) = Base.axes(A, dim)
         function ArrayInterface.device(::Type{T}) where {T<:OffsetArrays.OffsetArray}

src/axes.jl

@@ -41,10 +41,10 @@ function axes_types(::Type{T}) where {T<:VecAdjTrans}
     return Tuple{OptionallyStaticUnitRange{One,One},axes_types(parent_type(T), One())}
 end
 function axes_types(::Type{T}) where {T<:MatAdjTrans}
-    return eachop_tuple(_get_tuple, axes_types(parent_type(T)); iterator=to_parent_dims(T))
+    return eachop_tuple(_get_tuple, to_parent_dims(T), axes_types(parent_type(T)))
 end
 function axes_types(::Type{T}) where {T<:PermutedDimsArray}
-    return eachop_tuple(_get_tuple, axes_types(parent_type(T)); iterator=to_parent_dims(T))
+    return eachop_tuple(_get_tuple, to_parent_dims(T), axes_types(parent_type(T)))
 end
 function axes_types(::Type{T}) where {T<:AbstractRange}
     if known_length(T) === nothing
@@ -61,7 +61,7 @@ _int_or_static_int(::Nothing) = Int
 _int_or_static_int(x::Int) = StaticInt{x}
 
 @inline function axes_types(::Type{T}) where {N,P,I,T<:SubArray{<:Any,N,P,I}}
-    return eachop_tuple(_sub_axis_type, T; iterator=to_parent_dims(T))
+    return eachop_tuple(_sub_axis_type, to_parent_dims(T), T)
 end
 @inline function _sub_axis_type(::Type{A}, dim::StaticInt) where {T,N,P,I,A<:SubArray{T,N,P,I}}
     return OptionallyStaticUnitRange{
@@ -75,12 +75,12 @@ function axes_types(::Type{R}) where {T,N,S,A,R<:ReinterpretArray{T,N,S,A}}
         if sizeof(S) === sizeof(T)
             return axes_types(A)
         elseif sizeof(S) > sizeof(T)
-            return eachop_tuple(_reshaped_axis_type, R; iterator=to_parent_dims(R))
+            return eachop_tuple(_reshaped_axis_type, to_parent_dims(R), R)
         else
-            return eachop_tuple(axes_types, A; iterator=to_parent_dims(R))
+            return eachop_tuple(axes_types, to_parent_dims(R), A)
         end
     else
-        return eachop_tuple(_non_reshaped_axis_type, R; iterator=to_parent_dims(R))
+        return eachop_tuple(_non_reshaped_axis_type, to_parent_dims(R), R)
     end
 end
 

src/dimensions.jl

@@ -182,12 +182,12 @@ end
         if invariant_permutation(perm, perm) isa True
             return dimnames(parent_type(T))
         else
-            return eachop(dimnames, parent_type(T); iterator=perm)
+            return eachop(dimnames, perm, parent_type(T))
         end
     end
 end
 function dimnames(::Type{T}) where {T<:SubArray}
-    return eachop(dimnames, parent_type(T); iterator=to_parent_dims(T))
+    return eachop(dimnames, to_parent_dims(T), parent_type(T))
 end
 
 _to_int(x::Integer) = Int(x)
@@ -241,7 +241,7 @@ end
     inds::Tuple
 ) where {N}
 
-    out = eachop(order_named_inds, x, nd, inds; iterator=nstatic(Val(N)))
+    out = eachop(order_named_inds, nstatic(Val(N)), x, nd, inds)
     _order_named_inds_check(out, length(nd))
     return out
 end

src/indexing.jl

@@ -31,7 +31,7 @@ argdims(::ArrayStyle, ::Type{T}) where {N,T<:AbstractArray{<:Any,N}} = N
 argdims(::ArrayStyle, ::Type{T}) where {N,T<:LogicalIndex{<:Any,<:AbstractArray{Bool,N}}} = N
 _argdims(s::ArrayStyle, ::Type{I}, i::StaticInt) where {I} = argdims(s, _get_tuple(I, i))
 function argdims(s::ArrayStyle, ::Type{T}) where {N,T<:Tuple{Vararg{Any,N}}}
-    return eachop(_argdims, s, T; iterator=nstatic(Val(N)))
+    return eachop(_argdims, nstatic(Val(N)), s, T)
 end
 
 """
@@ -183,7 +183,7 @@ can_flatten(::Type{A}, ::Type{T}) where {A,T<:CartesianIndices} = true
 can_flatten(::Type{A}, ::Type{T}) where {A,N,T<:AbstractArray{Bool,N}} = N > 1
 can_flatten(::Type{A}, ::Type{T}) where {A,N,T<:CartesianIndex{N}} = true
 function can_flatten(::Type{A}, ::Type{T}) where {A,N,T<:Tuple{Vararg{Any,N}}}
-    return any(eachop(_can_flat, A, T; iterator=nstatic(Val(N))))
+    return any(eachop(_can_flat, nstatic(Val(N)), A, T))
 end
 function _can_flat(::Type{A}, ::Type{T}, i::StaticInt) where {A,T}
     if can_flatten(A, _get_tuple(T, i)) === true

src/ranges.jl

@@ -401,7 +401,6 @@ const OptionallyStaticRange = Union{<:OptionallyStaticUnitRange,<:OptionallyStat
 Base.eachindex(r::OptionallyStaticRange) = r
 @inline Base.iterate(r::OptionallyStaticRange) = (fi = Int(first(r)); (fi, fi))
 
-
 Base.to_shape(x::OptionallyStaticRange) = length(x)
 Base.to_shape(x::Slice{T}) where {T<:OptionallyStaticRange} = length(x)
 

src/size.jl

@@ -22,7 +22,7 @@ function size(a::A) where {A}
 end
 #size(a::AbstractVector) = (size(a, One()),)
 
-size(x::SubArray) = eachop(_sub_size, x.indices; iterator=to_parent_dims(x))
+size(x::SubArray) = eachop(_sub_size, to_parent_dims(x), x.indices)
 _sub_size(x::Tuple, ::StaticInt{dim}) where {dim} = static_length(getfield(x, dim))
 
 @inline size(B::VecAdjTrans) = (One(), length(parent(B)))
@@ -81,7 +81,7 @@ Returns the size of each dimension for `T` known at compile time. If a dimension
 have a known size along a dimension then `nothing` is returned in its position.
 """
 known_size(x) = known_size(typeof(x))
-known_size(::Type{T}) where {T} = eachop(known_size, T; iterator=nstatic(Val(ndims(T))))
+known_size(::Type{T}) where {T} = eachop(known_size, nstatic(Val(ndims(T))), T)
 
 """
     known_size(::Type{T}, dim)

src/stridelayout.jl

@@ -9,7 +9,7 @@ stride_preserving_index(::Type{T}) where {T<:AbstractRange} = True()
 stride_preserving_index(::Type{T}) where {T<:Int} = True()
 stride_preserving_index(::Type{T}) where {T} = False()
 function stride_preserving_index(::Type{T}) where {N,T<:Tuple{Vararg{Any,N}}}
-    if all(eachop(_stride_preserving_index, T; iterator=nstatic(Val(N))))
+    if all(eachop(_stride_preserving_index, nstatic(Val(N)), T))
         return True()
     else
         return False()
@@ -27,9 +27,17 @@ it should return them as `Static` numbers.
 For example, if `A isa Base.Matrix`, `offsets(A) === (StaticInt(1), StaticInt(1))`.
 """
 @inline offsets(x, i) = static_first(indices(x, i))
-# Explicit tuple needed for inference.
-offsets(x) = eachop(offsets, x; iterator=nstatic(Val(ndims(x))))
 offsets(::Tuple) = (One(),)
+offsets(x) = eachop(_offsets, nstatic(Val(ndims(x))), x)
+function _offsets(x::X, dim::StaticInt{D}) where {X,D}
+    start = known_first(axes_types(X, dim))
+    if start === nothing
+        return first(axes(x, dim))
+    else
+        return static(start)
+    end
+end
+
 
 """
     contiguous_axis(::Type{T}) -> StaticInt{N}
@@ -346,7 +354,7 @@ end
 
 known_offsets(x) = known_offsets(typeof(x))
 function known_offsets(::Type{T}) where {T}
-    return eachop(_known_offsets, axes_types(T); iterator=nstatic(Val(ndims(T))))
+    return eachop(_known_offsets, nstatic(Val(ndims(T))), axes_types(T))
 end
 _known_offsets(::Type{T}, dim::StaticInt) where {T} = known_first(_get_tuple(T, dim))
 
@@ -433,7 +441,7 @@ end
 
 getmul(x::Tuple, y::Tuple, ::StaticInt{i}) where {i} = getfield(x, i) * getfield(y, i)
 function strides(A::SubArray)
-    return eachop(getmul, map(maybe_static_step, A.indices), strides(parent(A)); iterator=to_parent_dims(A))
+    return eachop(getmul, to_parent_dims(A), map(maybe_static_step, A.indices), strides(parent(A)))
 end
 
 maybe_static_step(x::AbstractRange) = static_step(x)