Skip to main content

Mojo struct

MatmulKernels

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

Supported matmul kernels.

The configurations are named as: . BK, mma shape, and warp tile shape are decided internally.

Implemented traits

AnyType, Copyable, ExplicitlyCopyable, Movable, UnknownDestructibility

Aliases

ampere_128x128_4

alias ampere_128x128_4 = MatmulConfig(Index(128, 128, _bk_base[::DType,::Bool]()), Index(64, 64, _bk_base[::DType,::Bool]()), Index(1, 1, 1), UInt(4), UInt(1), UInt(1), UInt(1), UInt(1), False, Index(2, 2, 2), PDLLevel())

ampere_256x128_3

alias ampere_256x128_3 = MatmulConfig(Index(128, 256, (_bk_base[::DType,::Bool]() * 2)), Index(64, 64, (_bk_base[::DType,::Bool]() * 2)), Index(1, 1, 1), UInt(3), UInt(1), UInt(1), UInt(1), UInt(1), False, Index(2, 2, 2), PDLLevel())

ampere_256x64_4

alias ampere_256x64_4 = MatmulConfig(Index(64, 256, _bk_base[::DType,::Bool]()), Index(64, 64, _bk_base[::DType,::Bool]()), Index(1, 1, 1), UInt(4), UInt(1), UInt(1), UInt(1), UInt(1), False, Index(2, 2, 2), PDLLevel())

hopper_128x128_4

alias hopper_128x128_4 = MatmulConfig(Index(128, 128, _bk_base[::DType,::Bool]()), Index(64, 64, _bk_base[::DType,::Bool]()), Index(1, 1, 1), UInt(4), UInt(1), UInt(1), UInt(1), UInt(1), False, Index(2, 2, 2), PDLLevel())

tuning_config

alias tuning_config = MatmulConfig(Index(env_get_int[::StringSlice[::Bool(), env_get_int[::StringSlice[::Bool(), env_get_int[::StringSlice[::Bool()), Index(env_get_int[::StringSlice[::Bool(), env_get_int[::StringSlice[::Bool(), env_get_int[::StringSlice[::Bool()), Index(1, 1, 1), UInt(env_get_int[::StringSlice[::Bool()), UInt(env_get_int[::StringSlice[::Bool()), UInt(1), UInt(env_get_int[::StringSlice[::Bool()), UInt(1), False, Index(2, 2, 2), PDLLevel())

Was this page helpful?