xptycho blueprints › Hardware Mapping
The design for carrying out the algorithm blueprint's iteration on \(N_g\) GPUs of a single node, with all the data held in GPU memory. One GPU is the case \(N_g = 1\) of the same design. It is for the person writing or reviewing recon. It answers five questions:
ptycho_pmace, and with itself on different devices?\(J\) is the number of scan positions, \(K\) the number of probe modes, \(N_p \times N_p\) the frame size, \(N\) the number of object pixels.
| array | what it is | bytes | placement |
|---|---|---|---|
| \(y_j\) | measured amplitudes (float32) | \(4 N_p^2\) per position | sharded by position. Each GPU holds one block of positions. These arrays never move. |
| \(v_j\) | the position's copy of its object patch | \(8 N_p^2\) per position | |
| \(s_{j,k}\) | the position's copy of each probe mode; only when the probe is estimated | \(8 K N_p^2\) per position | |
| \(\bar x\), the image sums, \(\Lambda_k\) | the object image, the sums over positions that build it, the coverage maps | a few times \(8N\) in total | sharded by image row. Each GPU owns one band of rows, and also holds its window. |
| \(d_k\), \(u_k\), the deviation | the probe modes and the other probe-sized arrays | small | broadcast. |
How many positions fit on one 80 GB GPU, leaving 10 GB for its band, its window, and temporary arrays:
| case | bytes per position | positions per GPU, \(N_p = 256\) | positions per GPU, \(N_p = 512\) |
|---|---|---|---|
| known probe | \(12 N_p^2\) | 89,000 | 22,000 |
| estimated probe, 1 mode | \(20 N_p^2\) | 53,000 | 13,000 |
| estimated probe, 2 modes | \(28 N_p^2\) | 38,000 | 9,500 |
With \(N_g\) GPUs the capacity is \(N_g\) times these. The gold balls (794 positions, \(N_p = 512\), 2 modes) need 6 GB and fit on one GPU.
Each GPU computes over its own block and adds into a partial sum. Two operations then move data between GPUs. They are mbirtorch's, and every step uses them.
In the example above, GPU 0 sends its partial for rows 19–23 to GPU 1, and afterward copies the result for rows 19–23 back. A probe-sized array is the same with one owner, GPU 0, and a window that is the whole array.
| step | each GPU, over its own positions | summed to owner | the owner computes, and windows copy |
|---|---|---|---|
| 0. start (once) | \(v_j = P_j x^{(0)}\), cut from its window; \(s_{j,k} = d^{(0)}_k\); add \(|d_k|^\kappa\) at each position into an image sum per mode | \(K\) image sums | \(\Lambda_k\). The starting probe \(d^{(0)}\) and object \(x^{(0)}\) of the algorithm blueprint, section 4, are computed on the host before the run and sent once (section 3). |
| 1. fit patches | \(w_j = F^I_j(v_j)\); add \(|d_k|^\kappa(2w_j - v_j)\) into an image sum per mode; \(v_j \leftarrow v_j - 2\rho w_j\) | \(K\) image sums | the average image of \(2\mathbf{w}-\mathbf{v}\) |
| 2. update patches | \(v_j \leftarrow v_j + 2\rho z_j\), \(z_j\) the patch of that image; add \(|d_k|^\kappa v_j\) into an image sum per mode | \(K\) image sums | \(\bar x\), the reported image |
| 3. add a mode (scheduled iterations) | the residual, back-transformed and divided by the patch of \(\bar x\); add into a sum | one probe-sized sum | the new mode and the rescale factor; each GPU then makes its own copies and rescales its own arrays |
| 4. fit probe copies | \(r_{j,k} = F^P_{j,k}(s_{j,k})\), all modes, using the modes as they stand before this step; add \(2r_{j,k} - s_{j,k}\) into a sum per mode; \(s_{j,k} \leftarrow s_{j,k} - 2\rho r_{j,k}\) | \(K\) probe-sized sums | \(u_k\) |
| 5. update probe copies | \(s_{j,k} \leftarrow s_{j,k} + 2\rho u_k\); add \(s_{j,k}\) into a sum per mode. After the new modes return, add \(|d_k|^\kappa\) at each position into an image sum per mode | \(K\) probe-sized sums; then \(K\) image sums | the new modes \(d_k\); then the new \(\Lambda_k\) |
| 6. replace outliers | add \(|s_{j,k} - d_k|\) into a sum per mode; after the result returns, set outlying pixels of \(s_{j,k}\) to \(d_k\) | \(K\) probe-sized sums | the deviation per pixel |
| 7. data error (for the progress line) | forward-project the patches of \(\bar x\) with \(d_k\); add the squared difference from \(y_j\), and \(y_j^2\) | two numbers | the data error |
The algorithm blueprint's update \(v + 2\rho(z - w)\) is done in place in two parts, \(-2\rho w\) in step 1 and \(+2\rho z\) in step 2, so \(w\) is never stored; likewise for the probe copies. Steps 3 to 6 run only when the probe is estimated. With a known probe, \(\Lambda_k\) is computed once in step 0. Within a GPU, a step processes its block in batches (about 1,000 positions at \(N_p = 256\) with 9 GB free) so that the temporary arrays of the Fourier transforms fit.
Numbers that depend on all the positions or on the whole image. Each is formed from sums over all the GPUs. Computing one per GPU or per batch gives a different result (an error near \(10^{-3}\) in a small test).
| number | used in | how it is formed |
|---|---|---|
| \(\varepsilon\) of the division by \(\Lambda_k\) | steps 0, 1, 2 | each owner adds \(|\Lambda_k|^2\) over its band; the sums are added and divided by \(N\), so every pixel of the grid is counted once |
| \(\varepsilon\) of the division by the patch | steps 3, 4 | each GPU adds \(\|p_j\|^2\) over its positions, \(p_j\) the patch of \(\bar x\); the sums are added and divided by \(J N_p^2\); one number per iteration |
| every mean over positions (a new mode, \(u_k\), \(d_k\), the deviation) | steps 3, 4, 5, 6 | the GPUs' sums are added and divided by \(J\); never an average of the GPUs' own means |
Why the result matches one GPU. Each position is on exactly one GPU and each image row has exactly one owner, so every sum receives every contribution once. The owner adds the pieces in GPU order, so every copy of a value is identical on all GPUs. What differs from a one-GPU run is the order of the additions inside the sums. How far that moves the result is measured, not assumed; the measurements are in section 5.
| when | direction | what | gold balls |
|---|---|---|---|
| once, at the start | host to GPU | the amplitudes of each GPU's block, and the patch positions of that block; the starting probe, to every GPU; the starting object, each band to its owner | 0.8 GB |
| each iteration | GPU to host | a few numbers for the progress line | bytes |
| each iteration of the mode schedule | GPU to host, and back | the mean of the back-transformed residuals, which the host Fresnel propagates with numpy and returns as the new mode; with orthogonalize_modes, the modes, for the singular value decomposition on the host | a few MB |
| once, at the end | GPU to host | the bands of \(\bar x\) and of the coverage, joined into one image; \(d_k\) | about 10 MB |
The starting probe and object, when not given, are computed on the host CPU from the frames before anything goes to the GPUs. In the iteration, the GPUs wait on the host only for the data error of the progress line and, at an iteration of the mode schedule, for the new mode and the orthogonalization.
recon refuses before it starts.The GPUs are chosen with one method of the model, configure_devices, taken from mbirtorch so that the two packages are used the same way.
| call | what it does |
|---|---|
| no call | The model chooses automatically: every CUDA GPU on the node; if there is none, Apple's GPU (mps); if there is none, the CPU. The choice is printed. |
model.configure_devices(num_devices=4) | Use 4 CUDA GPUs. \(N_g = 4\). |
model.configure_devices(num_devices=1) | Use one device. |
model.configure_devices(devices=['cuda:0', 'cuda:2']) | Use exactly these devices. devices=['cpu'] runs on the CPU; a list of several CPU devices is how the tests exercise \(N_g > 1\) on a machine with no GPU. |
recon. Once made, the automatic choice is off for that model.ptycho_pmace, and with itself on different devices.Measured with dev_scripts/baseline/compare_baseline.py and the test suite. Differences are relative, on the object image, after removing one complex scale.
| case | xptycho against ptycho_pmace | xptycho, one device against several |
|---|---|---|
| tiny problem, no noise, known probe, 20 iterations | \(6 \times 10^{-7}\) | \(6 \times 10^{-7}\) (2, 3, 5 devices) |
| tiny problem, no noise, two modes estimated, 6 iterations | within \(10^{-3}\) (test tolerance) | within \(10^{-3}\) (3 devices) |
| gold balls, known probe, 100 iterations | \(2 \times 10^{-5}\); data error equal to 6 digits | |
| synthetic with noise, known probe, 100 iterations | \(3 \times 10^{-3}\); data error equal to 6 digits; error to the truth 0.036925 against 0.036927 | \(1 \times 10^{-3}\) (3 devices); \(3 \times 10^{-3}\) (CPU against Apple GPU) |
| blind, two modes, the second added at iteration 20, 200 iterations | about \(10^{-2}\); error to the truth 0.044376 against 0.044419; data error 0.292723 against 0.292830 | |
| blind, one mode, 200 iterations | data error equal to 7 digits for 3 iterations, then to about \(10^{-3}\); error to the truth 0.0494 against 0.0497 |
On noisy data the object image is sensitive to rounding: two runs that differ only in the order of additions drift apart by about \(10^{-3}\) over 100 iterations, while the data error and the error to the truth stay fixed to five digits. When the probe is estimated, the outlier test is a threshold, so a last-bit difference can change which probe copies are replaced. xptycho differs from ptycho_pmace by no more than it differs from itself on another device.