Split NamedTuples and add minbatch option to @batch

src/batch.jl

@@ -28,17 +28,21 @@ end
         setup_batch!(p, fptr, argtup, start, stop)
     end
 end
+_extract_params(::Type{T}) where {T<:Tuple} = T.parameters
+_extract_params(::Type{NamedTuple{S,T}}) where {S,T<:Tuple} = T.parameters
+push_tup!(x, ::Type{T}, t) where {T<:Tuple} = push!(x,t)
+push_tup!(x, ::Type{NamedTuple{S,T}}, t) where {S,T<:Tuple} = push!(x, Expr(:call, Expr(:curly, :NamedTuple, S), t))
 
 function add_var!(q, argtup, gcpres, ::Type{T}, argtupname, gcpresname, k) where {T}
     parg_k = Symbol(argtupname, :_, k)
     garg_k = Symbol(gcpresname, :_, k)
-    if T <: Tuple
+    if (T <: Tuple) || (T <: NamedTuple)
         push!(q.args, Expr(:(=), parg_k, Expr(:ref, argtupname, k)))
         t = Expr(:tuple)
-        for (j,p) ∈ enumerate(T.parameters)
+        for (j,p) ∈ enumerate(_extract_params(T))
             add_var!(q, t, gcpres, p, parg_k, garg_k, j)
         end
-        push!(argtup.args, t)
+        push_tup!(argtup.args, T, t)
     else
         push!(q.args, Expr(:(=), Expr(:tuple, parg_k, garg_k), Expr(:call, :object_and_preserve, Expr(:ref, argtupname, k))))
         push!(argtup.args, parg_k)
@@ -48,8 +52,10 @@ end
 
 @generated function _batch_no_reserve(
     f!::F, threadmask, nthread, torelease, Nr, Nd, ulen, args::Vararg{Any,K}
-) where {F,K}
+    ) where {F,K}
+    1+2
     q = quote
+        $(Expr(:meta,:inline))
         threads = UnsignedIteratorEarlyStop(threadmask, nthread)
         Ndp = Nd + one(Nd)
     end
@@ -159,7 +165,7 @@ end
 end
 
 
