Fix for loops without arrays, and improve performance of vmap_multithreaded!. It's a bit of a silly/experimental definition, but...

Author: chriselrod

src/lowering.jl

@@ -699,10 +699,52 @@ function gc_preserve(ls::LoopSet, q::Expr)
     # Expr(:block, gcp)
 end
 
+function typeof_outer_reduction_init(ls::LoopSet, op::Operation)
+    opid = identifier(op)
+    for (id, sym) ∈ ls.preamble_symsym
+        opid == id && return Expr(:call, :typeof, sym)
+    end
+    for (id,intval) ∈ ls.preamble_symint
+        opid == id && return :Int
+    end
+    for (id,floatval) ∈ ls.preamble_symfloat
+        opid == id && return :Float64
+    end
+    for (id,typ) ∈ ls.preamble_zeros
+        instruction(ops[id]) === LOOPCONSTANT || continue
+        opid == id || continue
+        if typ == IntOrFloat
+            return :Float64
+        elseif typ == HardInt
+            return :Int
+        else#if typ == HardFloat
+            return :Float64
+        end
+    end
+    throw("Could not find initializing constant.")
+end
+function typeof_outer_reduction(ls::LoopSet, op::Operation)
+    for opp ∈ operations(ls)
+        opp === op && continue
+        name(op) === name(opp) && return typeof_outer_reduction_init(ls, opp)
+    end
+    throw("Could not find initialization op.")
+end
 
-function determine_eltype(ls::LoopSet)
+function determine_eltype(ls::LoopSet)::Union{Symbol,Expr}
     if length(ls.includedactualarrays) == 0
-        return Expr(:call, lv(:typeof), 0)
+        if length(ls.outer_reductions) == 0
+            return Expr(:call, lv(:typeof), 0)
+        elseif length(ls.outer_reductions) == 1
+            op = ls.operations[only(ls.outer_reductions)]
+            return typeof_outer_reduction(ls, op)
+        else
+            pt = Expr(:call, lv(:promote_type))
+            for j ∈ ls.outer_reductions
+                push!(pt.args, typeof_outer_reduction(ls, ls.operations[j]))
+            end
+            return pt
+        end
     elseif length(ls.includedactualarrays) == 1
         return Expr(:call, lv(:eltype), first(ls.includedactualarrays))
     end
@@ -783,6 +825,7 @@ end
 function define_eltype_vec_width!(q::Expr, ls::LoopSet, vectorized)
     push!(q.args, Expr(:(=), ELTYPESYMBOL, determine_eltype(ls)))
     push!(q.args, Expr(:(=), VECTORWIDTHSYMBOL, determine_width(ls, vectorized)))
+    nothing
 end
 function setup_preamble!(ls::LoopSet, us::UnrollSpecification, Ureduct::Int)
     @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
@@ -791,7 +834,7 @@ function setup_preamble!(ls::LoopSet, us::UnrollSpecification, Ureduct::Int)
     u₂loopsym = order[u₂loopnum]
     vectorized = order[vectorizedloopnum]
     set_vector_width!(ls, vectorized)
-    iszero(length(ls.includedactualarrays)) || define_eltype_vec_width!(ls.preamble, ls, vectorized)
+    iszero(length(ls.includedactualarrays) + length(ls.outer_reductions)) || define_eltype_vec_width!(ls.preamble, ls, vectorized)
     lower_licm_constants!(ls)
     isone(num_loops(ls)) || pushpreamble!(ls, definemask(getloop(ls, vectorized)))#, u₁ > 1 && u₁loopnum == vectorizedloopnum))
     initialize_outer_reductions!(ls, 0, Ureduct, vectorized)

src/map.jl

