|
| 1 | +\documentclass[tikz,border=8pt]{standalone} |
| 2 | +\usepackage[T1]{fontenc} |
| 3 | +\usepackage{lmodern} |
| 4 | +\usepackage{amsmath} |
| 5 | +\usepackage{tikz} |
| 6 | +\usetikzlibrary{arrows.meta,calc,fit,positioning,backgrounds} |
| 7 | + |
| 8 | +\definecolor{ink}{HTML}{172033} |
| 9 | +\definecolor{muted}{HTML}{61708A} |
| 10 | +\definecolor{grid}{HTML}{D9E0EA} |
| 11 | +\definecolor{panel}{HTML}{F7F9FC} |
| 12 | +\definecolor{stagezero}{HTML}{2F80ED} |
| 13 | +\definecolor{stagezerofill}{HTML}{EAF3FF} |
| 14 | +\definecolor{stageone}{HTML}{E79A19} |
| 15 | +\definecolor{stageonefill}{HTML}{FFF4DD} |
| 16 | +\definecolor{fullbar}{HTML}{8E5CC7} |
| 17 | +\definecolor{emptybar}{HTML}{168B6A} |
| 18 | +\definecolor{danger}{HTML}{C94F5C} |
| 19 | + |
| 20 | +\begin{document} |
| 21 | +\sffamily |
| 22 | +\begin{tikzpicture}[ |
| 23 | + x=1.12cm, |
| 24 | + y=1cm, |
| 25 | + >={Latex[length=2.1mm,width=1.35mm]}, |
| 26 | + title/.style={font=\bfseries\fontsize{17}{20}\selectfont, text=ink}, |
| 27 | + subtitle/.style={font=\fontsize{9}{11}\selectfont, text=muted}, |
| 28 | + lane label/.style={font=\bfseries\fontsize{9}{11}\selectfont, text=ink, align=right}, |
| 29 | + tick/.style={font=\bfseries\fontsize{8}{10}\selectfont, text=muted}, |
| 30 | + bar/.style={rounded corners=2pt, minimum height=8mm, text width=2.15cm, inner xsep=4pt, |
| 31 | + font=\fontsize{8}{9.5}\selectfont, align=center, line width=.75pt}, |
| 32 | + s0/.style={bar, draw=stagezero, fill=stagezerofill, text=ink}, |
| 33 | + s1/.style={bar, draw=stageone, fill=stageonefill, text=ink}, |
| 34 | + neutral/.style={bar, draw=grid!90!ink, fill=white, text=ink}, |
| 35 | + note/.style={draw=grid!90!ink, fill=white, rounded corners=3pt, inner sep=6pt, |
| 36 | + font=\fontsize{8}{10}\selectfont, text=ink, align=left}, |
| 37 | + callout/.style={font=\fontsize{7.2}{8.8}\selectfont, text=muted, align=center}, |
| 38 | + flow/.style={->, line width=.7pt, color=muted}, |
| 39 | + dependency/.style={->, densely dashed, line width=.7pt, color=muted}, |
| 40 | + full marker/.style={circle, draw=fullbar, fill=fullbar!12, line width=.8pt, |
| 41 | + minimum size=5.2mm, inner sep=0pt, |
| 42 | + font=\bfseries\fontsize{6.6}{7}\selectfont, text=fullbar}, |
| 43 | + empty marker/.style={circle, draw=emptybar, fill=emptybar!12, line width=.8pt, |
| 44 | + minimum size=5.2mm, inner sep=0pt, |
| 45 | + font=\bfseries\fontsize{6.6}{7}\selectfont, text=emptybar}, |
| 46 | +] |
| 47 | + |
| 48 | +% Header |
| 49 | +\node[title, anchor=west] at (0,8.25) {Blockscaled TMA pipeline: how the mbarriers hand off each stage}; |
| 50 | +\node[subtitle, anchor=west] at (0,7.82) |
| 51 | + {Two-stage main TMA path from \texttt{BlockscaledTmaGemv}; optional input-scale pipeline omitted; not cycle-accurate}; |
| 52 | + |
| 53 | +% Legend |
| 54 | +\node[s0, minimum width=1.45cm, minimum height=5.5mm, text width=1.15cm, anchor=west] at (12.2,7.25) {stage 0}; |
| 55 | +\node[s1, minimum width=1.45cm, minimum height=5.5mm, text width=1.15cm, anchor=west] at (13.85,7.25) {stage 1}; |
| 56 | +\node[full marker] (legend-full) at (15.75,7.25) {F}; |
| 57 | +\node[callout, anchor=west] at (16.05,7.25) {full}; |
| 58 | +\node[empty marker] (legend-empty) at (17.05,7.25) {E}; |
| 59 | +\node[callout, anchor=west] at (17.35,7.25) {empty}; |
| 60 | + |
| 61 | +% Timeline panel and lanes |
| 62 | +\fill[panel, rounded corners=4pt] (0,1.15) rectangle (18.7,6.85); |
| 63 | +\draw[grid, rounded corners=4pt] (0,1.15) rectangle (18.7,6.85); |
| 64 | + |
| 65 | +\foreach \x/\label in {2.3/setup,4.6/prologue,7.2/$k_0$,10.0/$k_1$,12.8/$k_2$,15.6/drain} { |
| 66 | + \draw[grid, densely dashed] (\x,1.35) -- (\x,6.55); |
| 67 | + \node[tick] at (\x,6.7) {\label}; |
| 68 | +} |
| 69 | + |
| 70 | +\node[lane label, anchor=east] at (2.05,5.95) {producer warp\\\scriptsize control path}; |
| 71 | +\node[lane label, anchor=east] at (2.05,4.15) {TMA engine\\\scriptsize async proxy}; |
| 72 | +\node[lane label, anchor=east] at (2.05,2.35) {compute warps\\\scriptsize generic proxy}; |
| 73 | + |
| 74 | +\draw[grid!80] (2.3,5.35) -- (18.35,5.35); |
| 75 | +\draw[grid!80] (2.3,3.55) -- (18.35,3.55); |
| 76 | + |
| 77 | +% Producer lane |
| 78 | +\node[neutral, minimum width=1.75cm] (init) at (3.25,5.95) |
| 79 | + {initialize\\$2S$ barriers}; |
| 80 | +\node[s0, minimum width=2.05cm] (p0) at (5.35,5.95) |
| 81 | + {acquire $S_0$\\issue tile $k_0$}; |
| 82 | +\node[s1, minimum width=2.15cm] (p1) at (7.85,5.95) |
| 83 | + {acquire $S_1$\\issue tile $k_1$}; |
| 84 | +\node[s0, minimum width=2.3cm] (p2) at (10.65,5.95) |
| 85 | + {wait $E_0$; re-arm\\issue tile $k_2$}; |
| 86 | +\node[s1, minimum width=2.3cm] (p3) at (13.55,5.95) |
| 87 | + {wait $E_1$; re-arm\\issue tile $k_3$}; |
| 88 | +\node[neutral, minimum width=1.55cm] (tail) at (16.45,5.95) |
| 89 | + {producer\\tail}; |
| 90 | + |
| 91 | +% TMA lane |
| 92 | +\node[s0, minimum width=2.25cm] (t0) at (6.0,4.15) |
| 93 | + {copy $W+X(+S_{FW})$\\into $S_0$}; |
| 94 | +\node[s1, minimum width=2.25cm] (t1) at (8.6,4.15) |
| 95 | + {copy $W+X(+S_{FW})$\\into $S_1$}; |
| 96 | +\node[s0, minimum width=2.25cm] (t2) at (11.4,4.15) |
| 97 | + {copy next tile\\into reused $S_0$}; |
| 98 | +\node[s1, minimum width=2.25cm] (t3) at (14.3,4.15) |
| 99 | + {copy next tile\\into reused $S_1$}; |
| 100 | + |
| 101 | +\node[full marker] (f0a) at (7.05,4.15) {F$_0$}; |
| 102 | +\node[full marker] (f1a) at (9.65,4.15) {F$_1$}; |
| 103 | +\node[full marker] (f0b) at (12.45,4.15) {F$_0$}; |
| 104 | +\node[full marker] (f1b) at (15.35,4.15) {F$_1$}; |
| 105 | + |
| 106 | +% Consumer lane |
| 107 | +\node[s0, minimum width=2.35cm] (c0) at (7.45,2.35) |
| 108 | + {wait $F_0$\\load + accumulate $k_0$}; |
| 109 | +\node[s1, minimum width=2.35cm] (c1) at (10.15,2.35) |
| 110 | + {wait $F_1$\\load + accumulate $k_1$}; |
| 111 | +\node[s0, minimum width=2.35cm] (c2) at (12.95,2.35) |
| 112 | + {wait $F_0$\\load + accumulate $k_2$}; |
| 113 | +\node[s1, minimum width=2.35cm] (c3) at (15.85,2.35) |
| 114 | + {wait $F_1$\\load + accumulate $k_3$}; |
| 115 | + |
| 116 | +\node[empty marker] (e0a) at (8.55,2.35) {E$_0$}; |
| 117 | +\node[empty marker] (e1a) at (11.25,2.35) {E$_1$}; |
| 118 | +\node[empty marker] (e0b) at (14.05,2.35) {E$_0$}; |
| 119 | +\node[empty marker] (e1b) at (16.95,2.35) {E$_1$}; |
| 120 | + |
| 121 | +% Issue/completion/reuse dependencies |
| 122 | +\draw[dependency] (p0.south) -- node[callout, left] {barrier ptr} (t0.north west); |
| 123 | +\draw[dependency] (p1.south) -- (t1.north west); |
| 124 | +\draw[dependency] (p2.south) -- (t2.north west); |
| 125 | +\draw[dependency] (p3.south) -- (t3.north west); |
| 126 | + |
| 127 | +\draw[dependency, color=fullbar] (f0a.south) -- (c0.north); |
| 128 | +\draw[dependency, color=fullbar] (f1a.south) -- (c1.north); |
| 129 | +\draw[dependency, color=fullbar] (f0b.south) -- (c2.north); |
| 130 | +\draw[dependency, color=fullbar] (f1b.south) -- (c3.north); |
| 131 | + |
| 132 | +\draw[dependency, color=emptybar] (e0a.north) to[out=90,in=-90] (p2.south); |
| 133 | +\draw[dependency, color=emptybar] (e1a.north) to[out=90,in=-90] (p3.south); |
| 134 | + |
| 135 | +\node[callout, text width=2.4cm] at (17.45,4.15) |
| 136 | + {TMA completion performs\\\texttt{complete\_tx} against $F_i$}; |
| 137 | +\node[callout, text width=2.6cm] at (17.35,1.55) |
| 138 | + {before $E_i$: async-shared fence, warp sync, then release}; |
| 139 | + |
| 140 | +% Lower explanatory panels |
| 141 | +\node[note, anchor=north west, text width=5.35cm] (contract) at (0,0.72) { |
| 142 | + \textbf{One circular stage owns two 64-bit mbarriers}\\[2pt] |
| 143 | + $F_i$ tracks ``data has arrived''; $E_i$ tracks ``consumers are done.'' |
| 144 | + The kernel allocates \texttt{num\_stages * 2} \texttt{Int64} entries. |
| 145 | +}; |
| 146 | + |
| 147 | +\node[note, anchor=north west, text width=6.0cm] (api) at (6.05,0.72) { |
| 148 | + \textbf{CuTe pipeline API mapping}\\[2pt] |
| 149 | + \texttt{acquire\_and\_advance()} waits $E_i$ and arms $F_i$ with expected bytes. |
| 150 | + Each \texttt{cute.copy(..., tma\_bar\_ptr=...)} reports completion to $F_i$; |
| 151 | + TMA \texttt{commit()} is therefore a no-op. |
| 152 | +}; |
| 153 | + |
| 154 | +\node[note, anchor=north west, text width=5.75cm] (phase) at (12.75,0.72) { |
| 155 | + \textbf{Index/phase prevents stale completion}\\[2pt] |
| 156 | + For $S=2$: producer $(0,1)\!\to\!(1,1)\!\to\!(0,0)$; consumer |
| 157 | + $(0,0)\!\to\!(1,0)\!\to\!(0,1)$. The phase flips whenever the ring wraps. |
| 158 | +}; |
| 159 | + |
| 160 | +% Footer formula and source mapping |
| 161 | +\node[subtitle, anchor=north west, text width=18.4cm] at (0,-1.75) { |
| 162 | + \textbf{Expected transaction bytes in this kernel:} |
| 163 | + $\texttt{tx\_count}=(\texttt{block\_n}+m)\,\texttt{tile\_k\_u32}\!\times\!4 |
| 164 | + + \text{optional staged weight-scale bytes}$. |
| 165 | + Source anchors: barrier allocation and pipeline creation in |
| 166 | + \texttt{transformer\_nuggets/cute/blockscaled\_tma.py}; producer load in |
| 167 | + \texttt{tma\_producer\_load\_stage}; consumer handoff in |
| 168 | + \texttt{compute\_warp\_consume\_stage}. |
| 169 | +}; |
| 170 | + |
| 171 | +\end{tikzpicture} |
| 172 | +\end{document} |
0 commit comments