add utility for memory efficient maximum pairwise distance computation with GPU support #838
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Compute maximal pairwise distance without storing all distances Provides fallback to CPU implementation of optional cuVS/pylibraft dependency is not available
grlee77
added
improvement
gigony
reviewed
Feb 19, 2025
| if _distance_on_cpu: | ||
| warnings.warn( | ||
| "cuVS >= 25.02 or pylibraft < 24.12 must be installed to use " | ||
| "GPU-accelerated pairwaise distance computations. Falling back " |
There was a problem hiding this comment.
Suggested change
| "GPU-accelerated pairwaise distance computations. Falling back " | |
| "GPU-accelerated pairwise distance computations. Falling back " |
| Internally, calls to cdist will be made with subsets of coords where | ||
| the subset size is (coords_per_block, ndim). | ||
| compute_argmax : bool, optional | ||
| If True, the value of the coordate indices corresponding to the maxima |
There was a problem hiding this comment.
Suggested change
| If True, the value of the coordate indices corresponding to the maxima | |
| If True, the value of the cooridate indices corresponding to the maxima |
| requirement. The memory used at runtime will be proportional to | ||
| ``coords_per_block**2``. | ||
|
|
||
| A block size of >= 2000 is recommended to overhead poor GPU resource usage |
There was a problem hiding this comment.
Suggested change
| A block size of >= 2000 is recommended to overhead poor GPU resource usage | |
| A block size of >= 2000 is recommended to avoid poor GPU resource usage |
| ) | ||
| current_output = temp | ||
| else: | ||
| # omit out= for the last block as size may be |
There was a problem hiding this comment.
This comment sentence doesn't seem to be complete.
Comment on lines
+48
to
+49
| coords : np.ndarray (num_points, ndim) | ||
| The coordinates to process. |
There was a problem hiding this comment.
The code converts input coordinates to float32 by default, which could lead to precision loss. I think it would be good to note that this coords would be converted to float32 data type.
Comment on lines
+108
to
+109
|
|
||
| num_coords, _ = coords.shape |
There was a problem hiding this comment.
Input validation might be needed here.
Suggested change
| num_coords, _ = coords.shape | |
| if not isinstance(coords, (np.ndarray, cp.ndarray)): | |
| raise TypeError("coords must be a numpy or cupy array") | |
| if coords.ndim != 2: | |
| raise ValueError( | |
| f"coords must be a 2-dimensional array, got shape {coords.shape}" | |
| ) | |
| num_coords, _ = coords.shape |
| "to SciPy-based CPU implementation." | ||
| ) | ||
| xp = np | ||
| coords = cp.asnumpy(coords) |
There was a problem hiding this comment.
We don't need to use numpy for here?
Suggested change
| coords = cp.asnumpy(coords) | |
| coords = np.asnumpy(coords) |
|
|
||
| Parameters | ||
| ---------- | ||
| coords : np.ndarray (num_points, ndim) |
There was a problem hiding this comment.
Suggested change
| coords : np.ndarray (num_points, ndim) | |
| coords : numpy.ndarray or cupy.ndarray of shape (num_points, ndim) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The maximum Feret diameter computations compute pairwise distances between all points which can lead to out of memory errors if the number of points on the object boundary is large. (The memory used is quadratic in the number of points).
This MR implements a block-wise version that retains efficiency, but has much lower memory requirements. It also handles checking for optional GPU acceleration via CuPy's optional cuVS/pylibraft dependencies. CPU fallback is done if those are not available.