Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Nov 29, 2023

Conversation

zhanglx13
Copy link

@zhanglx13 zhanglx13 commented Nov 16, 2023

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

@zhanglx13 zhanglx13 merged commit 670ae80 into triton-mlir Nov 29, 2023
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants