tebdol

Simulation of ultracold atoms in optical lattices
git clone https://miroslavurbanek.com/git/tebdol.git
Log | Files | Refs | README

part.lisp (4367B)


      1 (defpackage :part
      2   (:use :common-lisp :mpi :mps)
      3   (:export :uniform-partition
      4 	   :balance-partition
      5 	   :change-partition))
      6 
      7 (in-package :part)
      8 
      9 (defun uniform-partition (size length)
     10   (multiple-value-bind (q r) (truncate length size)
     11     (loop
     12        with d = (max 2 q)
     13        with e = (if (< q 2) 0 r)
     14        for i below size
     15        for j from d by d
     16        when (or (< (* 2 i) e)
     17 		(<= (* 2 (- size i)) e))
     18        do (incf j)
     19        if (< j length)
     20        collect j
     21        else
     22        collect length)))
     23 
     24 (defun balanced-partition (size loads)
     25   (loop
     26      with l = (mapcar #'rationalize loads)
     27      with a = (/ (apply #'+ l) size)
     28      with i = 0
     29      for c in l
     30      count c into j
     31      sum c into s
     32      when (>= s a)
     33      if (and (> (- s a) (/ c 2))
     34 	     (> j (+ i 2)))
     35      collect (1- j) into p
     36      and
     37      do
     38        (decf s a)
     39        (setf i (1- j))
     40      else
     41      when (> j (1+ i))
     42      collect j into p
     43      and
     44      do
     45        (decf s a)
     46        (setf i j)
     47      finally
     48        (return
     49 	 (let ((q (butlast p)))
     50 	   (nconc
     51 	    q
     52 	    (make-list (- size (length q)) :initial-element (length l)))))))
     53 
     54 (defun send-subloads (subloads)
     55   (mpi-send-object subloads 0))
     56 
     57 (defun collect-subloads (size subloads)
     58   (loop
     59      with a = (make-array size)
     60      repeat (1- size)
     61      do
     62        (multiple-value-bind (s i) (mpi-receive-object +mpi-any-source+)
     63 	 (setf (svref a i) s))
     64      finally
     65        (setf (svref a 0) subloads)
     66        (return
     67 	 (loop
     68 	    for s across a
     69 	    nconc s))))
     70 
     71 (defun balance-partition (size rank subloads)
     72   (if (zerop rank)
     73       (let* ((l (collect-subloads size subloads))
     74 	     (p (balanced-partition size l)))
     75 	p)
     76       (send-subloads subloads)))
     77 
     78 (defun instructions (length oldpartition newpartition)
     79   (loop
     80      with op = 0
     81      with np = 0
     82      for i below length
     83      when (eql i (car oldpartition))
     84      do
     85        (incf op)
     86        (setf oldpartition (cdr oldpartition))
     87      when (eql i (car newpartition))
     88      do
     89        (incf np)
     90        (setf newpartition (cdr newpartition))
     91      unless (= op np)
     92      collect (list i op np)))
     93 
     94 (defun sort-instructions (instructions)
     95   (loop
     96      for i in instructions
     97      when (evenp (cadr i))
     98      collect i into e
     99      else
    100      collect i into o
    101      finally (return (nconc e o))))
    102 
    103 (defun send-instructions (size instructions)
    104   (let ((a (make-array size :initial-element nil)))
    105     (dolist (i instructions)
    106       (push i (svref a (cadr i)))
    107       (push i (svref a (caddr i))))
    108     (sb-sys:without-gcing
    109       (loop
    110 	 for i from 1 below size
    111 	 collect (mpi-issend-object (svref a i) i)
    112 	 into l
    113 	 finally (mapc #'mpi-wait l)))
    114     (svref a 0)))
    115 
    116 (defun receive-instructions ()
    117   (mpi-receive-object 0))
    118 
    119 (defun process-instructions (rank length submps instructions)
    120   (if instructions
    121       (let ((l (loop
    122 		  repeat (submps-size submps)
    123 		  for i from (submps-start submps)
    124 		  collect i))
    125 	    s r)
    126 	(loop
    127 	   for i in instructions
    128 	   if (= (cadr i) rank)
    129 	   do
    130 	     (setf l (delete (car i) l))
    131 	     (push i s)
    132 	   else
    133 	   do
    134 	     (push (car i) l)
    135 	     (push i r))
    136 	(let* ((os (submps-start submps))
    137 	       (ou (submps-tensors submps))
    138 	       (ov (submps-singvals submps))
    139 	       (ns (reduce #'min l :initial-value length))
    140 	       (nd (length l))
    141 	       (nu (make-array (1+ nd) :initial-element nil))
    142 	       (nv (make-array (1+ nd) :initial-element nil))
    143 	       m)
    144 	  (sb-sys:without-gcing
    145 	    (dolist (i (nreverse s))
    146 	      (let ((j (- (car i) os))
    147 		    (p (caddr i)))
    148 		(push (mpi-issend-object (svref ou j) p) m)
    149 		(push (mpi-issend-object (svref ov j) p) m)))
    150 	    (dolist (i (nreverse r))
    151 	      (let ((j (- (car i) ns))
    152 		    (p (cadr i)))
    153 		(setf (svref nu j) (mpi-receive-object p))
    154 		(setf (svref nv j) (mpi-receive-object p))))
    155 	    (mapc #'mpi-wait m))
    156 	  (dotimes (i nd)
    157 	    (unless (svref nu i)
    158 	      (let ((j (+ i (- ns os))))
    159 		(setf (svref nu i) (svref ou j))
    160 		(setf (svref nv i) (svref ov j)))))
    161 	  (make-submps :start ns :size nd :tensors nu :singvals nv)))
    162       submps))
    163 
    164 (defun change-partition (size rank length submps oldpartition newpartition)
    165   (if (zerop rank)
    166       (let* ((i (instructions length oldpartition newpartition))
    167 	     (s (sort-instructions i))
    168 	     (l (send-instructions size s)))
    169 	(process-instructions rank length submps l))
    170       (let ((l (receive-instructions)))
    171 	(process-instructions rank length submps l))))