AdvancedPhotonSource / tike

Repository for ptychography software
http://tike.readthedocs.io
Other
29 stars 15 forks source link

REF: Load batches into contiguous memory blocks #287

Closed carterbox closed 12 months ago

carterbox commented 1 year ago

Purpose

Load batched variables (scan positions, diffraction patterns, eigen probe weights) into contiguous blocks of memory according to their minibatch instead of reading them from the memory in the order that they were collected. This will enable more efficient data movement on and to/from the GPU because only contiguous blocks of pinned memory can be transferred to the GPU asynchronously.

Approach

Create a new splitting method to determine which positions are assigned to which GPU that just sorts the positions along one axis. Then immediately apply the user requested batch method to the divided positions instead of loading them to the GPU first. Once we have figured out all the batches and converted them into indices in the originally-ordered dataset, we can load the batches into contiguous blocks of GPU memory. The batch indices, then become ranges of consecutive integers.

I have also refactored the clustering methods so that they behave uniformly. i.e. When the number clusters is larger than the population being divided, some empty batches are now returned instead of an error being raised which complains that there are too many batches. Also, batches are expected to be ordered from largest batch to smallest batch when the number of batches does not divide the population evenly.

Pre-Merge Checklists

Submitter

Reviewer