search

ROCm / triton

Development repository for the Triton language and compiler
MIT License
92 stars 29 forks source link

Add a cute tool to plot blocked, dotOperand, and mfma layout #407

Closed zhanglx13 closed 11 months ago

zhanglx13 commented 11 months ago

To use, sudo apt install texlive-full.

python3 plot_layout.py
usage: Draw triton layouts [-h] [-shape SHAPE SHAPE SHAPE] [-plot {blocked,dot,wmma,lds}] [-nonKDim {32}] [-sizePerThread SIZEPERTHREAD SIZEPERTHREAD] [-threadsPerWarp THREADSPERWARP THREADSPERWARP] [-warpsPerCTA WARPSPERCTA WARPSPERCTA] [-order ORDER ORDER] [-kpack {4,8}]
                           [-lds_layout {swizzle,padding,none}] [-lds_access {read,write,none}] [-wave_size {32,64}] [-o O] [-mfmaTrans] [--keep]

optional arguments:
  -h, --help            show this help message and exit
  -shape SHAPE SHAPE SHAPE
                        Tensor shape in the form of M,N,K
  -plot {blocked,dot,wmma,lds}
                        choose plot mode
  -nonKDim {32}         mfma instruction dim, only 32 is supported for now
  -sizePerThread SIZEPERTHREAD SIZEPERTHREAD
  -threadsPerWarp THREADSPERWARP THREADSPERWARP
  -warpsPerCTA WARPSPERCTA WARPSPERCTA
  -order ORDER ORDER
  -kpack {4,8}          vector length during LDS load, same as vec
  -lds_layout {swizzle,padding,none}
                        choose the LDS data layout
  -lds_access {read,write,none}
                        choose LDS access mode
  -wave_size {32,64}    choose the wmma instruction mode
  -o O                  output pdf file name (without surfix)
  -mfmaTrans            If set, then use mfma.trans layout
  --keep                If set, keep the generated .tex file