Prepare for CheapThreads -> Polyester, add upperboundedintegers

Author: chriselrod

src/LoopVectorization.jl

@@ -27,6 +27,8 @@ using VectorizationBase: register_size, register_count, cache_linesize, cache_si
 using IfElse: ifelse
 
 using ThreadingUtilities, CheapThreads
+const Polyester = CheapThreads
+# using ThreadingUtilities, Polyester
 using Base.Broadcast: Broadcasted, DefaultArrayStyle
 using LinearAlgebra: Adjoint, Transpose, Diagonal
 using Base.Meta: isexpr

src/codegen/lower_load.jl

@@ -30,7 +30,8 @@ function prefetchisagoodidea(ls::LoopSet, op::Operation, td::UnrollArgs)
                 end
             end
             innermostloopind == -1 && return 0
-            if prod(s -> Float64(looplength(ls, s)), @view(indices[1:innermostloopind-1])) ≥ 120.0 && length(getloop(ls, innermostloopsym)) ≥ 120
+          # if prod(s -> Float64(looplength(ls, s)), @view(indices[1:innermostloopind-1])) ≥ 120.0 &&
+            if length(getloop(ls, innermostloopsym)) ≥ 120
                 if getoffsets(op)[innermostloopind] < 120
                     for opp ∈ operations(ls)
                         if iscompute(opp) && (innermostloopsym ∈ loopdependencies(opp)) &&

src/codegen/lower_threads.jl

