The nvmath.device module is experimental and potentially subject to future changes.

Accumulator

class nvmath.device.Accumulator(*args)

Accumulator is an abstraction that provides the link between the global memory and register layouts. It offers operations like partitioning, copying data, and mapping register indices to matrix coordinates.

Refer to the cuBLASDx documentation for more details on how to use this class: api/other_tensors.html#accumulator-and-register-fragment-tensors

Methods

axpby()

clear()

get_alignment() → int

get_results(

out=None,

) → OpaqueTensor

is_index_in_bounds(index: int) → bool

Checks if the given index is within the bounds of the partitioned tensor. This is used to prevent out-of-bounds access in the kernel.

is_predicated() → bool

Checks if the current thread is predicated. This is used to determine if the thread should execute the kernel.

is_thread_active() → bool

Checks if the current thread takes part in GEMM.

make_empty_fragment() → OpaqueTensor

Creates an empty fragment tensor in register memory. The fragment layout is the same as the accumulator layout.

make_partition_and_copy(

src: OpaqueTensor,

) → OpaqueTensor

Same as partition_and_copy but returns the partitioned rmem tensor.

map_fragment_index(fragment_index: int) → tuple[int, int]

Maps the given fragment index to the (row, column) coordinate of the corresponding C matrix element. This is used to access the correct element in the partitioned tensor.

partition_and_copy(

src: OpaqueTensor,

dst: OpaqueTensor,

)

Partition gmem tensor and copy to rmem fragment.

partition_and_store(

tensor: OpaqueTensor,

)

partition_like_C(

gmem_c: OpaqueTensor,

) → OpaqueTensor

Partitions the given global memory tensor gmem_c into a partitioned tensor. The partitioned tensor is used for accessing the C matrix when working with register fragment.

Note

Not yet implemented, use map_fragment_index() instead.

size()