Merge pull request #139 from GianlucaFuwa/master

Fix poor performance when one of `batch`'s arguments is a `Type{...}`

Author: chriselrod

src/batch.jl

@@ -1,7 +1,7 @@
 # S is a Val{Bool} indicating whether we will need to load the thread index
 # C is a Tuple{...} containing the types of the reduction variables
 struct BatchClosure{F,A,S,C}
-    f::F
+  f::F
 end
 function (b::BatchClosure{F,A,S,C})(p::Ptr{UInt}) where {F,A,S,C}
   (offset, args) = ThreadingUtilities.load(p, A, 2 * sizeof(UInt))
@@ -130,7 +130,7 @@ end
     start,
     stop,
     i,
-    reductup
+    reductup,
   ) do p, fptr, argtup, start, stop, i, reductup
     setup_batch!(p, fptr, argtup, start, stop, i, reductup)
   end
@@ -185,12 +185,9 @@ end
     # nthread_total = sum(nthread_tuple)
     Ndp = Nd + one(Nd)
   end
-  C !== 0 && push!(
-    q.args,
-    quote
-      @nexprs $C j -> RVAR_j = reducinits[j]
-    end
-  )
+  C !== 0 && push!(q.args, quote
+    @nexprs $C j -> RVAR_j = reducinits[j]
+  end)
   launch_quote = if S
     if C === 0
       :(launch_batched_thread!(cfunc, tid, argtup, start, stop, tid % UInt))
@@ -211,27 +208,23 @@ end
   if C !== 0
     push!(
       rem_quote.args,
-      :(@nexprs $C j -> RVAR_j = reducops[j](RVAR_j, thread_results[j]))
+      :(@nexprs $C j -> RVAR_j = reducops[j](RVAR_j, thread_results[j])),
     )
   end
   update_retv = if C === 0
-      Expr(:block)
+    Expr(:block)
   else
     quote
       thread_results = load_threadlocals(tid, argtup, needtid, reducinits)
       @nexprs $C j -> RVAR_j = reducops[j](RVAR_j, thread_results[j])
     end
   end
   ret_quote = Expr(:return)
-  if C === 0
-    push!(ret_quote.args, nothing)
-  else
-    redtup = Expr(:tuple)
-    for j in 1:C
-      push!(redtup.args, Symbol("RVAR_", j))
-    end
-    push!(ret_quote.args, redtup)
+  redtup = Expr(:tuple)
+  for j ∈ 1:C
+    push!(redtup.args, Symbol("RVAR_", j))
   end
+  push!(ret_quote.args, redtup)
 
   block = quote
     start = zero(UInt)
@@ -321,16 +314,16 @@ end
     @label SERIAL
     if S
       if C === 0
-        reducres = f!(args, one(Int), ulen % Int, 1)
-        return reducres
+        f!(args, one(Int), ulen % Int, 1)
+        return ()
       else
         reducres = f!(args, one(Int), ulen % Int, 1, reducinits)
         return reducres
       end
     else
       if C === 0
-        reducres = f!(args, one(Int), ulen % Int)
-        return reducres
+        f!(args, one(Int), ulen % Int)
+        return ()
       else
         reducres = f!(args, one(Int), ulen % Int, reducinits)
         return reducres

src/closure.jl

@@ -207,6 +207,13 @@ Base.@propagate_inbounds combine(x::AbstractArray, I, j) =
   x[combine(CombineIndices(), I, j)]
 Base.@propagate_inbounds combine(x::AbstractArray, ::NoLoop, j) = x[j]
 
+struct WrapType{T} end
+wrap_type(@nospecialize(x)) = x
+wrap_type(::Type{T}) where {T} = WrapType{T}()
+
+unwrap_type(@nospecialize(x)) = x
+unwrap_type(::WrapType{T}) where {T} = T
+
 function makestatic!(expr)
   expr isa Expr || return expr
   for i in eachindex(expr.args)
@@ -421,9 +428,11 @@ function enclose(exorig::Expr, minbatchsize, per, threadlocal, reduction, stride
     loop_stop_expr =
       :(var"##SUBSTOP##" * var"##LOOP_STEP##" + var"##LOOPOFFSET##" - var"##LOOP_STEP##")
   end
+  unwrap_args = Expr(:block)
   closureq = quote
     $closure = let
       @inline $closure_args -> begin
+        $unwrap_args
         local var"##STEP##" = $(
           stride ?
           :($loop_step * min(Threads.nthreads()::Int, Sys.CPU_THREADS::Int)) :
@@ -470,12 +479,15 @@ function enclose(exorig::Expr, minbatchsize, per, threadlocal, reduction, stride
   end
   for a ∈ arguments
     push!(args.args, get(defined, a, a))
-    push!(batchcall.args, esc(a))
+    push!(batchcall.args, :($wrap_type($(esc(a)))))
   end
   if threadlocal_val !== Symbol("")
     push!(args.args, threadlocal_accum)
     push!(batchcall.args, esc(threadlocal_accum))
   end
+  for a in args.args
+    push!(unwrap_args.args, :($a = $unwrap_type($a)))
+  end
   push!(q.args, batchcall)
   quote
     var"##NUM#THREADS##" = $(Threads.nthreads)()

test/runtests.jl

@@ -112,6 +112,12 @@ function issue25_but_with_strides!(dest, x, y)
   dest
 end
 
+function issue116!(y::Vector{T}, x::Vector{T}) where {T}
+  @batch for i in 1:length(x)
+      y[i] = exp(x[i] + one(T))
+  end
+end
+
 
 @testset "Range Map" begin
 
@@ -297,6 +303,15 @@ Base.eachindex(e::Iterators.Enumerate{LazyTree{T}}) where {T} = eachindex(e.itr)
   evt = 5
 end
 
+@testset "not-specializing-on-type heuristics" begin
+  allocated(f::F, args...) where {F} = @allocated f(args...)
+  x = rand(10000)
+  y = similar(x)
+  allocated(issue116!, y, x)
+  @test y ≈ exp.(x .+ 1.0)
+  @test allocated(issue116!, y, x) == 0
+end
+
 @testset "threadlocal storage" begin
   local1 = let
     @batch threadlocal = 0 for i = 0:9
@@ -338,15 +353,13 @@ end
   end
   @test local1 == local2 == local3 == local4 == local5 == local6
   # check that each thread has a separate init
-  myvar = 0
-  myinitC() = myvar += 1
   inits = let
-    @batch threadlocal = myinitC() for i = 0:9
+    @batch threadlocal = rand() for i = 0:9
       threadlocal += 1
     end
     threadlocal
   end
-  @test length(inits) == 1 || inits[1] != inits[end] # this test has a race condition and can (rarely) fail
+  @test length(inits) == 1 || inits[1] != inits[end]
   # check that types are respected
   myinitD() = Float16(1.0)
   settingtype = let
@@ -461,14 +474,15 @@ end
   # check for name interference with threadlocal (used to error on single threaded runs)
   function f()
     n = 1000
-    threadlocal = false
+    threadlocal = 0
     @batch minbatch = 10 reduction = (+,threadlocal) for i = 1:n
-      threadlocal += true
+      threadlocal += 1
     end
     return threadlocal
   end
   allocated(f::F) where {F} = @allocated f()
   inferred(f::F) where {F} = try @inferred f(); true catch; false end
+  allocated(f)
   @test allocated(f) == 0
   @test inferred(f) == true
   # remaining supported operations