@@ -378,12 +378,12 @@ function thread_one_loops_expr(
     $loopstart
     var"##do#thread##" = var"#nrequest#" ≠ 0x00000000
     if var"##do#thread##"
-      var"#threads#", var"#torelease#" = CheapThreads.request_threads(Threads.threadid()%UInt32, var"#nrequest#")
+      var"#threads#", var"#torelease#" = Polyester.request_threads(Threads.threadid()%UInt32, var"#nrequest#")
       var"#thread#factor#0#" = var"#nthreads#"
       $iterdef
       var"#thread#launch#count#" = 0x00000000
       var"#thread#id#" = 0x00000000
-      var"#thread#mask#" = CheapThreads.mask(var"#threads#")
+      var"#thread#mask#" = Polyester.mask(var"#threads#")
       var"#threads#remain#" = true
       while var"#threads#remain#"
         VectorizationBase.assume(var"#thread#mask#" ≠ zero(var"#thread#mask#"))
@@ -409,14 +409,14 @@ function thread_one_loops_expr(
         var"#threads#remain#" = var"#thread#launch#count#" ≠ var"#nrequest#"
       end
     else# eliminate undef var errors that the compiler should be able to figure out are unreachable, but doesn't
-      var"#torelease#" = zero(CheapThreads.worker_type())
-      var"#threads#" = CheapThreads.UnsignedIteratorEarlyStop(var"#torelease#", 0x00000000)
+      var"#torelease#" = zero(Polyester.worker_type())
+      var"#threads#" = Polyester.UnsignedIteratorEarlyStop(var"#torelease#", 0x00000000)
     end
     var"#avx#call#args#" = $avxcall_args
     $_avx_call_
     var"##do#thread##" || $retexpr
     var"#thread#id#" = 0x00000000
-    var"#thread#mask#" = CheapThreads.mask(var"#threads#")
+    var"#thread#mask#" = Polyester.mask(var"#threads#")
     var"#threads#remain#" = true
     while var"#threads#remain#"
       VectorizationBase.assume(var"#thread#mask#" ≠ zero(var"#thread#mask#"))
@@ -428,7 +428,7 @@ function thread_one_loops_expr(
       $update_return_values
       var"#threads#remain#" = var"#thread#mask#" ≠ 0x00000000
     end
-    CheapThreads.free_threads!(var"#torelease#")
+    Polyester.free_threads!(var"#torelease#")
     $retexpr
   end
   Expr(:block, ls.preamble, q)
@@ -559,15 +559,15 @@ function thread_two_loops_expr(
     var"#loop#1#start#init#" = var"#iter#start#0#"
     var"##do#thread##" = var"#nrequest#" ≠ 0x00000000
     if var"##do#thread##"
-      var"#threads#", var"#torelease#" = CheapThreads.request_threads(Threads.threadid(), var"#nrequest#")
+      var"#threads#", var"#torelease#" = Polyester.request_threads(Threads.threadid(), var"#nrequest#")
       $iterdef1
       $iterdef2
       # @show var"#base#block#size#thread#0#", var"#block#rem#step#0#" var"#base#block#size#thread#1#", var"#block#rem#step#1#"
       var"#thread#launch#count#" = 0x00000000
       var"#thread#launch#count#0#" = 0x00000000
       var"#thread#launch#count#1#" = 0x00000000
       var"#thread#id#" = 0x00000000
-      var"#thread#mask#" = CheapThreads.mask(var"#threads#")
+      var"#thread#mask#" = Polyester.mask(var"#threads#")
       var"#threads#remain#" = true
       while var"#threads#remain#"
         VectorizationBase.assume(var"#thread#mask#" ≠ zero(var"#thread#mask#"))
@@ -604,16 +604,16 @@ function thread_two_loops_expr(
         var"#threads#remain#" = var"#thread#launch#count#" ≠ var"#nrequest#"
       end
     else# eliminate undef var errors that the compiler should be able to figure out are unreachable, but doesn't
-      var"#torelease#" = zero(CheapThreads.worker_type())
-      var"#threads#" = CheapThreads.UnsignedIteratorEarlyStop(var"#torelease#", 0x00000000)
+      var"#torelease#" = zero(Polyester.worker_type())
+      var"#threads#" = Polyester.UnsignedIteratorEarlyStop(var"#torelease#", 0x00000000)
     end
     # @show $lastboundexpr
     var"#avx#call#args#" = $avxcall_args
     $_avx_call_
     var"##do#thread##" || $retexpr
     # @show $retv
     var"#thread#id#" = 0x00000000
-    var"#thread#mask#" = CheapThreads.mask(var"#threads#")
+    var"#thread#mask#" = Polyester.mask(var"#threads#")
     var"#threads#remain#" = true
     while var"#threads#remain#"
       VectorizationBase.assume(var"#thread#mask#" ≠ zero(var"#thread#mask#"))
@@ -625,7 +625,7 @@ function thread_two_loops_expr(
       $update_return_values
       var"#threads#remain#" = var"#thread#mask#" ≠ 0x00000000
     end
-    CheapThreads.free_threads!(var"#torelease#")
+    Polyester.free_threads!(var"#torelease#")
     $retexpr
   end
   # @show

src/codegen/lowering.jl

@@ -218,10 +218,17 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
     nisunrolled = isunrolled1(us, n)
     nisvectorized = isvectorized(us, n)
     W = nisvectorized ? ls.vector_width : 1
-    loopisstatic = isstaticloop(loop) & (!iszero(W))
     UFW = UF * W
     looplength = length(loop)
-    if loopisstatic & (UFW > looplength)
+  if W ≠ 0 & isknown(first(loop)) & isknown(step(loop))
+    loopisstatic = isknown(last(loop))
+    # something other than the default hint currently means an UpperBoundedInteger was passed as an argument
+    loopisbounded = (looplength < UFW) & (loopisstatic | (gethint(last(loop)) ≠ 1024))
+  else
+    loopisstatic = false
+    loopisbounded = false
+  end
+    if loopisstatic & loopisbounded
         UFWnew = cld(looplength, cld(looplength, UFW))
         UF = cld(UFWnew, W)
         UFW = UF * W
@@ -235,9 +242,11 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
     remfirst = loopisstatic & (!nisvectorized) & (UFt > 0) & !(unsigned(Ureduct) < unsigned(UF))
     tc = terminatecondition(ls, us, n, inclmask, remfirst ? 1 : UF)
     # usorig = ls.unrollspecification
-    # tc = (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !loopisstatic && !inclmask && !ls.loadelimination ? expect(tc) : tc
-    body = lower_block(ls, us, n, inclmask, UF)
-    if loopisstatic
+  # tc = (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !loopisstatic && !inclmask && !ls.loadelimination ? expect(tc) : tc
+  # Don't need to create the body if loop is dynamic and bounded
+  dynamicbounded = ((!loopisstatic) & loopisbounded)
+  body = dynamicbounded ? tc : lower_block(ls, us, n, inclmask, UF)
+  if loopisstatic
         iters = length(loop) ÷ UFW
         if (iters ≤ 1) || (iters*UF ≤ 16 && allinteriorunrolled(ls, us, n))# Let's set a limit on total unrolling
             q = Expr(:block)
@@ -254,16 +263,18 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
         remblock = init_remblock(loop, ls.lssm, n)#loopsym)
         # unroll_cleanup = Ureduct > 0 || (nisunrolled ? (u₂ > 1) : (u₁ > 1))
         # remblock = unroll_cleanup ? init_remblock(loop, ls.lssm, n)#loopsym) : Expr(:block)
-        q = if unsigned(Ureduct) < unsigned(UF)
-            # push!(body.args, Expr(:(||), tc, Expr(:break)))
-            # Expr(:while, true, body)
-            termcond = gensym(:maybeterm)
-            push!(body.args, Expr(:(=), termcond, tc))
-            Expr(:block, Expr(:(=), termcond, true), Expr(:while, termcond, body))
-        else
-            Expr(:while, tc, body)
-        end
-    end
+      q = if loopisbounded
+        Expr(:block)
+      elseif unsigned(Ureduct) < unsigned(UF)
+        # push!(body.args, Expr(:(||), tc, Expr(:break)))
+        # Expr(:while, true, body)
+        termcond = gensym(:maybeterm)
+        push!(body.args, Expr(:(=), termcond, tc))
+        Expr(:block, Expr(:(=), termcond, true), Expr(:while, termcond, body))
+      else
+        Expr(:while, tc, body)
+      end
+  end
     q = if unsigned(Ureduct) < unsigned(UF) # unsigned(-1) == typemax(UInt); 
         add_cleanup = true
         if isone(Ureduct)
@@ -290,7 +301,11 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
             UF_cleanup = UF - Ureduct
             blockhead = :if
         end
-        _q = Expr(:block, add_upper_outer_reductions(ls, q, Ureduct, UF, loop, nisvectorized))
+        _q = if dynamicbounded
+            initialize_outer_reductions!(q, ls, Ureduct); q
+        else
+            Expr(:block, add_upper_outer_reductions(ls, q, Ureduct, UF, loop, nisvectorized))
+        end
         if add_cleanup
             cleanup_expr = Expr(blockhead)
             blockhead === :block || push!(cleanup_expr.args, terminatecondition(ls, us, n, inclmask, UF_cleanup))

src/condense_loopset.jl

@@ -226,6 +226,7 @@ end
 
 @inline zerorangestart(r::Base.OneTo) = CloseOpen(maybestaticlast(r))
 @inline zerorangestart(r::CloseOpen) = CloseOpen(length(r))
+@inline zerorangestart(r::CloseOpen{Zero}) = r
 @inline zerorangestart(r::AbstractUnitRange) = Zero():One():(maybestaticlast(r)-maybestaticfirst(r))
 @inline zerorangestart(r::AbstractRange) = Zero():static_step(r):(maybestaticlast(r)-maybestaticfirst(r))
 @inline zerorangestart(r::CartesianIndices) = CartesianIndices(map(zerorangestart, r.indices))

src/reconstruct_loopset.jl

@@ -1,6 +1,25 @@
 const NOpsType = Int#Union{Int,Vector{Int}}
 
-function Loop(ls::LoopSet, ex::Expr, sym::Symbol, f, s, l)
+struct UpperBoundedInteger{N, T<: Integer} <: Integer
+  i::T
+end
+@inline UpperBoundedInteger{N}(i::T) where {N,T<:Integer} = UpperBoundedInteger{N,T}(i)
+@inline UpperBoundedInteger(i::T, ::StaticInt{N}) where {N,T<:Integer} = UpperBoundedInteger{N,T}(i)
+@inline Base.:(%)(a::UpperBoundedInteger, ::Type{T}) where {T<:Integer} = a.i % T
+Base.promote_rule(::Type{T}, ::Type{UpperBoundedInteger{N,S}}) where {N,T,S} = promote_rule(T,S)
+Base.promote_rule(::Type{UpperBoundedInteger{N,S}}, ::Type{T}) where {N,T,S} = promote_rule(S,T)
+Base.convert(::Type{T}, i::UpperBoundedInteger) where {T<:Number} = convert(T, i.i)
+Base.convert(::Type{UpperBoundedInteger{N,T}}, i::UpperBoundedInteger{N,T}) where {N,T<:Integer} = i
+Base.convert(::Type{Any}, i::UpperBoundedInteger) = i
+upper_bound(_) = typemax(Int)
+upper_bound(::Type{CloseOpen{T,UpperBoundedInteger{N,S}}}) where {T,N,S} = N - 1
+
+@inline Base.last(r::CloseOpen{<:Integer,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i) - One()
+@inline ArrayInterface.static_last(r::CloseOpen{<:Integer,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i) - One()
+@inline Base.length(r::CloseOpen{<:Integer,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i) - getfield(r,:start)
+@inline Base.length(r::CloseOpen{Zero,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i)
+
+function Loop(ls::LoopSet, ex::Expr, sym::Symbol, f, s, l, ub::Int)
     if (f !== nothing) && (s !== nothing) && (l !== nothing)
         return static_loop(sym, f, s, l)
     end
@@ -11,32 +30,33 @@ function Loop(ls::LoopSet, ex::Expr, sym::Symbol, f, s, l)
     pushpreamble!(ls, Expr(:(=), lensym, Expr(:call, lv(:maybestaticlength), rangesym)))
     F = if f === nothing
         start = gensym(ssym*"_loopstart")
-        pushpreamble!(ls, Expr(:(=), start, Expr(:call, lv(:first), rangesym)))
+        pushpreamble!(ls, Expr(:(=), start, Expr(:call, %, Expr(:call, lv(:first), rangesym), Int)))
         MaybeKnown(start, 1)
     else
         MaybeKnown(f)
     end
     S = if s === nothing
         step = gensym(ssym*"_loopstep")
-        pushpreamble!(ls, Expr(:(=), step, Expr(:call, lv(:step), rangesym)))
+        pushpreamble!(ls, Expr(:(=), step, Expr(:call, %, Expr(:call, lv(:step), rangesym), Int)))
         MaybeKnown(step, 1)
     else
         MaybeKnown(s)
     end
     L = if l === nothing
         stop = gensym(ssym*"_loopstop")
-        pushpreamble!(ls, Expr(:(=), stop, Expr(:call, lv(:last), rangesym)))
-        MaybeKnown(stop, 1024)
+        pushpreamble!(ls, Expr(:(=), stop, Expr(:call, %, Expr(:call, lv(:last), rangesym), Int)))
+        MaybeKnown(stop, min(ub, 1024))
     else
         MaybeKnown(l)
     end
     loopiteratesatleastonce!(ls, Loop(sym, F, L, S, rangesym, lensym))
 end
 function Loop(ls::LoopSet, ex::Expr, sym::Symbol, ::Type{R}) where {R<:AbstractRange}
-    f = ArrayInterface.known_first(R)
-    s = ArrayInterface.known_step(R)
-    l = ArrayInterface.known_last(R)
-    Loop(ls, ex, sym, f, s, l)
+  f = ArrayInterface.known_first(R)
+  s = ArrayInterface.known_step(R)
+  l = ArrayInterface.known_last(R)
+  ub = upper_bound(R)
+  Loop(ls, ex, sym, f, s, l, ub)
 end
 
 function static_loop(sym::Symbol, L::Int, S::Int, U::Int)
@@ -683,7 +703,7 @@ Execute an `@avx` block. The block's code is represented via the arguments:
 @aggressive_constprop @generated function _avx_!(
     ::Val{var"#UNROLL#"}, ::Val{var"#OPS#"}, ::Val{var"#ARF#"}, ::Val{var"#AM#"}, ::Val{var"#LPSYM#"}, ::Val{Tuple{var"#LB#",var"#V#"}}, var"#flattened#var#arguments#"::Vararg{Any,var"#num#vargs#"}
 ) where {var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#", var"#V#", var"#num#vargs#"}
-  # 1 + 1 # Irrelevant line you can comment out/in to force recompilation...
+  1 + 1 # Irrelevant line you can comment out/in to force recompilation...
   ls = _avx_loopset(var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#".parameters, var"#V#".parameters, var"#UNROLL#")
   pushfirst!(ls.preamble.args, :(var"#lv#tuple#args#" = reassemble_tuple(Tuple{var"#LB#",var"#V#"}, var"#flattened#var#arguments#")))
   # return @show avx_body(ls, var"#UNROLL#")

test/runtests.jl

@@ -73,6 +73,8 @@ const START_TIME = time()
 
   @time include("outer_reductions.jl")
 
+  @time include("upperboundedintegers.jl")
+  
   if VERSION ≥ v"1.6"
     @time include("quantum.jl")
   end

test/upperboundedintegers.jl

@@ -0,0 +1,33 @@
+using OffsetArrays, Test
+@testset "UpperboundedIntegers" begin
+  function ubsum(x)
+    @assert firstindex(x) == 0
+    r = LoopVectorization.CloseOpen(LoopVectorization.UpperBoundedInteger(length(x), StaticInt(15)))
+    s = zero(eltype(x))
+    @avx for i ∈ r
+      s += x[i]
+    end
+    s
+  end
+  function ubdouble!(y,x)
+    @assert firstindex(x) == 0
+    r = LoopVectorization.CloseOpen(LoopVectorization.UpperBoundedInteger(length(x), StaticInt(15)))
+    @avx for i ∈ r
+      y[i] = 2*x[i]
+    end
+    y
+  end
+  for l ∈ 1:15
+    x = OffsetVector(rand(l), -1)
+    @test ubsum(x) ≈ sum(x)
+    @test ubdouble!(similar(x), x) == x .* 2
+  end
+  for l ∈ 70:75
+    x = OffsetVector(rand(l), -1)
+    # @test ubsum(x) ≈ sum(@view(x[begin:14]))
+    @test ubsum(x) ≈ ubsum(x)
+    xs = similar(x)
+    @test @view(ubdouble!(xs, x)[begin:14]) == @view(x[begin:14]) .* 2;
+    @test xs[end] ≠ 2x[end]
+  end
+end