@@ -95,9 +95,12 @@ function vmap_multithreaded!(
     V = VectorizationBase.pick_vector_width_val(T)
     Wsh = Wshift + 2
     Niter = N >>> Wsh
-    Base.Threads.@threads :static for j ∈ 0:Niter-1
-        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((j << Wsh,))
-        vstorent!(ptry, f(vload.(ptrargs, index)...), index)
+    let Wsh = Wsh, ptry = ptry, ptrargs = ptrargs
+        Base.Threads.@threads :static for j ∈ 0:Niter-1
+            W = VectorizationBase.pick_vector_width(eltype(ptry))
+            index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((j << Wsh,))
+            vstorent!(ptry, f(vload.(ptrargs, index)...), index)
+        end
     end
     ii = Niter << Wsh
     while ii < N - (W - 1) # stops at 16 when
@@ -110,34 +113,59 @@ function vmap_multithreaded!(
     end
     y
 end
-function vmap_multithreaded!(
+struct VmapClosure{F,D,N,A<:Tuple{Vararg{Any,N}}}
+    f::F
+    dest::D
+    args::A
+end
+(m::VmapClosure)() = vmap_singlethread!(m.f, m.dest, Val{false}(), m.args...)
+@generated function vmap_multithreaded!(
     f::F,
     y::AbstractArray{T},
     ::Val{false},
     args::Vararg{AbstractArray,A}
 ) where {F,T,A}
-    N = length(y)
-    ptry = VectorizationBase.zstridedpointer(y)
-    ptrargs = VectorizationBase.zstridedpointer.(args)
-    N > 0 || return y
-    W, Wshift = VectorizationBase.pick_vector_width_shift(T)
-    V = VectorizationBase.pick_vector_width_val(T)
-    Wsh = Wshift + 2
-    Niter = N >>> Wsh
-    Base.Threads.@threads :static for j ∈ 0:Niter-1
-        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((j << Wsh,))
-        vnoaliasstore!(ptry, f(vload.(ptrargs, index)...), index)
-    end
-    ii = Niter << Wsh
-    while ii < N - (W - 1) # stops at 16 when
-        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(ii),),))...), (MM{W}(ii),))
-        ii = vadd_fast(ii, W)
-    end
-    if ii < N
-        m = mask(T, N & (W - 1))
-        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(ii),),), m)...), (MM{W}(ii),), m)
+    quote
+        N = length(y)
+        nt = min(Threads.nthreads(), $(Sys.CPU_THREADS))
+        W, Wshift = VectorizationBase.pick_vector_width_shift(T)
+        (((W * nt < N) & (nt > 1)) && iszero(ccall(:jl_in_threaded_region, Cint, ()))) || return vmap_singlethread!(f, y, Val{false}(), args...)
+        Nd, Nr = divrem(N >>> Wshift, nt)
+        Ndb = Nd << Wshift
+        Ndbr = Ndb + W
+        Nlast = N - Ndbr * Nr - Ndb * (nt - 1  - Nr)
+        yfi = firstindex(y);
+        Base.Cartesian.@nexprs $A a -> begin
+            args_a = args[a]
+            argsfi_a = firstindex(args_a);
+        end
+        lb = 0
+        # tasks = Vector{Task}(undef, nt)
+        # tasks = Base.Cartesian.@ntuple $(Sys.CPU_THREADS) t -> Ref{Task}()
+        Base.Cartesian.@nexprs $(Sys.CPU_THREADS) j -> begin
+        # for j ∈ Base.OneTo(nt)
+            Nlen = j == nt ? Nlast : (j > Nr ? Ndb : Ndbr)
+            ub = lb + Nlen
+            yv = view(y, yfi+lb:yfi+ub-1)
+            argsview = Base.Cartesian.@ntuple $A a -> view(args_a, argsfi_a+lb:argsfi_a+ub-1)
+            t_j = Task(VmapClosure(f, yv, argsview))
+            # tasks[j] = t
+            t_j.sticky = true
+            ccall(:jl_set_task_tid, Cvoid, (Any, Cint), t_j, (j == nt ? 0 : j) % Cint)
+            schedule(t_j)
+            j == nt && @goto WAIT
+            lb = ub
+        end
+        @label WAIT
+        Base.Cartesian.@nexprs $(Sys.CPU_THREADS) j -> begin
+            wait(t_j)
+            j == nt && return y
+        end
+        # for j ∈ Base.OneTo(nt)
+        #     wait(tasks[j])
+            # end
+        y
     end
-    y
 end
 
 Base.@pure _all_dense(::ArrayInterface.DenseDims{D}) where {D} = all(D)

src/reconstruct_loopset.jl

@@ -393,6 +393,7 @@ end
 # elbytes(::VectorizationBase.AbstractPointer{T}) where {T} = sizeof(T)::Int
 typeeltype(::Type{P}) where {T,P<:VectorizationBase.AbstractStridedPointer{T}} = T
 typeeltype(::Type{VectorizationBase.FastRange{T,F,S,O}}) where {T,F,S,O} = T
+typeeltype(::Type{T}) where {T<:Real} = T
 # typeeltype(::Any) = Int8
 
 function add_array_symbols!(ls::LoopSet, arraysymbolinds::Vector{Symbol}, offset::Int)

test/loopinductvars.jl

@@ -0,0 +1,31 @@
+
+# testset for using them in loops
+@testset "Loop Induction Variables" begin
+    f(x) = cos(x) * log(x)
+    function avxmax(v)
+        max_x = -Inf
+        @avx for i ∈ eachindex(v)
+            x = f(i)
+            max_x = max(max_x, x)
+        end
+        max_x
+    end
+    function avxextrema(v)
+        max_x = -Inf
+        min_x = Inf
+        @avx for i ∈ eachindex(v)
+            x = f(i)
+            max_x = max(max_x, x)
+            min_x = min(min_x, x)
+        end
+        min_x, max_x
+    end
+
+    v = 1:19
+    minref, maxref = extrema(f, v)
+    @test maxref ≈ avxmax(v)
+    minavx, maxavx = avxextrema(v);
+    @test minref ≈ minavx
+    @test maxref ≈ maxavx
+end
+

test/runtests.jl

@@ -27,6 +27,8 @@ const START_TIME = time()
 
     @time include("check_empty.jl")
 
+    @time include("loopinductvars.jl")
+    
     @time include("zygote.jl")
 
     @time include("offsetarrays.jl")