pylibraft.random.rmat_rectangular_generator

pylibraft.random.rmat_rectangular_generator#

4 min read time

Applies to Linux

Classes#

`cai_wrapper`	Simple wrapper around a CUDA array interface object to reduce
`Handle`	Handle is a lightweight python wrapper around the corresponding

Functions#

`auto_sync_handle`(f)	auto_sync_handle(f)
`rmat`(*args[, handle])	rmat(out, theta, r_scale, c_scale, seed=12345, handle=None)

Module Contents#

class pylibraft.random.rmat_rectangular_generator.cai_wrapper(cai_arr)#

Bases: pylibraft.common.ai_wrapper.ai_wrapper

Simple wrapper around a CUDA array interface object to reduce boilerplate for extracting common information from the underlying dictionary.

class pylibraft.random.rmat_rectangular_generator.Handle#

Bases: DeviceResources

Handle is a lightweight python wrapper around the corresponding C++ class of handle_t exposed by RAFT’s C++ interface. Refer to the header file raft/core/handle.hpp for interface level details of this struct

Note: This API is officially deprecated in favor of DeviceResources and will be removed in a future release.

Parameters#

streamOptional stream to use for ordering CUDA instructions: Accepts pylibraft.common.Stream() or uintptr_t (cudaStream_t)

Examples#

Basic usage:

>>> from pylibraft.common import Stream, Handle
>>> stream = Stream()
>>> handle = Handle(stream)
>>>
>>> # call algos here
>>>
>>> # final sync of all work launched in the stream of this handle
>>> # this is same as `raft.cuda.Stream.sync()` call, but safer in case
>>> # the default stream inside the `handle_t` is being used
>>> handle.sync()
>>> del handle  # optional!

Using a cuPy stream with RAFT device_resources:

>>> import cupy
>>> from pylibraft.common import Stream, Handle
>>>
>>> cupy_stream = cupy.cuda.Stream()
>>> handle = Handle(stream=cupy_stream.ptr)

Using a RAFT stream with CuPy ExternalStream:

>>> import cupy
>>> from pylibraft.common import Stream
>>>
>>> raft_stream = Stream()
>>> cupy_stream = cupy.cuda.ExternalStream(raft_stream.get_ptr())

static __reduce_cython__(*args, **kwargs)#: Handle.__reduce_cython__(self)

static __setstate_cython__(*args, **kwargs)#: Handle.__setstate_cython__(self, __pyx_state)

__getstate__()#: Handle.__getstate__(self)

__setstate__(state)#: Handle.__setstate__(self, state)

pylibraft.random.rmat_rectangular_generator.auto_sync_handle(f)#: auto_sync_handle(f) Decorator to automatically call sync on a raft handle when

it isn’t passed to a function.

When a handle=None is passed to the wrapped function, this decorator will automatically create a default handle for the function, and call sync on that handle when the function exits.

This will also insert the appropriate docstring for the handle parameter

pylibraft.random.rmat_rectangular_generator.rmat(*args, handle=None, **kwargs)#

rmat(out, theta, r_scale, c_scale, seed=12345, handle=None)

Generate RMAT adjacency list based on the input distribution.

out:
CUDA array interface compliant matrix shape (n_edges, 2). This will contain the src/dst node ids stored consecutively like a pair.

theta:
CUDA array interface compliant matrix shape (max(r_scale, c_scale) * 4) This stores the probability distribution at each RMAT level

r_scale: int
log2 of number of source nodes

c_scale: int
log2 of number of destination nodes

seed: int
random seed used for reproducibility

handleOptional RAFT resource handle for reusing CUDA resources.
If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.
>>> import cupy as cp
>>> from pylibraft.common import Handle
>>> from pylibraft.random import rmat
>>> n_edges = 5000
>>> r_scale = 16
>>> c_scale = 14
>>> theta_len = max(r_scale, c_scale) * 4
>>> out = cp.empty((n_edges, 2), dtype=cp.int32)
>>> theta = cp.random.random_sample(theta_len, dtype=cp.float32)
>>> # A single RAFT handle can optionally be reused across
>>> # pylibraft functions.
>>> handle = Handle()
>>> rmat(out, theta, r_scale, c_scale, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()