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 forget-cpu-count ()
  43   (setf *cpu-count* nil))
  44
  45 (pushnew 'forget-cpu-count sb-ext:*save-hooks*)
  46
  47 (defun cpu-count ()
  48   "Number of CPUs available in the system. Based on /proc/cpuinfo."
  49   (or *cpu-count*
  50       (setf *cpu-count*
  51             (let* ((key "processor")
  52                    (len (length key)))
  53               (with-open-file (f "/proc/cpuinfo")
  54                 (loop for line = (read-line f nil nil)
  55                       while line
  56                       count (when (> (length line) len)
  57                               (string= key line :end2 len))))))))
  58
  59 (defstruct cpu-affinity-mask
  60   "CPU affinity mask."
  61   %mask)
  62
  63 (defmethod print-object ((mask cpu-affinity-mask) stream)
  64   (print-unreadable-object (mask stream :type t)
  65     ;; Print the locally interesting part of the mask.
  66     (dotimes (i (cpu-count))
  67       (if (cpu-affinity-p i mask)
  68           (write-char #\1 stream)
  69           (write-char #\0 stream)))))
  70
  71 (defun make-cpu-mask ()
  72   (make-alien (unsigned 8) cpu-mask-size))
  73
  74 (defun get-cpu-affinity-mask ()
  75   "Returns the CPU affinity mask of the current thread. The mask can
  76 be inspected and mutated using CPU-AFFINITY-P, \(SETF CPU-AFFINITY-P),
  77 and CLEAR-CPU-AFFINITY-MASK. To make any changes take effect, the mask
  78 must be saved using SET-CPU-AFFINITY-MASK.
  79
  80 Using WITH-CPU-AFFINITY-MASK instead is recommended."
  81   ;; FIXME: Malloc'ed mask is nasty, but libc doesn't seem to like
  82   ;; stack allocated ones, nor does it seem to like freeing masks that
  83   ;; have been used. So we never do. Gah.
  84   (let ((mask (make-cpu-mask)))
  85     (unless (zerop (%get-cpu-affinity-mask mask))
  86       (error "Could not read CPU affinity mask: ~A" (sb-int:strerror)))
  87     (make-cpu-affinity-mask :%mask mask)))
  88
  89 (defun set-cpu-affinity-mask (mask)
  90   "Sets the CPU affinity mask for the current thread.
  91
  92 Using WITH-CPU-AFFINITY-MASK :SAVE T instead is recommended."
  93   (unless (zerop (%set-cpu-affinity-mask (cpu-affinity-mask-%mask mask)))
  94     (error "Coud not write CPU affinity mask: ~A" (sb-int:strerror))))
  95
  96 (defun cpu-affinity-p (cpu mask)
  97   "Returns T if the CPU \(a numeric indentifier between 0 and
  98  +CPU-LIMIT+) is part of the MASK."
  99   (plusp (%cpu-affinity-p cpu (cpu-affinity-mask-%mask mask))))
 100
 101 (defun (setf cpu-affinity-p) (bool cpu mask)
 102   "Toggles presence of the CPU \(a numeric identifier between 0 and +CPU-LIMIT+)
 103 in the MASK."
 104   (let ((%mask (cpu-affinity-mask-%mask mask)))
 105     (if bool
 106         (set-cpu-affinity cpu %mask)
 107         (clear-cpu-affinity cpu %mask)))
 108   bool)
 109
 110 (defun clear-cpu-affinity-mask (mask)
 111   "Removes all CPUs from the MASK."
 112   (zero-cpu-affinity-mask (cpu-affinity-mask-%mask mask))
 113   mask)
 114
 115 (defmacro with-cpu-affinity-mask ((mask &key save) &body body)
 116   "Reads the CPU affinity mask of the the current thread and binds it
 117 to MASK. The mask can be inspected and mutated using CPU-AFFINITY-P,
 118 \(SETF CPU-AFFINITY-P, and CLEAR-CPU-AFFINITY-MASK. Any changes take
 119 effect only if SAVE is true (default is NIL)."
 120   (let ((ok-n (gensym "OK")))
 121     `(let (,mask ,ok-n)
 122        (unwind-protect
 123             (progn
 124               (setf ,mask (get-cpu-affinity-mask))
 125               (multiple-value-prog1 (progn ,@body)
 126                 (setf ,ok-n t)))
 127          (when (and ,ok-n ,save)
 128            (set-cpu-affinity-mask ,mask))))))
 129
 130 ;;;; Usage examples
 131 #+nil
 132 (progn
 133
 134   (with-cpu-affinity-mask (mask)
 135     (print mask))
 136
 137   (with-cpu-affinity-mask (mask :save t)
 138     ;; Remove all
 139     (clear-cpu-affinity-mask mask)
 140     ;; Set CPU 0.
 141     (setf (cpu-affinity-p 0 mask) t))
 142
 143   (with-cpu-affinity-mask (mask)
 144     (print mask))
 145
 146   (with-cpu-affinity-mask (mask :save t)
 147     ;; Only odd CPUs in mask.
 148     (dotimes (cpu (cpu-count))
 149       (setf (cpu-affinity-p cpu mask) (oddp cpu))))
 150
 151   (with-cpu-affinity-mask (mask)
 152     (print mask)))