mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-14 19:27:20 +00:00
7c390fffd6
Add a public Granularity enum (MultiCastMinimum, MultiCastRecommended) and let GpuBuffer choose the NVLS multicast allocation granularity via a constructor argument, defaulting to MultiCastMinimum to minimize memory usage. Expose the same option through the C++ and Python (nanobind) APIs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# Licensed under the MIT License.
|
|
|
|
|
|
from typing import Union, Tuple
|
|
|
|
import cupy as cp
|
|
import numpy as np
|
|
from mscclpp._mscclpp import CppRawGpuBuffer, CppGpuBufferGranularity
|
|
|
|
__all__ = ["GpuBuffer", "GpuBufferGranularity"]
|
|
|
|
GpuBufferGranularity = CppGpuBufferGranularity
|
|
|
|
|
|
class GpuBuffer(cp.ndarray):
|
|
def __new__(
|
|
cls,
|
|
shape: Union[int, Tuple[int]],
|
|
dtype: cp.dtype = float,
|
|
strides: Tuple[int] = None,
|
|
order: str = "C",
|
|
granularity: CppGpuBufferGranularity = CppGpuBufferGranularity.MultiCastMinimum,
|
|
):
|
|
# Check if `shape` is valid
|
|
if isinstance(shape, int):
|
|
shape = (shape,)
|
|
try:
|
|
shape = tuple(shape)
|
|
except TypeError:
|
|
raise ValueError("Shape must be a tuple-like or an integer.")
|
|
if any(s <= 0 for s in shape):
|
|
raise ValueError("Shape must be positive.")
|
|
# Create the buffer
|
|
buffer = CppRawGpuBuffer(np.prod(shape) * np.dtype(dtype).itemsize, granularity)
|
|
memptr = cp.cuda.MemoryPointer(cp.cuda.UnownedMemory(buffer.data(), buffer.bytes(), buffer), 0)
|
|
return cp.ndarray(shape, dtype=dtype, strides=strides, order=order, memptr=memptr)
|