-function batch(
+@inline function batch(
     f!::F, (len, nbatches)::Tuple{Vararg{Integer,2}}, args::Vararg{Any,K}
 ) where {F,K}
     threads, torelease = request_threads(Base.Threads.threadid(), nbatches - one(nbatches))
@@ -170,7 +176,6 @@ function batch(
         return
     end
     nbatch = nthread + one(nthread)
-
     Nd = Base.udiv_int(ulen, nbatch % UInt) # reasonable for `ulen` to be ≥ 2^32
     Nr = ulen - Nd * nbatch
 

src/closure.jl

@@ -76,7 +76,7 @@ function extractargs!(arguments::Vector{Symbol}, defined::Set, expr::Expr)
     end
 end
 
-function enclose(exorig::Expr, reserve_per = 0)
+function enclose(exorig::Expr, reserve_per = 0, minbatchsize = 1)
     Meta.isexpr(exorig, :for, 2) || throw(ArgumentError("Expression invalid; should be a for loop."))
     ex = copy(exorig)
     loop_sym = Symbol("##LOOP##")
@@ -104,16 +104,25 @@ function enclose(exorig::Expr, reserve_per = 0)
         $loop_offs = CheapThreads.static_first($loop_sym)
     end
     threadtup = Expr(:tuple, iter_leng)
-    if reserve_per ≤ 0
-        push!(threadtup.args, :(min($iter_leng, CheapThreads.num_threads())))
+    if minbatchsize ≤ 1
+        if reserve_per ≤ 0
+            push!(threadtup.args, :(min($iter_leng, CheapThreads.num_threads())))
+        else
+            push!(threadtup.args, :(min($iter_leng, cld(CheapThreads.num_threads(), $reserve_per))), reserve_per)
+        end
     else
-        push!(threadtup.args, :(min($iter_leng, cld(CheapThreads.num_threads(), $reserve_per))), reserve_per)
+        il = :(div($iter_leng, $(StaticInt(minbatchsize))))
+        if reserve_per ≤ 0
+            push!(threadtup.args, :(min($il, CheapThreads.num_threads())))
+        else
+            push!(threadtup.args, :(min($il, cld(CheapThreads.num_threads(), $reserve_per))), reserve_per)
+        end
     end
     closure = Symbol("##closure##")
     args = Expr(:tuple, Symbol("##LOOPOFFSET##"), Symbol("##LOOP_STEP##"))
     closureq = quote
         $closure = let
-            ($args, var"##SUBSTART##"::Int, var"##SUBSTOP##") -> begin
+            @inline ($args, var"##SUBSTART##"::Int, var"##SUBSTOP##"::Int) -> begin
                 var"##LOOPSTART##" = var"##SUBSTART##" * var"##LOOP_STEP##" + var"##LOOPOFFSET##" - $(Static.One())
                 var"##LOOP_STOP##" = var"##SUBSTOP##" * var"##LOOP_STEP##" + var"##LOOPOFFSET##" - $(Static.One())
                 @inbounds begin
@@ -125,14 +134,20 @@ function enclose(exorig::Expr, reserve_per = 0)
     end
     push!(q.args, esc(closureq))
     batchcall = Expr(:call, GlobalRef(CheapThreads, :batch), esc(closure), threadtup, Symbol("##LOOPOFFSET##"), Symbol("##LOOP_STEP##"))
+    # batchcall = Expr(:call, esc(closure))
+    # t = Expr(:tuple, Symbol("##LOOPOFFSET##"), Symbol("##LOOP_STEP##"))
     for a ∈ arguments
         push!(args.args, a)
+        # push!(t.args, esc(a))
         push!(batchcall.args, esc(a))
     end
+    # push!(batchcall.args, t, 1, iter_leng)
     push!(q.args, batchcall)
     quote
         if CheapThreads.num_threads() == 1
-            $(esc(exorig))
+            let
+                $(esc(exorig))
+            end
         else
             let
                 $q
@@ -144,7 +159,27 @@ end
 macro batch(ex)
     enclose(macroexpand(__module__, ex))
 end
-macro batch(reserve_per, ex)
-    enclose(macroexpand(__module__, ex), reserve_per)
+function interpret_kwarg(arg, reserve_per = 0, minbatch = 1)
+    a = arg.args[1]
+    v = arg.args[2]
+    if a === :reserve
+        @assert v ≥ 0
+        reserve_per = v
+    elseif a === :minbatch
+        @assert v ≥ 1
+        minbatch = v
+    else
+        throw(ArgumentError("kwarg $(a) not recognized."))
+    end
+    reserve_per, minbatch
+end
+macro batch(arg1, ex)
+    reserve, minbatch = interpret_kwarg(arg1)
+    enclose(macroexpand(__module__, ex), reserve, minbatch)
+end
+macro batch(arg1, arg2, ex)
+    reserve, minbatch = interpret_kwarg(arg1)
+    reserve, minbatch = interpret_kwarg(arg2, reserve, minbatch)
+    enclose(macroexpand(__module__, ex), reserve, minbatch)
 end
 

src/request.jl

@@ -31,7 +31,7 @@ function free_local_threads!()
     free_threads!(r)
 end
 
-function _request_threads(id::UInt32, num_requested::UInt32)
+@inline function _request_threads(id::UInt32, num_requested::UInt32)
     reserved_threads = reserved(id)
     reserved_count = count_ones(reserved_threads)%UInt32
     no_threads = zero(worker_type())
@@ -65,7 +65,7 @@ function _request_threads(id::UInt32, num_requested::UInt32)
     ThreadingUtilities._atomic_store!(wp, _all_threads & (~all_threads))
     return UnsignedIteratorEarlyStop(reserved_threads | all_threads, num_requested), all_threads
 end
-function request_threads(id, num_requested)
+@inline function request_threads(id, num_requested)
     _request_threads(id % UInt32, num_requested % UInt32)
 end
 reserved_threads(id) = UnsignedIteratorEarlyStop(reserved(id))