nvmath.linalg.advanced.Algorithm

class nvmath.linalg.advanced.Algorithm(algorithm)

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

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

Methods

__init__(algorithm)

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.