Add CloseOpen to make working with zero-indexed arrays easier and allow avoiding some +1/-1s when using them (note it isn't used internally at the moment), refactor code related to interval definitions.

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.9.13"
+version = "0.9.14"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -14,13 +14,13 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
-ArrayInterface = "2.14.9"
+ArrayInterface = "2.14.10"
 DocStringExtensions = "0.8"
 IfElse = "0.1"
 OffsetArrays = "1.4.1"
 SLEEFPirates = "0.6"
 UnPack = "1"
-VectorizationBase = "0.14"
+VectorizationBase = "0.14.10"
 julia = "1.5"
 
 [extras]

src/LoopVectorization.jl

@@ -33,7 +33,7 @@ import LinearAlgebra # for check_args
 using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast
 
 using ArrayInterface
-using ArrayInterface: OptionallyStaticUnitRange, Zero, One
+using ArrayInterface: OptionallyStaticUnitRange, Zero, One, static_length
 const Static = ArrayInterface.StaticInt
 
 
@@ -52,6 +52,7 @@ const VECTORWIDTHSYMBOL, ELTYPESYMBOL = Symbol("##Wvecwidth##"), Symbol("##Tloop
 
 include("vectorizationbase_compat/contract_pass.jl")
 include("vectorizationbase_compat/subsetview.jl")
+include("closeopen.jl")
 include("getconstindexes.jl")
 include("predicates.jl")
 include("map.jl")

src/broadcast.jl

@@ -102,14 +102,16 @@ function add_broadcast!(
     @nospecialize(prod::Type{<:Product}), elementbytes::Int
 )
     A, B = prod.parameters
-    K = gensym!(ls, "K")
+    Krange = gensym!(ls, "K")
+    Klen = gensym!(ls, "K")
     mA = gensym!(ls, "Aₘₖ")
     mB = gensym!(ls, "Bₖₙ")
     pushprepreamble!(ls, Expr(:(=), mA, Expr(:(.), bcname, QuoteNode(:a))))
     pushprepreamble!(ls, Expr(:(=), mB, Expr(:(.), bcname, QuoteNode(:b))))
-    pushprepreamble!(ls, Expr(:(=), K, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__,Symbol(@__FILE__)), Expr(:ref, Expr(:call, :size, mB), 1))))
+    pushprepreamble!(ls, Expr(:(=), Klen, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__,Symbol(@__FILE__)), Expr(:ref, Expr(:call, :size, mB), 1))))
+    pushpreamble!(ls, Expr(:(=), Krange, Expr(:call, :(:), staticexpr(1), Klen)))
     k = gensym!(ls, "k")
-    add_loop!(ls, Loop(k, 1, K), k)
+    add_loop!(ls, Loop(k, 1, Klen, Krange, Klen), k)
     m = loopsyms[1];
     if numdims(B) == 1
         bloopsyms = Symbol[k]
