Algorithm#

class nvmath.linalg.advanced.Algorithm(algorithm)[source]#

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.

Methods

__init__(algorithm)[source]#

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.