|
| 1 | +# This file is a part of Julia. License is MIT: https://julialang.org/license |
| 2 | + |
| 3 | +module Partr |
| 4 | + |
| 5 | +using ..Threads: SpinLock |
| 6 | + |
| 7 | +# a task heap |
| 8 | +mutable struct taskheap |
| 9 | + const lock::SpinLock |
| 10 | + const tasks::Vector{Task} |
| 11 | + @atomic ntasks::Int32 |
| 12 | + @atomic priority::UInt16 |
| 13 | + taskheap() = new(SpinLock(), Vector{Task}(undef, tasks_per_heap), zero(Int32), typemax(UInt16)) |
| 14 | +end |
| 15 | + |
| 16 | +# multiqueue parameters |
| 17 | +const heap_d = UInt32(8) |
| 18 | +const heap_c = UInt32(2) |
| 19 | + |
| 20 | +# size of each heap |
| 21 | +const tasks_per_heap = Int32(65536) # TODO: this should be smaller by default, but growable! |
| 22 | + |
| 23 | +# the multiqueue's heaps |
| 24 | +global heaps::Vector{taskheap} |
| 25 | +global heap_p::UInt32 = 0 |
| 26 | + |
| 27 | +# unbias state for the RNG |
| 28 | +global cong_unbias::UInt32 = 0 |
| 29 | + |
| 30 | + |
| 31 | +cong(max::UInt32, unbias::UInt32) = ccall(:jl_rand_ptls, UInt32, (UInt32, UInt32), max, unbias) + UInt32(1) |
| 32 | + |
| 33 | +function unbias_cong(max::UInt32) |
| 34 | + return typemax(UInt32) - ((typemax(UInt32) % max) + UInt32(1)) |
| 35 | +end |
| 36 | + |
| 37 | + |
| 38 | +function multiq_init(nthreads) |
| 39 | + global heap_p = heap_c * nthreads |
| 40 | + global cong_unbias = unbias_cong(UInt32(heap_p)) |
| 41 | + global heaps = Vector{taskheap}(undef, heap_p) |
| 42 | + for i = UInt32(1):heap_p |
| 43 | + heaps[i] = taskheap() |
| 44 | + end |
| 45 | + nothing |
| 46 | +end |
| 47 | + |
| 48 | + |
| 49 | +function sift_up(heap::taskheap, idx::Int32) |
| 50 | + while idx > Int32(1) |
| 51 | + parent = (idx - Int32(2)) ÷ heap_d + Int32(1) |
| 52 | + if heap.tasks[idx].priority < heap.tasks[parent].priority |
| 53 | + t = heap.tasks[parent] |
| 54 | + heap.tasks[parent] = heap.tasks[idx] |
| 55 | + heap.tasks[idx] = t |
| 56 | + idx = parent |
| 57 | + else |
| 58 | + break |
| 59 | + end |
| 60 | + end |
| 61 | +end |
| 62 | + |
| 63 | + |
| 64 | +function sift_down(heap::taskheap, idx::Int32) |
| 65 | + if idx <= heap.ntasks |
| 66 | + for child = (heap_d * idx - heap_d + Int32(2)):(heap_d * idx + Int32(1)) |
| 67 | + child > tasks_per_heap && break |
| 68 | + if isassigned(heap.tasks, child) && |
| 69 | + heap.tasks[child].priority < heap.tasks[idx].priority |
| 70 | + t = heap.tasks[idx] |
| 71 | + heap.tasks[idx] = heap.tasks[child] |
| 72 | + heap.tasks[child] = t |
| 73 | + sift_down(heap, child) |
| 74 | + end |
| 75 | + end |
| 76 | + end |
| 77 | +end |
| 78 | + |
| 79 | + |
| 80 | +function multiq_insert(task::Task, priority::UInt16) |
| 81 | + task.priority = priority |
| 82 | + |
| 83 | + rn = cong(heap_p, cong_unbias) |
| 84 | + while !trylock(heaps[rn].lock) |
| 85 | + rn = cong(heap_p, cong_unbias) |
| 86 | + end |
| 87 | + |
| 88 | + if heaps[rn].ntasks >= tasks_per_heap |
| 89 | + unlock(heaps[rn].lock) |
| 90 | + # multiq insertion failed, increase #tasks per heap |
| 91 | + return false |
| 92 | + end |
| 93 | + |
| 94 | + ntasks = heaps[rn].ntasks + Int32(1) |
| 95 | + @atomic :monotonic heaps[rn].ntasks = ntasks |
| 96 | + heaps[rn].tasks[ntasks] = task |
| 97 | + sift_up(heaps[rn], ntasks) |
| 98 | + priority = heaps[rn].priority |
| 99 | + if task.priority < priority |
| 100 | + @atomic :monotonic heaps[rn].priority = task.priority |
| 101 | + end |
| 102 | + unlock(heaps[rn].lock) |
| 103 | + return true |
| 104 | +end |
| 105 | + |
| 106 | + |
| 107 | +function multiq_deletemin() |
| 108 | + local rn1, rn2 |
| 109 | + local prio1, prio2 |
| 110 | + |
| 111 | + @label retry |
| 112 | + GC.safepoint() |
| 113 | + for i = UInt32(1):heap_p |
| 114 | + if i == heap_p |
| 115 | + return nothing |
| 116 | + end |
| 117 | + rn1 = cong(heap_p, cong_unbias) |
| 118 | + rn2 = cong(heap_p, cong_unbias) |
| 119 | + prio1 = heaps[rn1].priority |
| 120 | + prio2 = heaps[rn2].priority |
| 121 | + if prio1 > prio2 |
| 122 | + prio1 = prio2 |
| 123 | + rn1 = rn2 |
| 124 | + elseif prio1 == prio2 && prio1 == typemax(UInt16) |
| 125 | + continue |
| 126 | + end |
| 127 | + if trylock(heaps[rn1].lock) |
| 128 | + if prio1 == heaps[rn1].priority |
| 129 | + break |
| 130 | + end |
| 131 | + unlock(heaps[rn1].lock) |
| 132 | + end |
| 133 | + end |
| 134 | + |
| 135 | + task = heaps[rn1].tasks[1] |
| 136 | + tid = Threads.threadid() |
| 137 | + if ccall(:jl_set_task_tid, Cint, (Any, Cint), task, tid-1) == 0 |
| 138 | + unlock(heaps[rn1].lock) |
| 139 | + @goto retry |
| 140 | + end |
| 141 | + ntasks = heaps[rn1].ntasks |
| 142 | + @atomic :monotonic heaps[rn1].ntasks = ntasks - Int32(1) |
| 143 | + heaps[rn1].tasks[1] = heaps[rn1].tasks[ntasks] |
| 144 | + Base._unsetindex!(heaps[rn1].tasks, Int(ntasks)) |
| 145 | + prio1 = typemax(UInt16) |
| 146 | + if ntasks > 1 |
| 147 | + sift_down(heaps[rn1], Int32(1)) |
| 148 | + prio1 = heaps[rn1].tasks[1].priority |
| 149 | + end |
| 150 | + @atomic :monotonic heaps[rn1].priority = prio1 |
| 151 | + unlock(heaps[rn1].lock) |
| 152 | + |
| 153 | + return task |
| 154 | +end |
| 155 | + |
| 156 | + |
| 157 | +function multiq_check_empty() |
| 158 | + for i = UInt32(1):heap_p |
| 159 | + if heaps[i].ntasks != 0 |
| 160 | + return false |
| 161 | + end |
| 162 | + end |
| 163 | + return true |
| 164 | +end |
| 165 | + |
| 166 | +end |
0 commit comments