Crime doesn't pay. (Stop stealing.)

Author: chriselrod

src/ThreadingUtilities.jl

@@ -4,10 +4,9 @@ using VectorizationBase:
     pause, StaticInt, StridedPointer, stridedpointer, offsets, cache_linesize, align, __vload, __vstore!, num_threads, assume, False, register_size, NativeTypes
 
 @enum ThreadState::UInt32 begin
-  TASK = 0   # 3: task available
-  EXEC = 1   # 2: task executed
-  WAIT = 2   # 1: waiting
-  SPIN = 3   # 0: spinning
+  TASK = 0   # 0: task available
+  WAIT = 1   # 1: waiting
+  SPIN = 2   # 2: spinning
 end
 const TASKS = Task[]
 const THREADBUFFERSIZE = 512

src/threadtasks.jl

@@ -1,51 +1,51 @@
 struct ThreadTask
-    p::Ptr{UInt}
+  p::Ptr{UInt}
 end
 Base.pointer(tt::ThreadTask) = tt.p
 
 @inline taskpointer(tid::T) where {T} = THREADPOOLPTR[] + tid*(THREADBUFFERSIZE%T)
 
 @inline function _call(p::Ptr{UInt})
-    fptr = load(p + sizeof(UInt), Ptr{Cvoid})
-    assume(fptr ≠ C_NULL)
-    ccall(fptr, Cvoid, (Ptr{UInt},), p)
+  fptr = load(p + sizeof(UInt), Ptr{Cvoid})
+  assume(fptr ≠ C_NULL)
+  ccall(fptr, Cvoid, (Ptr{UInt},), p)
 end
 @inline function launch(f::F, tid::Integer, args::Vararg{Any,K}) where {F,K}
-    p = taskpointer(tid)
-    f(p, args...)
-    state = _atomic_xchg!(p, TASK)
-    state == WAIT && wake_thread!(tid)
-    return nothing
+  p = taskpointer(tid)
+  f(p, args...)
+  state = _atomic_xchg!(p, TASK) # exchange must happen atomically, to prevent it from switching to `WAIT` after reading
+  state == WAIT && wake_thread!(tid)
+  return nothing
 end
 
 function (tt::ThreadTask)()
-    p = pointer(tt)
-    max_wait = one(UInt32) << 20
-    wait_counter = max_wait
-    GC.@preserve THREADPOOL begin
-        while true
-            # if _atomic_state(p) == TASK
-            if _atomic_cas_cmp!(p, TASK, EXEC)
-                _call(p)
-              # store!(p, SPIN)
-              _atomic_store!(p, SPIN)
-                wait_counter = zero(UInt32)
-                continue
-            end
-            pause()
-            if (wait_counter += one(UInt32)) > max_wait
-                wait_counter = zero(UInt32)
-                _atomic_cas_cmp!(p, SPIN, WAIT) && Base.wait()
-            end
-        end
+  p = pointer(tt)
+  max_wait = one(UInt32) << 20
+  wait_counter = max_wait
+  GC.@preserve THREADPOOL begin
+    while true
+      if _atomic_state(p) == TASK
+        # if _atomic_cas_cmp!(p, TASK, EXEC)
+        _call(p)
+        # store!(p, SPIN)
+        _atomic_store!(p, SPIN)
+        wait_counter = zero(UInt32)
+        continue
+      end
+      pause()
+      if (wait_counter += one(UInt32)) > max_wait
+        wait_counter = zero(UInt32)
+        _atomic_cas_cmp!(p, SPIN, WAIT) && Base.wait()
+      end
     end
+  end
 end
 
 # 1-based tid, pushes into task 2-nthreads()
 # function wake_thread!(tid::T) where {T <: Unsigned}
 @noinline function wake_thread!(_tid::T) where {T <: Integer}
   tid = _tid % Int
-  store!(taskpointer(_tid), TASK)
+  # store!(taskpointer(_tid), TASK)
   tidp1 = tid + one(tid)
   assume(unsigned(length(Base.Workqueues)) > unsigned(tid))
   assume(unsigned(length(TASKS)) > unsigned(tidp1))
@@ -56,29 +56,10 @@ end
 # 1-based tid
 @inline wait(tid::Integer) = wait(taskpointer(tid))
 @inline function wait(p::Ptr{UInt})
-  # TASK = 0
-  # EXEC = 1
-  # WAIT = 2
-  # SPIN = 3
-  s = _atomic_umax!(p, EXEC) # s = old state, state gets set to EXEC if s == TASK or s == EXEC
-  if s == TASK # thread hasn't begun yet for some reason, so we steal the work.
-    _call(p)
-    store!(p, SPIN)
-    return
-  end
   counter = 0x00000000
-  while _atomic_state(p) == EXEC
+  while _atomic_state(p) == TASK
     pause()
     ((counter += 0x00000001) > 0x00010000) && yield()
   end
 end
 
-
-# function launch_thread(f::F, tid) where {F}
-#     cfunc = @cfunction($mapper, Cvoid, ());
-
-#     fptr = Base.unsafe_convert(Ptr{Cvoid}, cfunc)
-
-#     ccall(fptr, Cvoid, ())
-
-# end

test/threadingutilities.jl

@@ -53,7 +53,7 @@ end
 
 @testset "ThreadingUtilities.jl" begin
   for tid ∈ eachindex(ThreadingUtilities.TASKS)
-    @test unsafe_load(Ptr{UInt32}(ThreadingUtilities.taskpointer(tid))) == 0x00000002
+    @test unsafe_load(Ptr{UInt32}(ThreadingUtilities.taskpointer(tid))) == 0x00000001
   end
   @test all(eachindex(ThreadingUtilities.TASKS)) do tid
     ThreadingUtilities.load(ThreadingUtilities.taskpointer(tid), ThreadingUtilities.ThreadState) === ThreadingUtilities.WAIT
@@ -67,8 +67,8 @@ end
   GC.@preserve x begin
     ThreadingUtilities._atomic_store!(pointer(x), zero(UInt))
     @test ThreadingUtilities._atomic_xchg!(pointer(x), ThreadingUtilities.WAIT) == ThreadingUtilities.TASK
+    @test ThreadingUtilities._atomic_umax!(pointer(x), ThreadingUtilities.TASK) == ThreadingUtilities.WAIT
     @test ThreadingUtilities._atomic_umax!(pointer(x), ThreadingUtilities.SPIN) == ThreadingUtilities.WAIT
-    @test ThreadingUtilities._atomic_umax!(pointer(x), ThreadingUtilities.EXEC) == ThreadingUtilities.SPIN
     @test ThreadingUtilities.load(pointer(x), ThreadingUtilities.ThreadState) == ThreadingUtilities.SPIN
   end
   for tid ∈ eachindex(ThreadingUtilities.TASKS)