cpu-affinity.lisp

   1 ;; Copyright (c) 2008 Nikodemus Siivola <nikodemus@random-state.net>
   2 ;;
   3 ;; Permission is hereby granted, free of charge, to any person
   4 ;; obtaining a copy of this software and associated documentation
   5 ;; files (the "Software"), to deal in the Software without
   6 ;; restriction, including without limitation the rights to use,
   7 ;; copy, modify, merge, publish, distribute, sublicense, and/or sell
   8 ;; copies of the Software, and to permit persons to whom the
   9 ;; Software is furnished to do so, subject to the following
  10 ;; conditions:
  11 ;;
  12 ;; The above copyright notice and this permission notice shall be
  13 ;; included in all copies or substantial portions of the Software.
  14 ;;
  15 ;; THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  16 ;; EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  17 ;; OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  18 ;; NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  19 ;; HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  20 ;; WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  21 ;; FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  22 ;; OTHER DEALINGS IN THE SOFTWARE.
  23
  24 (in-package :sb-cpu-affinity)
  25
  26 ;;;; Alien definitions
  27
  28 ;; This is defined as a #DEFINE constant, but the .so stuff the value
  29 ;; into a variable for us.
  30 (eval-when (:compile-toplevel :load-toplevel)
  31   (define-alien-variable cpu-setsize int))
  32
  33 ;; This is sizeof(cpu_set_t)
  34 (define-alien-variable cpu-mask-size int)
  35
  36 ;; These are convenience wrappers around sched_set/getaffinity.
  37 (define-alien-routine ("get_cpu_affinity_mask" %get-cpu-affinity-mask) int
  38   (mask (* (unsigned 8))))
  39 (define-alien-routine ("set_cpu_affinity_mask" %set-cpu-affinity-mask) int
  40   (mask (* (unsigned 8))))
  41
  42 ;; These are wrappers around libc macros that manipulate the mask. Why
  43 ;; the hell can't C people provide functions in addition to macros
  44 ;; like this? Or just use inline functions?
  45 (define-alien-routine zero-cpu-affinity-mask void
  46   (mask (* (unsigned 8))))
  47 (define-alien-routine ("cpu_affinity_p" %cpu-affinity-p) int
  48   (cpu int)
  49   (mask (* (unsigned 8))))
  50 (define-alien-routine set-cpu-affinity  void
  51   (cpu int)
  52   (mask (* (unsigned 8))))
  53 (define-alien-routine clear-cpu-affinity void
  54   (cpu int)
  55   (mask (* (unsigned 8))))
  56
  57 ;;;; Nice lispy interface
  58
  59 (defconstant +cpu-limit+ cpu-setsize
  60   "Upper exclusive limit on the number of CPUs. Based on CPU_SETSIZE
  61 from sched.h.")
  62
  63 (defvar *cpu-count* nil)
  64
  65 (defun forget-cpu-count ()
  66   (setf *cpu-count* nil))
  67
  68 (pushnew 'forget-cpu-count sb-ext:*save-hooks*)
  69
  70 (defun cpu-count ()
  71   "Number of CPUs available in the system. Based on /proc/cpuinfo."
  72   (or *cpu-count*
  73       (setf *cpu-count*
  74             (let* ((key "processor")
  75                    (len (length key)))
  76               (with-open-file (f "/proc/cpuinfo")
  77                 (loop for line = (read-line f nil nil)
  78                       while line
  79                       count (when (> (length line) len)
  80                               (string= key line :end2 len))))))))
  81
  82 (defstruct cpu-affinity-mask
  83   "CPU affinity mask."
  84   %mask)
  85
  86 (defmethod print-object ((mask cpu-affinity-mask) stream)
  87   (print-unreadable-object (mask stream :type t)
  88     ;; Print the locally interesting part of the mask.
  89     (dotimes (i (cpu-count))
  90       (if (cpu-affinity-p i mask)
  91           (write-char #\1 stream)
  92           (write-char #\0 stream)))))
  93
  94 (defun make-cpu-mask ()
  95   (make-alien (unsigned 8) cpu-mask-size))
  96
  97 (defun get-cpu-affinity-mask ()
  98   "Returns the CPU affinity mask of the current thread. The mask can
  99 be inspected and mutated using CPU-AFFINITY-P, \(SETF CPU-AFFINITY-P),
 100 and CLEAR-CPU-AFFINITY-MASK. To make any changes take effect, the mask
 101 must be saved using SET-CPU-AFFINITY-MASK.
 102
 103 Using WITH-CPU-AFFINITY-MASK instead is recommended."
 104   ;; FIXME: Malloc'ed mask is nasty, but libc doesn't seem to like
 105   ;; stack allocated ones, nor does it seem to like freeing masks that
 106   ;; have been used. So we never do. Gah.
 107   (let ((mask (make-cpu-mask)))
 108     (unless (zerop (%get-cpu-affinity-mask mask))
 109       (error "Could not read CPU affinity mask: ~A" (sb-int:strerror)))
 110     (make-cpu-affinity-mask :%mask mask)))
 111
 112 (defun set-cpu-affinity-mask (mask)
 113   "Sets the CPU affinity mask for the current thread.
 114
 115 Using WITH-CPU-AFFINITY-MASK :SAVE T instead is recommended."
 116   (unless (zerop (%set-cpu-affinity-mask (cpu-affinity-mask-%mask mask)))
 117     (error "Coud not write CPU affinity mask: ~A" (sb-int:strerror))))
 118
 119 (defun cpu-affinity-p (cpu mask)
 120   "Returns T if the CPU \(a numeric indentifier between 0 and
 121  +CPU-LIMIT+) is part of the MASK."
 122   (plusp (%cpu-affinity-p cpu (cpu-affinity-mask-%mask mask))))
 123
 124 (defun (setf cpu-affinity-p) (bool cpu mask)
 125   "Toggles presence of the CPU \(a numeric identifier between 0 and +CPU-LIMIT+)
 126 in the MASK."
 127   (let ((%mask (cpu-affinity-mask-%mask mask)))
 128     (if bool
 129         (set-cpu-affinity cpu %mask)
 130         (clear-cpu-affinity cpu %mask)))
 131   bool)
 132
 133 (defun clear-cpu-affinity-mask (mask)
 134   "Removes all CPUs from the MASK."
 135   (zero-cpu-affinity-mask (cpu-affinity-mask-%mask mask))
 136   mask)
 137
 138 (defmacro with-cpu-affinity-mask ((mask &key save) &body body)
 139   "Reads the CPU affinity mask of the the current thread and binds it
 140 to MASK. The mask can be inspected and mutated using CPU-AFFINITY-P,
 141 \(SETF CPU-AFFINITY-P, and CLEAR-CPU-AFFINITY-MASK. Any changes take
 142 effect only if SAVE is true (default is NIL)."
 143   (let ((ok-n (gensym "OK")))
 144     `(let (,mask ,ok-n)
 145        (unwind-protect
 146             (progn
 147               (setf ,mask (get-cpu-affinity-mask))
 148               (multiple-value-prog1 (progn ,@body)
 149                 (setf ,ok-n t)))
 150          (when (and ,ok-n ,save)
 151            (set-cpu-affinity-mask ,mask))))))
 152
 153 ;;;; Usage examples
 154 #+nil
 155 (progn
 156
 157   (with-cpu-affinity-mask (mask)
 158     (print mask))
 159
 160   (with-cpu-affinity-mask (mask :save t)
 161     ;; Remove all
 162     (clear-cpu-affinity-mask mask)
 163     ;; Set CPU 0.
 164     (setf (cpu-affinity-p 0 mask) t))
 165
 166   (with-cpu-affinity-mask (mask)
 167     (print mask))
 168
 169   (with-cpu-affinity-mask (mask :save t)
 170     ;; Only odd CPUs in mask.
 171     (dotimes (cpu (cpu-count))
 172       (setf (cpu-affinity-p cpu mask) (oddp cpu))))
 173
 174   (with-cpu-affinity-mask (mask)
 175     (print mask)))