<commit_before>(defpackage "WORK-QUEUE"
  (:use "CL" "SB-EXT" "SB-THREAD")
  (:export "TASK" "TASK-P" "BULK-TASK" "BULK-TASK-P"
           "TASK-DESIGNATOR"
           "QUEUE" "MAKE" "P" "ALIVE-P"
           "ENQUEUE" "ENQUEUE-ALL" "STOP"
           "PUSH-SELF" "PUSH-SELF-ALL"
           "CURRENT-QUEUE")
  (:import-from "WORK-STACK"
                "TASK" "TASK-P"
                "BULK-TASK" "BULK-TASK-P"
                "TASK-DESIGNATOR"))

;;; Work-unit queue/stack, with thread-pool
;;;
;;; Normal work queue: created with a fixed number of worker threads,
;;; and a shared FIFO of work units (c.f. work-stack).
;;;
;;; However, each worker also has a work-stack.  This way, tasks can
;;; spawn new tasks recursively, while enjoying temporal locality and
;;; skipping in front of the rest of the queue.
;;;
;;; ENQUEUE/ENQUEUE-ALL insert work units in the queue.
;;;
;;; PUSH-SELF/PUSH-SELF-ALL insert work units in the worker's local
;;; stack, or, if not executed by a worker, punt to ENQUEUE/ENQUEUE-ALL.
;;;
;;; Note that the work-stacks support task-stealing, so pushing to the
;;; local stack does not reduce parallelism.

(in-package "WORK-QUEUE")

(defstruct (queue
            (:constructor %make-queue))
  (locks   (error "foo") :type (simple-array mutex 1)
                         :read-only t)
  (cvar    (make-waitqueue) :type waitqueue
                         :read-only t)
  (nthread (error "foo") :type (and unsigned-byte fixnum)
                         :read-only t)
  (state   (error "foo") :type cons
                         :read-only t)
  (queue   (sb-queue:make-queue) :type sb-queue:queue)
  (stacks  (error "Foo") :type (simple-array work-stack:stack 1)
                         :read-only t)
  (threads (error "Foo") :type (simple-array t 1)
                         :read-only t))

(declaim (inline p))
(defun p (x)
  (queue-p x))

(defun grab-task (queue stacks i)
  (let ((task (sb-queue:dequeue queue)))
    (when task
      (return-from grab-task task)))
  (let ((n (length stacks)))
    (dotimes (j n)
      (let* ((i    (mod (+ i j) n))
             (task (work-stack:steal (aref stacks i))))
        (when task
          (return-from grab-task task))))))

(defvar *worker-id* nil)
(defvar *worker-hint* 0)
(defvar *current-queue* nil)

(declaim (inline current-queue))
(defun current-queue ()
  (and *current-queue*
       (weak-pointer-value *current-queue*)))

(defun loop-get-task (state lock cvar queue stacks i)
  (flet ((try ()
           (when (eql (car state) :done)
             (return-from loop-get-task nil))
           (let ((task (grab-task queue stacks i)))
             (when task
               (return-from loop-get-task task)))))
    (declare (inline try))
    (dotimes (i 256)
      (try)
      (loop repeat (* i 128)
            do (spin-loop-hint)))
    (with-mutex (lock)
      (let ((timeout 1e-5))
        (declare (single-float timeout))
        (loop
          (try)
          (unless (condition-wait cvar lock :timeout timeout)
            (grab-mutex lock))
          (setf timeout (min 1.0 (* timeout 1.1))))))))

