Skip to main content

Mojo struct

MatmulConfig

@register_passable(trivial) struct MatmulConfig[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = True]

Static configuration of SM90 GPU matmul.

Fields

  • block_tile_shape (IndexList[3]):
  • mma_shape (IndexList[3]):
  • cluster_shape (IndexList[3]):
  • num_pipeline_stages (UInt):
  • num_k_partitions (UInt):
  • num_consumer (UInt):
  • partitioned_multicast (Bool):

Implemented traits

AnyType, Copyable, Equatable, Hashable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, Stringable, Writable

comptime members

__copyinit__is_trivial

comptime __copyinit__is_trivial = True

__del__is_trivial

comptime __del__is_trivial = True

__moveinit__is_trivial

comptime __moveinit__is_trivial = True

Methods

__init__

__init__(M: Int, N: Int, K: Int, num_k_partitions: UInt = 1, partitioned_multicast: Bool = False, pdl_level: PDLLevel = PDLLevel.OFF) -> Self

Initialize MatmulConfig by computing optimal values from M, N, K.

Args:

  • M (Int): The M dimension of the matmul.
  • N (Int): The N dimension of the matmul.
  • K (Int): The K dimension of the matmul.
  • num_k_partitions (UInt): Number of K partitions.
  • partitioned_multicast (Bool): Whether to use partitioned multicast.
  • pdl_level (PDLLevel): PDL level for grid controls.

__eq__

__eq__(self, other: Self) -> Bool

Returns:

Bool

pdl_level

pdl_level(self) -> PDLLevel

Returns:

PDLLevel

to_base_config

to_base_config(self) -> MatmulConfig[a_type, b_type, c_type, transpose_b]

Convert to base MatmulConfig from utils_gpu.

Returns:

MatmulConfig

__str__

__str__(self) -> String

Returns:

String

write_to

write_to(self, mut writer: T)

__repr__

__repr__(self) -> String

Returns:

String

__hash__

__hash__[H: Hasher](self, mut hasher: H)

Updates hasher with the underlying bytes.

Parameters:

  • H (Hasher): The hasher type.

Args:

  • hasher (H): The hasher instance.

Was this page helpful?