nvmath-python Device APIs#

Overview#

The device module of nvmath-python nvmath.device offers integration with NVIDIA’s high-performance computing libraries through device APIs for cuFFTDx, cuBLASDx, and cuRAND. Detailed documentation for these libraries can be found at cuFFTDx, cuBLASDx, and cuRAND device APIs respectively.

Users may take advantage of the device module via the two approaches below:

  • Numba Extensions: Users can access these device APIs via Numba by utilizing specific extensions that simplify the process of defining functions, querying device traits, and calling device functions.

  • Third-party JIT Compilers: The APIs are also available through low-level interfaces in other JIT compilers, allowing advanced users to work directly with the raw device code.

Note

The fft and matmul device APIs in module nvmath.device currently supports cuFFTDx 1.2.0 and cuBLASDx 0.1.1, also available as part of MathDx 24.04. All functionalities from the C++ libraries are supported with the exception of cuFFTDx C++ APIs with a workspace argument, which are currently not available in nvmath-python.

API Reference#

Utility APIs (nvmath.device)#

current_device_lto()

A helper function to get the default code type for link time optimization (LTO) on the current device.

float16x2(x, y)

Create a Numba compliant vector object for float16 with vector length 2.

float16x4(x, y, z, w)

Create a Numba compliant vector object for float16 with vector length 4.

float32x2(x, y)

Create a Numba compliant vector object for float32 with vector length 2.

float64x2(x, y)

Create a Numba compliant vector object for float64 with vector length 2.

float16x2_type

A Numba compliant vector type object for float16 with vector length 2

float16x4_type

A Numba compliant vector type object for float16 with vector length 4

float32x2_type

A Numba compliant vector type object for float32 with vector length 2

float64x2_type

A Numba compliant vector type object for float64 with vector length 2

ISAVersion(major, minor)

A namedtuple class that encapsulates the code version.

Code(code_type, isa_version, data)

A namedtuple class that encapsulates code type, version, and buffer.

CodeType(kind, cc)

A namedtuple class that encapsulates code kind and compute capability.

ComputeCapability(major, minor)

A namedtuple class that encapsulates the major and minor compute capability.

CodeType(kind, cc)

A namedtuple class that encapsulates code kind and compute capability.

Symbol(variant, name)

A namedtuple class that encapsulates a device function symbol and which API it maps to.

Dim3([x, y, z])

A namedtuple class that encapsulates the dimensions for grids and blocks.

cuBLASDx APIs (nvmath.device)#

matmul(*[, compiler])

Create an BlasOptions object that encapsulates a compiled and ready-to-use device function for matrix multiplication.

BlasOptions(size, precision, data_type, *[, ...])

A class that encapsulates a partial BLAS device function.

LeadingDimension(a, b, c)

A namedtuple class that encapsulates the three leading dimensions in matrix multiplication \(C = \alpha Op(A) Op(B) + \beta C\).

TransposeMode(a, b)

A namedtuple class that encapsulates the transpose mode for input matrices A and B in matrix multiplication.

cuFFTDx APIs (nvmath.device)#

fft(*[, compiler])

Create an FFTOptions object that encapsulates a compiled and ready-to-use FFT device function.

FFTOptions(size, precision, fft_type, ...[, ...])

A class that encapsulates a partial FFT device function.

cuRAND Device APIs (nvmath.device.random)#

Compile([cc])

Compile the random device APIs with the specified compute capability.

Bit Generator and State APIs (nvmath.device.random)#

init(..., state)

Initialize the RNG state.

rand(state)

Generate 32 or 64 bits of randomness depending on the provided bit generator.

rand4(state)

Generate a 4-tuple of 32-bits of randomness from a Philox4_32_10 generator.

StatesMRG32k3a(n)

States of MRG32k3a bit generator.

StatesPhilox4_32_10(n)

States of Philox4_32_10 bit generator.

StatesSobol32(n)

States of Sobol32 bit generator.

StatesSobol64(n)

States of Sobol64 bit generator.

StatesScrambledSobol32(n)

States of ScrambledSobol32 bit generator.

StatesScrambledSobol64(n)

States of ScrambledSobol64 bit generator.

StatesXORWOW(n)

States of XORWOW bit generator.

Distribution Sampling APIs (nvmath.device.random)#

normal(state))

Sample a float from normal distribution using the specified bit generator state.

normal_double(state))

Sample a double from normal distribution using the specified bit generator state.

normal2(state))

Sample two floats from normal distribution using the specified bit generator state.

normal2_double(state))

Sample two doubles from normal distribution using the specified bit generator state.

normal4(state))

Sample four floats from normal distribution using the specified bit generator state.

log_normal(state, mean, stddev))

Sample a float from log-normal distribution using the specified bit generator state.

log_normal_double(state, mean, stddev))

Sample a double from log-normal distribution using the specified bit generator state.

log_normal2(state, mean, stddev))

Sample two floats from log-normal distribution using the specified bit generator state.

log_normal2_double(state, mean, stddev))

Sample two doubles from log-normal distribution using the specified bit generator state.

log_normal4(state, mean, stddev))

Sample four floats from log-normal distribution using the specified bit generator state.

poisson(state, Lambda))

Sample a uint from Poisson distribution using the specified bit generator state.

poisson4(state, Lambda))

Sample four uints from Poisson distribution using the specified bit generator state.

uniform(state))

Sample a float from uniform distribution using the specified bit generator state.

uniform_double(state))

Sample a double from uniform distribution using the specified bit generator state.

uniform2_double(state))

Sample two doubles from uniform distribution using the specified bit generator state.

uniform4(state))

Sample four floats from uniform distribution using the specified bit generator state.

Skip Ahead APIs (nvmath.device.random)#

skipahead(n, state)

Update the bit generator state to skip ahead n elements.

skipahead_sequence(n, state)

Update the bit generator state to skip ahead n sequences.

skipahead_subsequence(n, state)

Update the bit generator state to skip ahead n subsequences.

Helper Host APIs (nvmath.device.random_helpers)#

get_direction_vectors32(int set_, size_t size)

Get direction vectors for 32-bit quasirandom number generation.

get_direction_vectors64(int set_, size_t size)

Get direction vectors for 64-bit quasirandom number generation.

get_scramble_constants32(size_t size)

Get scramble constants for 32-bit scrambled Sobol' .

get_scramble_constants64(size_t size)

Get scramble constants for 64-bit scrambled Sobol' .

DirectionVectorSet(value[, names, module, ...])

See curandDirectionVectorSet.