compile_matmul_atomic_add#

nvmath.sparse.compile_matmul_atomic_add(add, *, dtype, compute_capability=None)[source]#

Compile a binary Python function with arguments of type dtype returning a value of type dtype to device code. This function must satisfy the mathematical requirements of a semiring addition operator and must perform an atomic addition.

Parameters:

  • add – The binary Python function to compile to device code.

  • dtype – The dtype of the epilog’s argument. Any scalar Numba type can be provided.

  • compute_capability – The target compute capability, specified as a string ('80', '89', …). The default is the compute capability of the current device.

Returns:

The compiled device code as a bytes object.