@@ -335,24 +337,34 @@ function add_broadcast!(
     pushop!(ls, op, destname)
 end
 
+function add_broadcast_loops!(ls::LoopSet, loopsyms::Vector{Symbol}, destsym::Symbol)
+    axes_tuple = Expr(:tuple)
+    pushpreamble!(ls, Expr(:(=), axes_tuple, Expr(:call, :axes, destsym)))
+    for (n,itersym) ∈ enumerate(loopsyms)
+        Nrange = gensym!(ls, "N")
+        Nlower = gensym!(ls, "N")
+        Nupper = gensym!(ls, "N")
+        Nlen = gensym!(ls, "N")
+        add_loop!(ls, Loop(itersym, Nlower, Nupper, Nrange, Nlen), itersym)
+        push!(axes_tuple.args, Nrange)
+        pushpreamble!(ls, Expr(:(=), Nlower, Expr(:call, lv(:maybestaticfirst), Nrange)))
+        pushpreamble!(ls, Expr(:(=), Nupper, Expr(:call, lv(:maybestaticlast), Nrange)))
+        pushpreamble!(ls, Expr(:(=), Nlen, Expr(:call, lv(:static_length), Nrange)))
+    end
+end
 # size of dest determines loops
 # function vmaterialize!(
 @generated function vmaterialize!(
     dest::AbstractArray{T,N}, bc::BC, ::Val{Mod}
 ) where {T <: NativeTypes, N, BC <: Union{Broadcasted,Product}, Mod}
+    # 2+1
     # we have an N dimensional loop.
     # need to construct the LoopSet
     # @show typeof(dest)
     ls = LoopSet(Mod)
     loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
     ls.isbroadcast[] = true
-    sizes = Expr(:tuple)
-    for (n,itersym) ∈ enumerate(loopsyms)
-        Nsym = gensym!(ls, "N")
-        add_loop!(ls, Loop(itersym, 1, Nsym), itersym)
-        push!(sizes.args, Nsym)
-    end
-    pushpreamble!(ls, Expr(:(=), sizes, Expr(:call, :size, :dest)))
+    add_broadcast_loops!(ls, loopsyms, :dest)
     elementbytes = sizeof(T)
     add_broadcast!(ls, :dest, :bc, loopsyms, BC, elementbytes)
     add_simple_store!(ls, :dest, ArrayReference(:dest, loopsyms), elementbytes)
@@ -381,13 +393,7 @@ end
     loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
     ls.isbroadcast[] = true
     pushprepreamble!(ls, Expr(:(=), :dest, Expr(:call, :parent, :dest′)))
-    sizes = Expr(:tuple)
-    for (n,itersym) ∈ enumerate(loopsyms)
-        Nsym = gensym!(ls, "N")
-        add_loop!(ls, Loop(itersym, 1, Nsym), itersym)
-        push!(sizes.args, Nsym)
-    end
-    pushpreamble!(ls, Expr(:(=), sizes, Expr(:call, :size, :dest′)))
+    add_broadcast_loops!(ls, loopsyms, :dest′)
     elementbytes = sizeof(T)
     add_broadcast!(ls, :dest, :bc, loopsyms, BC, elementbytes)
     add_simple_store!(ls, :dest, ArrayReference(:dest, reverse(loopsyms)), elementbytes)

src/closeopen.jl

@@ -0,0 +1,31 @@
+
+struct CloseOpen{L <: Union{Int,StaticInt}, U <: Union{Int,StaticInt}} <: AbstractUnitRange{Int}
+    start::L
+    upper::U
+    @inline CloseOpen(s::StaticInt{L}, u::StaticInt{U}) where {L,U} = new{StaticInt{L},StaticInt{U}}(s, u)
+    @inline CloseOpen(s::Integer, u::StaticInt{U}) where {U} = new{Int,StaticInt{U}}(s % Int, u)
+    @inline CloseOpen(s::StaticInt{L}, u::Integer) where {L} = new{StaticInt{L},Int}(s, u % Int)
+    @inline CloseOpen(s::Integer, u::Integer) = new{Int,Int}(s % Int, u % Int)
+end
+@inline CloseOpen(len::Integer) = CloseOpen(Zero(), len)
+
+@inline Base.first(r::CloseOpen) = r.start
+@inline Base.step(::CloseOpen) = One()
+@inline Base.last(r::CloseOpen) = r.upper - One()
+@inline Base.length(r::CloseOpen) = r.upper - r.start
+@inline Base.length(r::CloseOpen{Zero}) = r.upper
+
+@inline Base.iterate(r::CloseOpen) = (i = Int(first(r)); (i, i))
+@inline Base.iterate(r::CloseOpen, i::Int) = (i += 1) == r.upper ? nothing : (i, i)
+
+ArrayInterface.known_first(::Type{<:CloseOpen{StaticInt{F}}}) where {F} = F
+ArrayInterface.known_step(::Type{<:CloseOpen}) = 1
+ArrayInterface.known_last(::Type{<:CloseOpen{<:Any,StaticInt{L}}}) where {L} = L - 1
+ArrayInterface.known_length(::Type{CloseOpen{StaticInt{F},StaticInt{L}}}) where {F,L} = L - F
+
+@inline canonicalize_range(r::OptionallyStaticUnitRange) = r
+@inline canonicalize_range(r::CloseOpen) = r
+@inline canonicalize_range(r::AbstractUnitRange) = maybestaticfirst(r):maybestaticlast(r)
+@inline canonicalize_range(r::CartesianIndices) = CartesianIndices(map(canonicalize_range, r.indices))
+
+

src/condense_loopset.jl

@@ -122,23 +122,26 @@ function OperationStruct!(varnames::Vector{Symbol}, ids::Vector{Int}, ls::LoopSe
 end
 ## turn a LoopSet into a type object which can be used to reconstruct the LoopSet.
 
-function loop_boundary(loop::Loop)
-    startexact = loop.startexact
-    stopexact = loop.stopexact
-    if startexact & stopexact
-        Expr(:call, lv(:OptionallyStaticUnitRange), staticexpr(loop.starthint), staticexpr(loop.stophint))
-    elseif startexact
-        Expr(:call, lv(:OptionallyStaticUnitRange), staticexpr(loop.starthint), loop.stopsym)
-    elseif stopexact
-        Expr(:call, lv(:OptionallyStaticUnitRange), loop.startsym, staticexpr(loop.stophint))
+function loop_boundary!(q::Expr, loop::Loop)
+    if loop.startexact & loop.stopexact
+        push!(q.args, Expr(:call, lv(:OptionallyStaticUnitRange), staticexpr(loop.starthint), staticexpr(loop.stophint)))
+    elseif loop.rangesym === Symbol("")
+        lb = if startexact
+            Expr(:call, lv(:OptionallyStaticUnitRange), staticexpr(loop.starthint), loop.stopsym)
+        elseif stopexact
+            Expr(:call, lv(:OptionallyStaticUnitRange), loop.startsym, staticexpr(loop.stophint))
+        else
+            Expr(:call, :(:), loop.startsym, loop.stopsym)
+        end
+        push!(q.args, lb)
     else
-        Expr(:call, :(:), loop.startsym, loop.stopsym)
+        push!(q.args, loop.rangesym)
     end
 end
 
 function loop_boundaries(ls::LoopSet)
     lbd = Expr(:tuple)
-    foreach(loop -> push!(lbd.args, loop_boundary(loop)), ls.loops)
+    foreach(loop -> loop_boundary!(lbd, loop), ls.loops)
     lbd
 end