vol/fftw-fft.lisp

   1 (declaim (optimize (speed 2) (debug 3) (safety 3)))
   2 (in-package :vol)
   3
   4 (defconstant +forward+ 1)
   5 (defconstant +backward+ -1)
   6 (defconstant +estimate+ (ash 1 6))
   7
   8 (defmacro def-ffi (&optional (library "libfftw3.so") (prefix "fftw"))
   9   `(progn
  10      (load-shared-object ,(format nil "~a" library))
  11      (define-alien-routine (,(format nil "~a_execute" prefix)
  12                              ,(vol:format-symbol "~a-execute" prefix))
  13          void
  14        (plan (* int)))
  15      ,@(loop for rank in '(1 2 3) collect
  16             `(define-alien-routine
  17                  (,(format nil "~a_plan_dft_~ad" prefix rank)
  18                    ,(format-symbol "~a-plan-dft-~ad" prefix rank))
  19                  (* int)
  20                ,@(loop for i below rank collect
  21                       `(,(format-symbol "n~d" i) int))
  22                (in (* double-float))
  23                (out (* double-float))
  24                (sign int)
  25                (flags unsigned-int)))))
  26
  27 ;; define the foreign functions
  28 ;; fftw{f}-execute, fftw{f}-plan-dft-{1,2,3}d
  29 (def-ffi "libfftw3.so.3" "fftw")
  30 (def-ffi "libfftw3f.so.3" "fftwf")
  31
  32 #+nil
  33 (progn
  34   (load-shared-object "/usr/lib/libfftw3_threads.so")
  35
  36   (define-alien-routine ("fftw_init_threads" init-threads)
  37       int)
  38   (define-alien-routine ("fftw_plan_with_nthreads" plan-with-nthreads)
  39       void
  40     (nthreads int))
  41
  42   (defun init-ft ()
  43     (init-threads)
  44     (plan-with-nthreads 4)))
  45
  46 ;; to clean up completely call void fftw_cleanup_threads(void)
  47
  48
  49
  50 (def-generator (ft (rank type))
  51   (let* ((prefix (ecase type
  52                   (cdf 'fftw)
  53                   (csf 'fftwf)))
  54          (rlong-type (ecase type
  55                        (csf (get-long-type 'sf))
  56                        (cdf (get-long-type 'df))))
  57          (plan (format-symbol "~a-plan-dft-~ad" prefix rank))
  58          (execute (format-symbol "~a-execute" prefix)))
  59     `(defun ,name (in &key (forward t))
  60        (declare ((simple-array ,long-type ,rank) in)
  61                 (boolean forward)
  62                (values (simple-array ,long-type ,rank) &optional))
  63       (let* ((dims (array-dimensions in))
  64              (out (make-array dims :element-type ',long-type))
  65              (in-sap (sb-sys:vector-sap
  66                       (sb-ext:array-storage-vector in)))
  67              (out-sap (sb-sys:vector-sap
  68                        (sb-ext:array-storage-vector out)))
  69              (dir (if forward +forward+ +backward+)))
  70         (sb-sys:with-pinned-objects (in out)
  71           ,(ecase rank
  72                   (1 `(destructuring-bind (x) dims
  73                         (let ((p (,plan x in-sap out-sap
  74                                         dir +estimate+)))
  75                           (,execute p))))
  76                   (2 `(destructuring-bind (y x) dims
  77                         (let ((p (,plan y x in-sap out-sap
  78                                         dir +estimate+)))
  79                           (,execute p))))
  80                   (3 `(destructuring-bind (z y x) dims
  81                         (let ((p (,plan z y x in-sap out-sap
  82                                         dir +estimate+)))
  83                           (,execute p))))))
  84         ;; normalize so that a=ift(ft(a)), after forward transform
  85         ;; divide by n
  86         (if forward
  87             (s* (/ ,(coerce 1 rlong-type)
  88                    (array-total-size out))
  89                 out)
  90             out)))))
  91
  92 #+nil
  93 (def-ft-rank-type 1 csf)
  94
  95 (defmacro def-ft-functions (ranks types)
  96   (let* ((specifics nil)
  97          (cases nil)
  98          (name (format-symbol "ft")))
  99     (loop for rank in ranks do
 100          (loop for type in types do
 101               (let ((def-name (format-symbol "def-~a-rank-type" name))
 102                     (specific-name (format-symbol "~a-~a-~a" name rank type)))
 103                 (push `(,def-name ,rank ,type) specifics)
 104                 (push `((simple-array ,(get-long-type type) ,rank)
 105                         (,specific-name a :forward forward))
 106                       cases))))
 107     (store-new-function name)
 108     `(progn ,@specifics
 109             (defun ,name (a &key (forward t))
 110                (etypecase a
 111                  ,@cases
 112                  (t (error "The given type can't be handled with a generic ~a function." ',name)))))))
 113
 114 (def-ft-functions (1 2 3) (csf cdf))
 115
 116 (defmacro ift (in)
 117   (alexandria:with-gensyms (input)
 118     (let ((input in))
 119       `(ft ,input :forward nil))))
 120
 121 #+nil
 122 (progn
 123   (time
 124    (let* ((nx 128)
 125           (ny nx)
 126           (nz ny)
 127           (a (convert-3-ub8/csf-mul
 128               (draw-sphere-ub8 20.0 nz ny nx))))
 129      (write-pgm "/home/martin/tmp/fftw.pgm"
 130                 (normalize-2-csf/ub8-abs
 131                  (cross-section-xz-csf (ft a)))))))