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

Matmul#

class nvmath.device.Matmul(
size,
precision,
data_type,
*,
sm=None,
block_size=None,
block_dim=None,
leading_dimension=None,
transpose_mode=None,
arrangement=None,
alignment=None,
function='MM',
static_block_dim=False,
execution='Block',
with_pipeline: bool = False,
enable_input_streaming: bool = False,
)[source]#

A class that encapsulates a partial Matmul device function. A partial device function can be queried for available or optimal values for some knobs (such as leading_dimension or block_dim).

Changed in version 0.7.0: Matmul has replaced BlasOptions and BlasOptionsComplete.

Parameters:

  • size – A sequence of integers denoting the three dimensions (m, n, k) for the matrix multiplication problem.

  • precision – The computation precision, either a single numpy dtype or a 3-sequence of numpy dtypes (a, b, c) for the A, B, and C matrices. Supported dtypes are numpy.float16, numpy.float32, numpy.float64 and the signed and unsigned integer dtypes of 8, 16, 32 or 64 bits.

  • data_type – The data type of the input matrices, can be either 'real' or 'complex'.

  • sm (ComputeCapability) – Target mathdx compute-capability.

  • block_size (int) – The total block size, optional. If not provided or set to 'suggested', will be set to a suggested value for 1D block dim.

  • block_dim (Dim3) – The block dimension for launching the CUDA kernel, optional. If not provided or set to 'suggested', will be set to a suggested value. Cannot be used when block_size is explicitly specified.

  • leading_dimension (LeadingDimension) – The leading dimensions for the input matrices, optional. If not provided, will be set to match the matrix row/column dimension. Alternatively, if provided as 'suggested', will be set to a suggested value for optimal performance.

  • transpose_mode (TransposeMode) – The transpose mode for all input matrices ; transpose_mode or arrangement must be provided.

  • arrangement (Arrangement) – The arrangement for all input matrices ; transpose_mode or arrangement must be provided.

  • alignment (Alignment) – The alignment for the input matrices in shared memory. Defines the alignments (in bytes) of the input matrices A, B, and C (either arrays or wrapped in opaque tensors) that are passed to the execute(…) method. The default alignment is equal to the element size of the matrix. When a suggested layout is used, the alignment is greater than or equal to the element size.

  • function (str) – A string specifying the name of the function. Currently supports 'MM' (default) for matrix multiplication.

  • static_block_dim (bool) – If set to True, cuBLASDx assumes the kernel is launched with a block dimension exactly equal to block_dim, which enables additional optimizations. The default is False.

  • execution (str) – A string specifying the execution method, can be 'Block' or 'Thread'.

  • with_pipeline (bool) – If set to True, the device function is generated for pipelined execution and must be used only through nvmath.device.DevicePipeline or nvmath.device.TilePipeline (see nvmath.device.Matmul.suggest_device_pipeline()). Executing such a matmul with the regular execute(...) interface is incorrect. The default is False.

  • enable_input_streaming (bool) – If set to True, informs the library that no per-element load transforms are applied to the inputs, which allows input data to be streamed from shared memory directly to the MMA units without passing through registers. This may improve performance. The default is False.

See also

The attributes of this class provide a 1:1 mapping with the CUDA C++ cuBLASDx APIs. For further details, please refer to cuBLASDx documentation.

Attributes

a_dim#
a_size#
a_value_type#
alignment#

The alignment for the input matrices in shared memory. Defines the alignments (in bytes) of the input matrices A, B, and C (either arrays or wrapped in opaque tensors) that are passed to the execute(…) method. The default alignment is equal to the element size of the matrix. When a suggested layout is used, the alignment is greater than or equal to the element size.

arrangement#

The arrangement for all input matrices ; transpose_mode or arrangement must be provided.

b_dim#
b_size#
b_value_type#
block_dim#

The block dimension for launching the CUDA kernel, optional. If not provided or set to 'suggested', will be set to a suggested value. Cannot be used when block_size is explicitly specified.

block_size#

The total block size, optional. If not provided or set to 'suggested', will be set to a suggested value for 1D block dim.

c_dim#
c_size#
c_value_type#
data_type#

The data type of the input matrices, can be either 'real' or 'complex'.

enable_input_streaming#

If set to True, informs the library that no per-element load transforms are applied to the inputs, which allows input data to be streamed from shared memory directly to the MMA units without passing through registers. This may improve performance. The default is False.

execution#

A string specifying the execution method, can be 'Block' or 'Thread'.

function#

A string specifying the name of the function. Currently supports 'MM' (default) for matrix multiplication.

leading_dimension#

The leading dimensions for the input matrices, optional. If not provided, will be set to match the matrix row/column dimension. Alternatively, if provided as 'suggested', will be set to a suggested value for optimal performance.

max_threads_per_block#
precision#

The computation precision, either a single numpy dtype or a 3-sequence of numpy dtypes (a, b, c) for the A, B, and C matrices. Supported dtypes are numpy.float16, numpy.float32, numpy.float64 and the signed and unsigned integer dtypes of 8, 16, 32 or 64 bits.

size#

A sequence of integers denoting the three dimensions (m, n, k) for the matrix multiplication problem.

sm#

Target mathdx compute-capability.

static_block_dim#

If set to True, cuBLASDx assumes the kernel is launched with a block dimension exactly equal to block_dim, which enables additional optimizations. The default is False.

transpose_mode#

The transpose mode for all input matrices ; transpose_mode or arrangement must be provided.

with_pipeline#

If set to True, the device function is generated for pipelined execution and must be used only through nvmath.device.DevicePipeline or nvmath.device.TilePipeline (see nvmath.device.Matmul.suggest_device_pipeline()). Executing such a matmul with the regular execute(...) interface is incorrect. The default is False.

Methods

execute(*args)[source]#
get_accumulator() → Accumulator[source]#
get_layout_gmem_a(
leading_dimension: int | None = None,
) → Layout[source]#
get_layout_gmem_b(
leading_dimension: int | None = None,
) → Layout[source]#
get_layout_gmem_c(
leading_dimension: int | None = None,
) → Layout[source]#
get_layout_rmem_c() → Layout[source]#
get_layout_smem_a() → Layout[source]#
get_layout_smem_b() → Layout[source]#
get_layout_smem_c() → Layout[source]#
get_shared_storage_size() → int[source]#
get_shared_storage_size(lda: int, ldb: int, ldc: int) → int
get_shared_storage_size(
matrix_a_layout: Layout,
matrix_b_layout: Layout,
matrix_c_layout: Layout,
) → int
get_shared_storage_size_ab() → int[source]#
get_shared_storage_size_ab(lda: int, ldb: int) → int
get_shared_storage_size_ab(
matrix_a_layout: Layout,
matrix_b_layout: Layout,
) → int
suggest_accumulator() → Accumulator[source]#
suggest_device_pipeline(
pipeline_depth: int,
a: ndarray,
b: ndarray,
) → DevicePipeline[source]#
suggest_layout_rmem_c() → Layout[source]#
suggest_layout_smem_a() → Layout[source]#
suggest_layout_smem_b() → Layout[source]#
suggest_layout_smem_c() → Layout[source]#
valid(*knobs)[source]#