Mojo struct
MMAOperandOffsetFn
@register_passable(trivial)
struct MMAOperandOffsetFn[dtype: DType, BMN: Int, BK: Int, swizzle: TensorMapSwizzle, is_k_major: Bool, WMMA_MN: Int, WMMA_K: Int]
Implemented traits
AnyType,
Copyable,
Movable,
UnknownDestructibility
Aliases
canonical_K
alias canonical_K = 0 if (sizeof[::DType,__mlir_type.!kgen.target]() == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int apply(:!lit.generator<("self": !lit.struct<@stdlib::@gpu::@host::@_nvidia_cuda::@TensorMapSwizzle>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@gpu::@host::@_nvidia_cuda::@TensorMapSwizzle::@"bytes(::TensorMapSwizzle)", swizzle), "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@sys::@info::@"sizeof[::DType,__mlir_type.!kgen.target]()"<:@stdlib::@builtin::@dtype::@DType dtype, :target apply(:!lit.generator<() -> !kgen.target> @stdlib::@sys::@info::@"_current_target()")>), "value">, 0), {1}, apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@sys::@info::@"sizeof[::DType,__mlir_type.!kgen.target]()"<:@stdlib::@builtin::@dtype::@DType dtype, :target apply(:!lit.generator<() -> !kgen.target> @stdlib::@sys::@info::@"_current_target()")>)), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int apply(:!lit.generator<("self": !lit.struct<@stdlib::@gpu::@host::@_nvidia_cuda::@TensorMapSwizzle>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@gpu::@host::@_nvidia_cuda::@TensorMapSwizzle::@"bytes(::TensorMapSwizzle)", swizzle), "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@sys::@info::@"sizeof[::DType,__mlir_type.!kgen.target]()"<:@stdlib::@builtin::@dtype::@DType dtype, :target apply(:!lit.generator<() -> !kgen.target> @stdlib::@sys::@info::@"_current_target()")>), "value">, 0), {1}, apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@sys::@info::@"sizeof[::DType,__mlir_type.!kgen.target]()"<:@stdlib::@builtin::@dtype::@DType dtype, :target apply(:!lit.generator<() -> !kgen.target> @stdlib::@sys::@info::@"_current_target()")>)), "value">) == 0) ^ True) & ((swizzle.bytes() < 0) ^ (sizeof[::DType,__mlir_type.!kgen.target]() < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int apply(:!lit.generator<("self": !lit.struct<@stdlib::@gpu::@host::@_nvidia_cuda::@TensorMapSwizzle>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@gpu::@host::@_nvidia_cuda::@TensorMapSwizzle::@"bytes(::TensorMapSwizzle)", swizzle), "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@sys::@info::@"sizeof[::DType,__mlir_type.!kgen.target]()"<:@stdlib::@builtin::@dtype::@DType dtype, :target apply(:!lit.generator<() -> !kgen.target> @stdlib::@sys::@info::@"_current_target()")>), "value">, 0), {1}, apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@sys::@info::@"sizeof[::DType,__mlir_type.!kgen.target]()"<:@stdlib::@builtin::@dtype::@DType dtype, :target apply(:!lit.generator<() -> !kgen.target> @stdlib::@sys::@info::@"_current_target()")>)), "value">) if (swizzle != TensorMapSwizzle(__init__[__mlir_type.!pop.int_literal](0))) else BK
canonical_layout
alias canonical_layout = tile_to_descriptor[::DType,::Layout,::Bool]()
canonical_layout_flat
alias canonical_layout_flat = tile_layout_k_major[::DType,::Int,::Int,::TensorMapSwizzle]() if is_k_major else tile_layout_k_major[::DType,::Int,::Int,::TensorMapSwizzle]() if is_k_major else tile_layout_mn_major[::DType,::Int,::Int,::TensorMapSwizzle]()
canonical_layout_size
alias canonical_layout_size = tile_to_descriptor[::DType,::Layout,::Bool]().size()
layout
alias layout = tile_layout_k_major[::DType,::Int,::Int,::TensorMapSwizzle]() if is_k_major else tile_layout_mn_major[::DType,::Int,::Int,::TensorMapSwizzle]()
layout_size
alias layout_size = tile_layout_k_major[::DType,::Int,::Int,::TensorMapSwizzle]() if is_k_major else tile_layout_mn_major[::DType,::Int,::Int,::TensorMapSwizzle]().size()
Methods
__init__
__init__() -> Self
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!