For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo function
shmem_malloc
shmem_malloc[dtype: DType](size: Int) -> UnsafePointer[Scalar[dtype], MutExternalOrigin]
Collectively allocate symmetric memory.
The shmem_malloc routine is a collective operation on the world team and
returns the symmetric address of a block of at least size bytes, which shall
be suitably aligned so that it may be assigned to a pointer to any type of
object. This space is allocated from the symmetric heap (in contrast to
malloc, which allocates from the private heap). When size is zero, the
shmem_malloc routine performs no action and returns a null pointer;
otherwise, shmem_malloc calls a procedure that is semantically equivalent to
shmem_barrier_all on exit. This ensures that all PEs participate in the
memory allocation, and that the memory on other PEs can be used as soon as
the local PE returns. The value of the size argument must be identical on
all PEs; otherwise, the behavior is undefined.
Note: on ROCSHMEM this will raise an error if exceeding the default static symmetric heap size of 1GB across all GPUs, while NVSHMEM uses dynamic symmetric heap allocation and won't error.
Parameters:
- dtype (
DType): The data type of elements to allocate memory for.
Args:
- size (
Int): The number of elements to be allocated from the symmetric heap.
Returns:
UnsafePointer[Scalar[dtype], MutExternalOrigin]: The symmetric address of the allocated space; otherwise, it returns a
null pointer.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!