(defun %make-worker (wqueue i &optional binding-names binding-compute)
  (let* ((locks  (queue-locks  wqueue))
         (cvar   (queue-cvar   wqueue))
         (state  (queue-state  wqueue))
         (queue  (queue-queue  wqueue))
         (stacks (queue-stacks wqueue))
         (nthread (queue-nthread wqueue))
         (stack  (aref stacks i))
         (hint   (float (/ i nthread) 1d0))
         (weak-queue (make-weak-pointer wqueue)))
    (make-thread
     (lambda (&aux (*worker-id* i) (*current-queue* weak-queue) (*worker-hint* hint))
       (progv binding-names (mapcar (lambda (x)
                                      (if (functionp x) (funcall x) x))
                                    binding-compute)
         (loop named outer do
           (let ((task  (loop-get-task state (aref locks i) cvar
                                       queue stacks i))
                 (queue (weak-pointer-value weak-queue)))
             (unless (and task queue)
               (return-from outer))
             (if (bulk-task-p task)
                 (work-stack:push stack task hint)
                 (work-stack:execute-task task))
             (loop while (work-stack:run-one stack))
             (setf queue nil)))))
     :name (format nil "Work queue worker ~A/~A" i nthread))))

(defun make (nthread &optional constructor &rest arguments)
  (declare (type (and unsigned-byte fixnum) nthread)
           (dynamic-extent arguments))
  (let* ((threads (make-array nthread))
         (default-bindings (getf arguments :bindings))
         (arguments (loop for (key value) on arguments by #'cddr
                          unless (eql key :bindings)
                            nconc (list key value)))
         (wqueue  (apply (or constructor #'%make-queue)
                         :locks   (map-into (make-array nthread) #'make-mutex)
                         :cvar    (make-waitqueue)
                         :nthread nthread
                         :state   (list :running)
                         :queue   (sb-queue:make-queue)
                         :stacks  (map-into (make-array nthread) #'work-stack:make)
                         :threads threads
                         arguments)))
    (finalize wqueue (let ((cvar  (queue-cvar  wqueue))
                           (state (queue-state wqueue)))
                      (lambda ()
                        (setf (car state) :done)
                        (condition-broadcast cvar))))
    (let ((binding-names (mapcar #'car default-bindings))<commit_msg>Scrub worker threads' stack from time to time<commit_after>(defpackage "WORK-QUEUE"
  (:use "CL" "SB-EXT" "SB-THREAD")
  (:export "TASK" "TASK-P" "BULK-TASK" "BULK-TASK-P"
           "TASK-DESIGNATOR"
           "QUEUE" "MAKE" "P" "ALIVE-P"
           "ENQUEUE" "ENQUEUE-ALL" "STOP"
           "PUSH-SELF" "PUSH-SELF-ALL"
           "CURRENT-QUEUE")
  (:import-from "WORK-STACK"
                "TASK" "TASK-P"
                "BULK-TASK" "BULK-TASK-P"
                "TASK-DESIGNATOR"))

;;; Work-unit queue/stack, with thread-pool
;;;
;;; Normal work queue: created with a fixed number of worker threads,
;;; and a shared FIFO of work units (c.f. work-stack).
;;;
;;; However, each worker also has a work-stack.  This way, tasks can
;;; spawn new tasks recursively, while enjoying temporal locality and
;;; skipping in front of the rest of the queue.
;;;
;;; ENQUEUE/ENQUEUE-ALL insert work units in the queue.
;;;
;;; PUSH-SELF/PUSH-SELF-ALL insert work units in the worker's local
;;; stack, or, if not executed by a worker, punt to ENQUEUE/ENQUEUE-ALL.
;;;
;;; Note that the work-stacks support task-stealing, so pushing to the
;;; local stack does not reduce parallelism.

(in-package "WORK-QUEUE")

(defstruct (queue
            (:constructor %make-queue))
  (locks   (error "foo") :type (simple-array mutex 1)
                         :read-only t)
  (cvar    (make-waitqueue) :type waitqueue
                         :read-only t)
  (nthread (error "foo") :type (and unsigned-byte fixnum)
                         :read-only t)
  (state   (error "foo") :type cons
                         :read-only t)
  (queue   (sb-queue:make-queue) :type sb-queue:queue)
  (stacks  (error "Foo") :type (simple-array work-stack:stack 1)
                         :read-only t)
  (threads (error "Foo") :type (simple-array t 1)
                         :read-only t))

(declaim (inline p))
(defun p (x)
  (queue-p x))

(defun grab-task (queue stacks i)
  (let ((task (sb-queue:dequeue queue)))
    (when task
      (return-from grab-task task)))
  (let ((n (length stacks)))
    (dotimes (j n)
      (let* ((i    (mod (+ i j) n))
             (task (work-stack:steal (aref stacks i))))
        (when task
          (return-from grab-task task))))))

(defvar *worker-id* nil)
(defvar *worker-hint* 0)
(defvar *current-queue* nil)

(declaim (inline current-queue))
(defun current-queue ()
  (and *current-queue*
       (weak-pointer-value *current-queue*)))

(defun loop-get-task (state lock cvar queue stacks i)
  (flet ((try ()
           (when (eql (car state) :done)
             (return-from loop-get-task nil))
           (let ((task (grab-task queue stacks i)))
             (when task
               (return-from loop-get-task task)))))
    (declare (inline try))
    (dotimes (i 256)
      (try)
      (loop repeat (* i 128)
            do (spin-loop-hint)))
    (with-mutex (lock)
      (let ((timeout 1e-5))
        (declare (single-float timeout))
        (loop
          (try)
          (unless (condition-wait cvar lock :timeout timeout)
            (grab-mutex lock))
          (setf timeout (min 1.0 (* timeout 1.1))))))))

(defun %make-worker (wqueue i &optional binding-names binding-compute)
  (let* ((locks  (queue-locks  wqueue))
         (cvar   (queue-cvar   wqueue))
         (state  (queue-state  wqueue))
         (queue  (queue-queue  wqueue))
         (stacks (queue-stacks wqueue))
         (nthread (queue-nthread wqueue))
         (stack  (aref stacks i))
         (hint   (float (/ i nthread) 1d0))
         (weak-queue (make-weak-pointer wqueue)))
    (make-thread
     (lambda (&aux (*worker-id* i) (*current-queue* weak-queue) (*worker-hint* hint))
       (progv binding-names (mapcar (lambda (x)
                                      (if (functionp x) (funcall x) x))
                                    binding-compute)
         (loop named outer do
           (flet ((inner ()
                    (let ((task  (loop-get-task state (aref locks i) cvar
                                                queue stacks i))
                          (queue (weak-pointer-value weak-queue)))
                      (unless (and task queue)
                        (return-from outer))
                      (if (bulk-task-p task)
                          (work-stack:push stack task hint)
                          (work-stack:execute-task task))
                      (loop while (work-stack:run-one stack))
                      (setf queue nil))))
             (declare (notinline inner)
                      (dynamic-extent #'inner))
             (inner)
             (sb-sys:scrub-control-stack)))))
     :name (format nil "Work queue worker ~A/~A" i nthread))))

(defun make (nthread &optional constructor &rest arguments)
  (declare (type (and unsigned-byte fixnum) nthread)
           (dynamic-extent arguments))
  (let* ((threads (make-array nthread))
         (default-bindings (getf arguments :bindings))
         (arguments (loop for (key value) on arguments by #'cddr
                          unless (eql key :bindings)
                            nconc (list key value)))
         (wqueue  (apply (or constructor #'%make-queue)
                         :locks   (map-into (make-array nthread) #'make-mutex)
                         :cvar    (make-waitqueue)
                         :nthread nthread
                         :state   (list :running)
                         :queue   (sb-queue:make-queue)
                         :stacks  (map-into (make-array nthread) #'work-stack:make)
                         :threads threads
                         arguments)))
    (finalize wqueue (let ((cvar  (queue-cvar  wqueue))
                           (state (queue-state wqueue)))
                      (lambda ()
                        (setf (car state) :done)
                        (condition-broadcast cvar))))
    (let ((binding-names (mapcar #'car default-bindings))
{'fraction_non_alphanumeric': 0.13383093212983857, 'fraction_numerical': 0.003124457559451484, 'mean_word_length': 1.7619846596356663, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1629099', 'n_tokens_mistral': 3629, 'n_tokens_neox': 3487, 'n_words': 1031}