Skip to content

Commit 1fb5901

Browse files
Rewrite ch04 around 2D tensor contraction + cross-chapter content moves
- Ch04: New opening from general 2D tensor contraction concept, GPU tensor cores as one example, Croktile's abstract 4-step MMA syntax. Added comprehensive MMA variants section (layout, sparse, scale, load/swizzle, store/transpose, commit). Renamed "lifecycle" to "syntax". Three new Manim figures (tensor contraction, register loading, MMA syntax flow), SM86/SM90 renumbered to fig4. - Ch05: Moved parallel.async/stream s section from ch07 (host launch control fits better with control flow chapter). - Ch06: Added note about swap example's simplified kernel style. - Ch07: Removed parallel.async aside, added ch04 back-reference in swizzle matching rule. - All changes translated to ZH with consistent terminology. Made-with: Cursor
1 parent e53d83f commit 1fb5901

20 files changed

Lines changed: 585 additions & 141 deletions
36.8 KB
Loading
37.4 KB
Loading
28.9 KB
Loading
29.2 KB
Loading
35.1 KB
Loading
35.8 KB
Loading
147 KB
Loading
144 KB
Loading
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
"""
2+
Figure 1: 2D tensor contraction as a general concept.
3+
A[M,K] × B[K,N] → C[M,N] at tile level, with multiple hardware backends.
4+
"""
5+
import os, sys
6+
sys.path.insert(0, os.path.dirname(__file__))
7+
from manim import *
8+
from theme import parse_theme
9+
10+
C, THEME = parse_theme()
11+
12+
13+
class TensorContraction(Scene):
14+
def construct(self):
15+
self.camera.background_color = C["bg"]
16+
17+
title = Text(
18+
"2D Tensor Contraction: C += A × B",
19+
font_size=22, color=C["fg"], font="Monospace",
20+
)
21+
title.to_edge(UP, buff=0.3)
22+
self.add(title)
23+
24+
def make_tile(label, w, h, color, pos):
25+
r = Rectangle(width=w, height=h, fill_color=color,
26+
fill_opacity=0.2, stroke_color=color, stroke_width=2)
27+
t = Text(label, font_size=14, color=color, font="Monospace")
28+
g = VGroup(r, t)
29+
t.move_to(r.get_center())
30+
g.move_to(pos)
31+
return g
32+
33+
y_top = 0.8
34+
a = make_tile("A [M, K]", 1.4, 1.8, C["lhs_c"], LEFT * 3.5 + UP * y_top)
35+
b = make_tile("B [K, N]", 1.8, 1.4, C["rhs_c"], LEFT * 0.3 + UP * y_top)
36+
c_tile = make_tile("C [M, N]", 1.8, 1.8, C["out_c"], RIGHT * 3.0 + UP * y_top)
37+
38+
times = Text("×", font_size=28, color=C["fg"], font="Monospace")
39+
times.move_to((a.get_right() + b.get_left()) / 2)
40+
arrow = Arrow(b.get_right() + RIGHT * 0.15, c_tile.get_left() + LEFT * 0.15,
41+
buff=0, stroke_width=2.5, color=C["arrow_c"],
42+
max_tip_length_to_length_ratio=0.15)
43+
eq_label = Text("+=", font_size=16, color=C["fg2"], font="Monospace")
44+
eq_label.next_to(arrow, UP, buff=0.08)
45+
46+
self.add(a, b, c_tile, times, arrow, eq_label)
47+
48+
dim_a = Text("M rows, K cols", font_size=10, color=C["dim"], font="Monospace")
49+
dim_a.next_to(a, DOWN, buff=0.12)
50+
dim_b = Text("K rows, N cols", font_size=10, color=C["dim"], font="Monospace")
51+
dim_b.next_to(b, DOWN, buff=0.12)
52+
dim_c = Text("M rows, N cols", font_size=10, color=C["dim"], font="Monospace")
53+
dim_c.next_to(c_tile, DOWN, buff=0.12)
54+
self.add(dim_a, dim_b, dim_c)
55+
56+
sub_title = Text(
57+
"Same operation, different hardware implementations",
58+
font_size=14, color=C["fg3"], font="Monospace",
59+
)
60+
sub_title.move_to(DOWN * 0.8)
61+
self.add(sub_title)
62+
63+
backends = [
64+
("GPU Tensor Core", "16×16×16 FP16", C["green"]),
65+
("TPU MXU", "128×128 systolic", C["blue"]),
66+
("Intel AMX", "16×64 tiles", C["orange"]),
67+
("Custom DSA", "vendor-specific", C["purple"]),
68+
]
69+
70+
x_start = -4.5
71+
x_step = 3.0
72+
y_hw = -2.0
73+
74+
for i, (name, detail, color) in enumerate(backends):
75+
x = x_start + i * x_step
76+
box = RoundedRectangle(
77+
corner_radius=0.1, width=2.6, height=1.1,
78+
fill_color=color, fill_opacity=0.12,
79+
stroke_color=color, stroke_width=1.5,
80+
)
81+
box.move_to([x, y_hw, 0])
82+
n = Text(name, font_size=12, color=color, font="Monospace")
83+
d = Text(detail, font_size=10, color=C["dim"], font="Monospace")
84+
VGroup(n, d).arrange(DOWN, buff=0.08).move_to(box)
85+
self.add(box, n, d)
86+
87+
foot = Text(
88+
"Croktile's mma syntax targets the abstract operation, not a specific chip.",
89+
font_size=11, color=C["fg3"], font="Monospace",
90+
)
91+
foot.to_edge(DOWN, buff=0.3)
92+
self.add(foot)
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
"""
2+
Figure 2: GPU tensor core register layout — fragmented lane ownership.
3+
Threads in a warp own scattered pieces of the tile; Croktile hides this.
4+
"""
5+
import os, sys
6+
sys.path.insert(0, os.path.dirname(__file__))
7+
from manim import *
8+
from theme import parse_theme
9+
10+
C, THEME = parse_theme()
11+
12+
13+
class RegisterLoading(Scene):
14+
def construct(self):
15+
self.camera.background_color = C["bg"]
16+
17+
title = Text(
18+
"Tensor Core Register Layout (simplified)",
19+
font_size=20, color=C["fg"], font="Monospace",
20+
)
21+
title.to_edge(UP, buff=0.3)
22+
self.add(title)
23+
24+
sub = Text(
25+
"Each thread in a warp owns scattered register fragments of the tile",
26+
font_size=12, color=C["fg3"], font="Monospace",
27+
)
28+
sub.next_to(title, DOWN, buff=0.15)
29+
self.add(sub)
30+
31+
cell = 0.38
32+
rows, cols = 8, 8
33+
colors_per_thread = [
34+
C["blue"], C["orange"], C["green"], C["purple"],
35+
C["red"], C["teal"], C["yellow"], C["pink"],
36+
]
37+
ownership = [
38+
[0,0,1,1,2,2,3,3],
39+
[0,0,1,1,2,2,3,3],
40+
[4,4,5,5,6,6,7,7],
41+
[4,4,5,5,6,6,7,7],
42+
[0,0,1,1,2,2,3,3],
43+
[0,0,1,1,2,2,3,3],
44+
[4,4,5,5,6,6,7,7],
45+
[4,4,5,5,6,6,7,7],
46+
]
47+
48+
grid_center = LEFT * 2.0 + DOWN * 0.6
49+
grid = VGroup()
50+
for r in range(rows):
51+
for c_idx in range(cols):
52+
sq = Square(side_length=cell,
53+
fill_color=colors_per_thread[ownership[r][c_idx]],
54+
fill_opacity=0.3,
55+
stroke_color=colors_per_thread[ownership[r][c_idx]],
56+
stroke_width=1.2)
57+
sq.move_to(grid_center + RIGHT * (c_idx - cols/2 + 0.5) * cell
58+
+ DOWN * (r - rows/2 + 0.5) * cell)
59+
grid.add(sq)
60+
self.add(grid)
61+
62+
tile_label = Text("8×8 MMA tile", font_size=12, color=C["fg2"], font="Monospace")
63+
tile_label.next_to(grid, DOWN, buff=0.15)
64+
self.add(tile_label)
65+
66+
legend_title = Text("Thread ownership:", font_size=12, color=C["fg"], font="Monospace")
67+
legend_title.move_to(RIGHT * 2.8 + UP * 0.3)
68+
self.add(legend_title)
69+
for i in range(8):
70+
sq = Square(side_length=0.2,
71+
fill_color=colors_per_thread[i], fill_opacity=0.4,
72+
stroke_color=colors_per_thread[i], stroke_width=1)
73+
label = Text(f"T{i*4}–T{i*4+3}", font_size=9, color=C["fg2"], font="Monospace")
74+
row = VGroup(sq, label).arrange(RIGHT, buff=0.1)
75+
row.move_to(RIGHT * 2.8 + DOWN * (0.0 + i * 0.28))
76+
self.add(row)
77+
78+
brace_l = Text("Croktile:", font_size=14, color=C["green"], font="Monospace")
79+
brace_r = Text("mma.load / mma.store hide this", font_size=12, color=C["fg2"], font="Monospace")
80+
note = VGroup(brace_l, brace_r).arrange(RIGHT, buff=0.15)
81+
note.to_edge(DOWN, buff=0.3)
82+
self.add(note)

0 commit comments

Comments
 (0)