Skip to main content

/

Docs

Nightly

v24.6

Log in

Mojo module

tensor_core_async

This module provides abstractions for using Async Tensor Cores to perform asynchronous matrix multiplication operations.

Aliases​

  • ​supported_mma_shape = Tuple(VariadicPack(<store_to_mem(apply(:!lit.generator<[3]("x": !lit.ref<@stdlib::@builtin::@int::@Int, imm {}> read_mem, "y": !lit.ref<@stdlib::@builtin::@int::@Int, imm {}> read_mem, "z": !lit.ref<@stdlib::@builtin::@int::@Int, imm {}> read_mem) -> !lit.struct<@stdlib::@utils::@index::@IndexList<:@stdlib::@builtin::@int::@Int {3}, :@stdlib::@builtin::@int::@Int apply(:!lit.generator<("value": !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__init__(::IntLiteral)", apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>> @stdlib::@sys::@info::@"bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()"<:type @stdlib::@builtin::@int::@Int, :target current_target()>)), :@stdlib::@builtin::@bool::@Bool {:i1 0}>>> rebind(:!lit.generator<[3]("x": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, "y": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,1]> read_mem, "z": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,2]> read_mem) -> !lit.struct<@stdlib::@utils::@index::@IndexList<:@stdlib::@builtin::@int::@Int {3}, :@stdlib::@builtin::@int::@Int apply(:!lit.generator<("value": !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__init__(::IntLiteral)", apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>> @stdlib::@sys::@info::@"bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()"<:type @stdlib::@builtin::@int::@Int, :target current_target()>)), :@stdlib::@builtin::@bool::@Bool {:i1 0}>>> @stdlib::@utils::@index::@"Index[::Intable,::Intable,::Intable,::Int,::Bool]($0,$1,$2)"<:trait<@stdlib::@builtin::@int::@Intable> [@stdlib::@builtin::@int::@Int, {"__int__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @"fn(::Int, /) -> ::Int|fn(::Int, /) -> ::Int|s2ide7te5CPYGyhMuF+8dq+Xgtv1lztRWrKrzokOkcw=[fn(::Int, /) -> ::Int](::Int)"<:!lit.generator<(!lit.struct<@stdlib::@builtin::@int::@Int>, |) -> !lit.struct<@stdlib::@builtin::@int::@Int>> rebind(:!lit.generator<("self": !lit.struct<@stdlib::@builtin::@int::@Int>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__int__(::Int)")>), "__copyinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__copyinit__(::Int)_thunk"), "__del__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__del__(::Int)_thunk"), "__moveinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__moveinit__(::Int)_thunk")}], :trait<@stdlib::@builtin::@int::@Intable> [@stdlib::@builtin::@int::@Int, {"__int__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @"fn(::Int, /) -> ::Int|fn(::Int, /) -> ::Int|s2ide7te5CPYGyhMuF+8dq+Xgtv1lztRWrKrzokOkcw=[fn(::Int, /) -> ::Int](::Int)"<:!lit.generator<(!lit.struct<@stdlib::@builtin::@int::@Int>, |) -> !lit.struct<@stdlib::@builtin::@int::@Int>> rebind(:!lit.generator<("self": !lit.struct<@stdlib::@builtin::@int::@Int>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__int__(::Int)")>), "__copyinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__copyinit__(::Int)_thunk"), "__del__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__del__(::Int)_thunk"), "__moveinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__moveinit__(::Int)_thunk")}], :trait<@stdlib::@builtin::@int::@Intable> [@stdlib::@builtin::@int::@Int, {"__int__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @"fn(::Int, /) -> ::Int|fn(::Int, /) -> ::Int|s2ide7te5CPYGyhMuF+8dq+Xgtv1lztRWrKrzokOkcw=[fn(::Int, /) -> ::Int](::Int)"<:!lit.generator<(!lit.struct<@stdlib::@builtin::@int::@Int>, |) -> !lit.struct<@stdlib::@builtin::@int::@Int>> rebind(:!lit.generator<("self": !lit.struct<@stdlib::@builtin::@int::@Int>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__int__(::Int)")>), "__copyinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__copyinit__(::Int)_thunk"), "__del__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__del__(::Int)_thunk"), "__moveinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__moveinit__(::Int)_thunk")}], :@stdlib::@builtin::@int::@Int apply(:!lit.generator<("value": !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__init__(::IntLiteral)", apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>> @stdlib::@sys::@info::@"bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()"<:type @stdlib::@builtin::@int::@Int, :target current_target()>)), :@stdlib::@builtin::@bool::@Bool {:i1 0}>), store_to_mem({64}), store_to_mem({8}), store_to_mem({8}))), store_to_mem(apply(:!lit.generator<[3]("x": !lit.ref<@stdlib::@builtin::@int::@Int, imm {}> read_mem, "y": !lit.ref<@stdlib::@builtin::@int::@Int, imm {}> read_mem, "z": !lit.ref<@stdlib::@builtin::@int::@Int, imm {}> read_mem) -> !lit.struct<@stdlib::@utils::@index::@IndexList<:@stdlib::@builtin::@int::@Int {3}, :@stdlib::@builtin::@int::@Int apply(:!lit.generator<("value": !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__init__(::IntLiteral)", apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>> @stdlib::@sys::@info::@"bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()"<:type @stdlib::@builtin::@int::@Int, :target current_target()>)), :@stdlib::@builtin::@bool::@Bool {:i1 0}>>> rebind(:!lit.generator<[3]("x": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, "y": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,1]> read_mem, "z": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,2]> read_mem) -> !lit.struct<@stdlib::@utils::@index::@IndexList<:@stdlib::@builtin::@int::@Int {3}, :@stdlib::@builtin::@int::@Int apply(:!lit.generator<("value": !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__init__(::IntLiteral)", apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>> @stdlib::@sys::@info::@"bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()"<:type @stdlib::@builtin::@int::@Int, :target current_target()>)), :@stdlib::@builtin::@bool::@Bool {:i1 0}>>> @stdlib::@utils::@index::@"Index[::Intable,::Intable,::Intable,::Int,::Bool]($0,$1,$2)"<:trait<@stdlib::@builtin::@int::@Intable> [@stdlib::@builtin::@int::@Int, {"__int__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @"fn(::Int, /) -> ::Int|fn(::Int, /) -> ::Int|s2ide7te5CPYGyhMuF+8dq+Xgtv1lztRWrKrzokOkcw=[fn(::Int, /) -> ::Int](::Int)"<:!lit.generator<(!lit.struct<@stdlib::@builtin::@int::@Int>, |) -> !lit.struct<@stdlib::@builtin::@int::@Int>> rebind(:!lit.generator<("self": !lit.struct<@stdlib::@builtin::@int::@Int>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__int__(::Int)")>), "__copyinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__copyinit__(::Int)_thunk"), "__del__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__del__(::Int)_thunk"), "__moveinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__moveinit__(::Int)_thunk")}], :trait<@stdlib::@builtin::@int::@Intable> [@stdlib::@builtin::@int::@Int, {"__int__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @"fn(::Int, /) -> ::Int|fn(::Int, /) -> ::Int|s2ide7te5CPYGyhMuF+8dq+Xgtv1lztRWrKrzokOkcw=[fn(::Int, /) -> ::Int](::Int)"<:!lit.generator<(!lit.struct<@stdlib::@builtin::@int::@Int>, |) -> !lit.struct<@stdlib::@builtin::@int::@Int>> rebind(:!lit.generator<("self": !lit.struct<@stdlib::@builtin::@int::@Int>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__int__(::Int)")>), "__copyinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__copyinit__(::Int)_thunk"), "__del__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__del__(::Int)_thunk"), "__moveinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__moveinit__(::Int)_thunk")}], :trait<@stdlib::@builtin::@int::@Intable> [@stdlib::@builtin::@int::@Int, {"__int__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @"fn(::Int, /) -> ::Int|fn(::Int, /) -> ::Int|s2ide7te5CPYGyhMuF+8dq+Xgtv1lztRWrKrzokOkcw=[fn(::Int, /) -> ::Int](::Int)"<:!lit.generator<(!lit.struct<@stdlib::@builtin::@int::@Int>, |) -> !lit.struct<@stdlib::@builtin::@int::@Int>> rebind(:!lit.generator<("self": !lit.struct<@stdlib::@builtin::@int::@Int>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__int__(::Int)")>), "__copyinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, imm *[0,0]> read_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__copyinit__(::Int)_thunk"), "__del__" : !lit.generator<[1]("self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> = rebind(:!lit.generator<[1](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__del__(::Int)_thunk"), "__moveinit__" : !lit.generator<[2]("existing": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, "self": !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> = rebind(:!lit.generator<[2](!lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,0]> owned_in_mem, |, ?, !lit.ref<@stdlib::@builtin::@int::@Int, mut *[0,1]> byref_result) -> !kgen.none> @stdlib::@builtin::@int::@Int::@"__moveinit__(::Int)_thunk")}], :@stdlib::@builtin::@int::@Int apply(:!lit.generator<("value": !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>) -> !lit.struct<@stdlib::@builtin::@int::@Int>> @stdlib::@builtin::@int::@Int::@"__init__(::IntLiteral)", apply(:!lit.generator<() -> !lit.struct<@stdlib::@builtin::@int_literal::@IntLiteral>> @stdlib::@sys::@info::@"bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()"<:type @stdlib::@builtin::@int::@Int, :target current_target()>)), :@stdlib::@builtin::@bool::@Bool {:i1 0}>), store_to_mem({64}), store_to_mem({8}), store_to_mem({16})))>, True)):
  • ​WGMMA_K_BYTES = 32:

Structs​

Functions​