cpu-affinity.lisp

   1 (in-package :sb-cpu-affinity)
   2
   3 ;;;; Alien definitions
   4
   5 ;; This is defined as a #DEFINE constant, but the .so stuff the value
   6 ;; into a variable for us.
   7 (eval-when (:compile-toplevel :load-toplevel)
   8   (define-alien-variable cpu-setsize int))
   9
  10 ;; This is sizeof(cpu_set_t)
  11 (define-alien-variable cpu-mask-size int)
  12
  13 ;; These are convenience wrappers around sched_set/getaffinity.
  14 (define-alien-routine ("get_cpu_affinity_mask" %get-cpu-affinity-mask) int
  15   (mask (* (unsigned 8))))
  16 (define-alien-routine ("set_cpu_affinity_mask" %set-cpu-affinity-mask) int
  17   (mask (* (unsigned 8))))
  18
  19 ;; These are wrappers around libc macros that manipulate the mask. Why
  20 ;; the hell can't C people provide functions in addition to macros
  21 ;; like this? Or just use inline functions?
  22 (define-alien-routine zero-cpu-affinity-mask void
  23   (mask (* (unsigned 8))))
  24 (define-alien-routine ("cpu_affinity_p" %cpu-affinity-p) int
  25   (cpu int)
  26   (mask (* (unsigned 8))))
  27 (define-alien-routine set-cpu-affinity  void
  28   (cpu int)
  29   (mask (* (unsigned 8))))
  30 (define-alien-routine clear-cpu-affinity void
  31   (cpu int)
  32   (mask (* (unsigned 8))))
  33
  34 ;;;; Nice lispy interface
  35
  36 (defconstant +cpu-limit+ cpu-setsize
  37   "Upper exclusive limit on the number of CPUs. Based on CPU_SETSIZE
  38 from sched.h.")
  39
  40 (defvar *cpu-count* nil)
  41
  42 (defun cpu-count ()
  43   "Number of CPUs available in the system. Based on /proc/cpuinfo."
  44   (or *cpu-count*
  45       (setf *cpu-count*
  46             (let* ((key "processor")
  47                    (len (length key)))
  48               (with-open-file (f "/proc/cpuinfo")
  49                 (loop for line = (read-line f nil nil)
  50                       while line
  51                       count (when (> (length line) len)
  52                               (string= key line :end2 len))))))))
  53
  54 (defstruct cpu-affinity-mask
  55   "CPU affinity mask."
  56   %mask)
  57
  58 (defmethod print-object ((mask cpu-affinity-mask) stream)
  59   (print-unreadable-object (mask stream :type t)
  60     ;; Print the locally interesting part of the mask.
  61     (dotimes (i (cpu-count))
  62       (if (cpu-affinity-p i mask)
  63           (write-char #\1 stream)
  64           (write-char #\0 stream)))))
  65
  66 (defun make-cpu-mask ()
  67   (make-alien (unsigned 8) cpu-mask-size))
  68
  69 (defun get-cpu-affinity-mask ()
  70   "Returns the CPU affinity mask of the current thread. The mask can
  71 be inspected and mutated using CPU-AFFINITY-P, \(SETF CPU-AFFINITY-P),
  72 and CLEAR-CPU-AFFINITY-MASK. To make any changes take effect, the mask
  73 must be saved using SET-CPU-AFFINITY-MASK.
  74
  75 Using WITH-CPU-AFFINITY-MASK instead is recommended."
  76   ;; FIXME: Malloc'ed mask is nasty, but libc doesn't seem to like
  77   ;; stack allocated ones, nor does it seem to like freeing masks that
  78   ;; have been used. So we never do. Gah.
  79   (let ((mask (make-cpu-mask)))
  80     (unless (zerop (%get-cpu-affinity-mask mask))
  81       (error "Could not read CPU affinity mask: ~A" (sb-int:strerror)))
  82     (make-cpu-affinity-mask :%mask mask)))
  83
  84 (defun set-cpu-affinity-mask (mask)
  85   "Sets the CPU affinity mask for the current thread.
  86
  87 Using WITH-CPU-AFFINITY-MASK :SAVE T instead is recommended."
  88   (unless (zerop (%set-cpu-affinity-mask (cpu-affinity-mask-%mask mask)))
  89     (error "Coud not write CPU affinity mask: ~A" (sb-int:strerror))))
  90
  91 (defun cpu-affinity-p (cpu mask)
  92   "Returns T if the CPU \(a numeric indentifier between 0 and
  93  +CPU-LIMIT+) is part of the MASK."
  94   (plusp (%cpu-affinity-p cpu (cpu-affinity-mask-%mask mask))))
  95
  96 (defun (setf cpu-affinity-p) (bool cpu mask)
  97   "Toggles presence of the CPU \(a numeric identifier between 0 and +CPU-LIMIT+)
  98 in the MASK."
  99   (let ((%mask (cpu-affinity-mask-%mask mask)))
 100     (if bool
 101         (set-cpu-affinity cpu %mask)
 102         (clear-cpu-affinity cpu %mask)))
 103   bool)
 104
 105 (defun clear-cpu-affinity-mask (mask)
 106   "Removes all CPUs from the MASK."
 107   (zero-cpu-affinity-mask (cpu-affinity-mask-%mask mask))
 108   mask)
 109
 110 (defmacro with-cpu-affinity-mask ((mask &key save) &body body)
 111   "Reads the CPU affinity mask of the the current thread and binds it
 112 to MASK. The mask can be inspected and mutated using CPU-AFFINITY-P,
 113 \(SETF CPU-AFFINITY-P, and CLEAR-CPU-AFFINITY-MASK. Any changes take
 114 effect only if SAVE is true (default is NIL)."
 115   (let ((ok-n (gensym "OK")))
 116     `(let (,mask ,ok-n)
 117        (unwind-protect
 118             (progn
 119               (setf ,mask (get-cpu-affinity-mask))
 120               (multiple-value-prog1 (progn ,@body)
 121                 (setf ,ok-n t)))
 122          (when (and ,ok-n ,save)
 123            (set-cpu-affinity-mask ,mask))))))
 124
 125 ;;;; Usage examples
 126 #+nil
 127 (progn
 128
 129   (with-cpu-affinity-mask (mask)
 130     (print mask))
 131
 132   (with-cpu-affinity-mask (mask :save t)
 133     ;; Remove all
 134     (clear-cpu-affinity-mask mask)
 135     ;; Set CPU 0.
 136     (setf (cpu-affinity-p 0 mask) t))
 137
 138   (with-cpu-affinity-mask (mask)
 139     (print mask))
 140
 141   (with-cpu-affinity-mask (mask :save t)
 142     ;; Only odd CPUs in mask.
 143     (dotimes (cpu (cpu-count))
 144       (setf (cpu-affinity-p cpu mask) (oddp cpu))))
 145
 146   (with-cpu-affinity-mask (mask)
 147     (print mask)))