tilelang.tileop.gemm

子模組

類別

GemmPy

函式

gemm_py_infer_layout(gemm_py, target, thread_bounds)

gemm_py_lower(gemm_py, layout_map, target, ...)

Package Contents

tilelang.tileop.gemm.gemm_py_infer_layout(gemm_py, target, thread_bounds)
參數:

  • gemm_py (gemm_mma.GemmMMA)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

tilelang.tileop.gemm.gemm_py_lower(gemm_py, layout_map, target, thread_bounds, thread_var)
參數:

  • gemm_py (gemm_mma.GemmMMA)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_var (tvm.tir.Var)

class tilelang.tileop.gemm.GemmPy

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
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)

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)