tilelang.tileop.gemm¶

子模組¶

類別¶

函式¶

`gemm_infer_layout`(gemm, target, thread_bounds)
`gemm_lower`(gemm, layout_map, target, thread_bounds, ...)

Package Contents¶

tilelang.tileop.gemm.gemm_infer_layout(gemm, target, thread_bounds)¶

參數:

gemm (gemm_mma.GemmMMA)
target (tvm.target.Target)
thread_bounds (tvm.ir.Range)

tilelang.tileop.gemm.gemm_lower(gemm, layout_map, target, thread_bounds, thread_var, mbar_phase_expr)¶

參數:

gemm (gemm_mma.GemmMMA)
target (tvm.target.Target)
thread_bounds (tvm.ir.Range)
thread_var (tvm.tir.Var)
mbar_phase_expr (tvm.tir.PrimExpr)

class tilelang.tileop.gemm.Gemm¶

Bases: tvm.ir.base.Node, tvm.runtime.Scriptable

property A¶

property B¶

property C¶

property APtr¶

property BPtr¶

property CPtr¶

property M¶

property N¶

property K¶

property trans_A¶

property trans_B¶

property stride_A¶

property stride_B¶

property offset_A¶

property offset_B¶

property clear_accum¶

property k_pack¶

property wg_wait¶

property is_tcgen05¶

property sf_a_id¶

property sf_b_id¶

infer_layout(target, thread_nums)¶

Infer the layout for the GEMM operation based on target architecture.

參數:

target (tvm.target.Target)
thread_nums (int)

lower(layout_map, target, thread_bounds, thread_var, mbar_phase_expr)¶

Lower the GEMM operation to TIR statements based on target architecture.

參數:

layout_map (dict)
target (tvm.target.Target)
thread_bounds (tvm.ir.Range)
thread_var (tvm.tir.Var)
mbar_phase_expr (tvm.tir.PrimExpr)