Algorithm

class nvmath.linalg.advanced.Algorithm(algorithm: MatmulHeuristicResult)

An interface class to query algorithm capabilities and configure the algorithm.

Note that objects of this type should not be constructed directly by the user.

Attributes

algorithm_id

The ID of the algorithm (integer).

capabilities

Return the capabilities of this algorithm as a nvmath.linalg.advanced.AlgorithmCapabilities dataclass.

cluster_shape

A tuple representing the cluster shape (see MatmulAlgoConfigAttribute.CLUSTER_SHAPE_ID).

The value provided must be one of the cluster_shape_ids in the algorithm capabilities.

cta_swizzling

A flag indicating CTA swizzling (see MatmulAlgoConfigAttribute.CTA_SWIZZLING).

This can be set only if cta_swizzling is 1 in the algorithm capabilities.

custom_option

A value indicating the custom option (see MatmulAlgoConfigAttribute.CUSTOM_OPTION).

The value provided must be less than custom_option_max in the algorithm capabilities.

inner_shape

A value indicating the inner shape (see MatmulAlgoConfigAttribute.INNER_SHAPE_ID).

reduction_scheme

The reduction scheme used (see MatmulAlgoConfigAttribute.REDUCTION_SCHEME).

The value provided must be consistent with the reduction_scheme_mask in the algorithm capabilities.

split_k

The number of split-k steps (see MatmulAlgoConfigAttribute.SPLITK_NUM).

This can be set only if splitk_support is 1 in the algorithm capabilities.

stages

A tuple representing the stages (see MatmulAlgoConfigAttribute.STAGES_ID). The value provided must be one of the stages_ids in the algorithm capabilities.

tile

A tuple representing the tile (see MatmulAlgoConfigAttribute.TILE_ID). The value provided must be one of the tile_ids in the algorithm capabilities.

Methods

as_numpy() → ndarray[tuple[Any, ...], dtype[_ScalarT]]

Return a copy of this Algorithm as a NumPy array.

Converts the Algorithm to a NumPy array of type matmul_heuristic_result_dtype which may then be saved to disk using np.save(..., allow_pickle=False).

classmethod from_numpy(

array: ndarray[tuple[Any, ...], dtype[_ScalarT]],

) → Algorithm

Create an Algorithm from a NumPy array.

Creates an Algorithm by copying a NumPy array of type matmul_heuristic_result_dtype. The array must be of size 1, and may have been loaded from disk using np.load(..., allow_pickle=False).