wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
a40d2b4ec36500b4913244c3da669e2f2686be88
d4r/kernels/rr/descriptor_dec0.py
T
woutermans 0fe3a28cb7 kernels/rr: staged-image Swin layer engine prototypes 
Per-layer staged-image engine for the eleven translated denoiser Swin
layers: descriptors (verified offsets per phase), cut recipes that splice an
isolated native phase into NVIDIA's PTX, wave templates, epilogue notes and
GPU probes. The isolated phase-8 splice produces a byte-identical final
4K frame but is slower than the translated body (14.90 vs 6.59 ms) - a
correctness experiment, not a shipping improvement; docs/rr-engine.md
records the design and the next direct-converter win (B operands built
from transposed accumulators).
2026-10-12 12:47:53 +02:00

703 lines
35 KiB
Python
Raw Blame History

"""Machine-readable descriptor for `cuda_dldn_engine_swin_dec0_kernel`.
The WMMA engine's per-layer input. Every number here is read out of
`~/.cache/d4r-rr-corpus/0023-PREPASS_ENTRYPOINT_NAME.ptx` (entry lines 1018-47791)
with `kernels/rr/rr_layer_spec.py`'s parser, or measured on the GPU; the source of
each family of numbers is named in the comments and anything not established that
way carries an `"unverified"` key.
Conventions used throughout:
* `m_tiles` counts 16-row mma tiles, `n_tiles` counts 8-column n tiles, `k` is the
contraction depth of the phase. `mma` is the shape NVIDIA emits,
`mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16`.
* The engine's unit is `k32_e4m3` (one `wmma_f32_16x16x16_fp8_fp8` pair), so a
phase of depth K is `k // 16` WMMA steps per output tile, two per k32 slice.
* The weight image is arg040, 47376 bytes. A *tile* is 512 bytes: 32 lanes x 16
bytes, holding two n-tiles' B fragments. For a phase with `k_slices` k32 steps
the tile index of n-tile `nt`, step `ks` is `k_slices * (nt >> 1) + ks`, and
inside a tile the fragment for n-tile `nt` sits at bytes `0..7` when `nt` is
even and `8..15` when it is odd (both read by the same `ld.weak.global.ca.v4`
at `W + 16*laneid + base`). The column a lane owns inside a tile is
`8*(nt & 1) + 2*t + r` for register `r`, and the address bit assignment inside
the tile is
`bit0 = k&1, bit1 = (k>>4)&1, bit2 = column&1, bit3 = column>>3,
bits4-5 = (column>>1)&3, bits6-8 = (k>>1)&7`
for the interleaved byte->k map, or
`bit0-1 = k&3, bit2 = column&1, bit3 = column>>3, bits4-5 = (column>>1)&3,
bits6-8 = k>>2`
for the consecutive one; see `rrswin_dec0.hip` for both and for the A side.
`"unverified"`: which of the two byte maps the mma uses was not settled --
the first attempt to prove it was invalidated by an ambient-memory comparison
(see the "Validation" section of rrswin_dec0.hip), so both are carried.
* A *column table* (E1's gain, the per-column addends) is read with the table
address `W + ((laneid << 2) & 12) + off`: the four lanes with equal `laneid & 3`
share a 16-byte word, so one load covers 32 columns.
* Register names (`%rNNNN`) are cited so each field can be checked against the
PTX with `sed -n`.
"""
def layer():
"""Launch geometry, argument map and token geometry."""
return dict(
entry="cuda_dldn_engine_swin_dec0_kernel",
corpus="0023-PREPASS_ENTRYPOINT_NAME.ptx",
entry_lines=(1018, 47791),
role="first Swin decoder (encoders' outputs -> feature planes)",
# one warp per 8x8 token patch; blockDim 32x1x1 is a single wave, so no
# %tid.z term appears in any address
grid=(161, 93, 1),
block=(32, 1, 1),
wave=32,
blocks=14973,
token_patch=(8, 8),
param_bytes=152,
# param_0+80 is the block origin as a v2.u16 (s5, L1033); the token
# geometry is 1280x720 -> 161x93 patches of 8x8
image=(1280, 736),
args=dict(
input=8, # u64, e4m3 input plane; read once, 24 ld.global.u32
skip=24, # u64, e4m3 skip input; 24 ld.global.v2.u16
weights=40, # u64, prepared weight image, 47376 bytes
plane=48, # u64, output e4m3 surface (see writes below)
origin=80, # v2.u16 block origin
texture_a=88, # u64, 18 tex.base.2d.v4.f16.f32 samples
texture_b=96, # u64, tex.base.2d.v4.f16.s32 at (x, y)
texture_c=104, # u64, tex.base.2d.v4.f16.s32 at (x, y)
surface_a=112, # u64, sust.b.2d.v4.b16.zero destination
coord_scale=136, # v2.f32, texture coordinate scale/origin
surface_b=144, # u64, sust.p.2d.v4.b32.zero destination
),
# measured on the GPU: the launch dump files the 65,945,600-byte block
# that holds these pointers as arg008 at base+0x1cc0000, so a pointer at
# `block + N` is dumped at arg008[N + 0x1cc0000]
arena=dict(
dump_arg=8,
dump_offset=0x1CC0000,
dump_bytes=65945600,
input=30146560, # param+8 (Main's measurement, shared)
plane=45219840, # param+48 (measured here: the p1 oracle's
# payload at D4R_RRSWIN_DEBUG_BASE lands at
# arg008 offset 61,997,056, i.e. 45,219,840 +
# 16,777,216, and shifting the base by 8192
# shifts the first differing byte by 8192)
param56=60293120, # param+56 (Main's measurement, shared)
),
writes=(
# (statement, instruction, argument, address)
(18502, "st.global.v4.u16", 48, "4*(ex*y + x)"), # e4m3 planes
(18603, "sust.p.2d.v4.b32.zero", 144, "(x, y)"), # four f32
(18813, "sust.b.2d.v4.b16.zero", 112, "(8*x, y)"), # texture fuse
),
unverified="param_0+32, +56, +72, +120 and +128 are never read by the "
"entry; the 4-byte value of param+0 is the (ex, ey) extents "
"pair (s13), and the offset of param+48 inside its block is "
"the only arena figure measured for this layer.",
)
def phases():
"""The 19 GEMM phases in program order, each with its epilogue."""
return [p1(), p2(), p3(), p4(), p5(), p6(), p7(), p8(), p9(), p10(),
p11(), p12(), p13(), p14(), p15(), p16(), p17(), p18(), p19()]
# ---------------------------------------------------------------------------
# shared pieces
# ---------------------------------------------------------------------------
def _e1():
"""E1: patch expand, skip residual, RMS norm. Returns the normalised tile."""
return dict(
name="E1",
stmts=(422, 5936),
steps=[
dict(op="requantise", detail="D1 (96 f16x2 registers, %r464..%r775 "
"written by the 96 mma) -> 37 `mov.b32` of two "
"cvt.rn.satfinite.e4m3x2.f16x2 each plus 3 v2 pairs; 4 e4m3 "
"codes per b32, the two codes of two adjacent n-tiles for one "
"row"),
dict(op="stage", detail="st.shared into the entry's 3200-byte smem as "
"two 400-cell planes of a 10x10 grid of 4-byte cells, cell = "
"row*40 + col*4 + t, plane stride 1600 bytes; 25 cells per lane "
"are written and all 800 cells are covered exactly once"),
dict(op="gather", detail="24 ld.shared.v2.u16 per lane at "
"cell = plane*400 + row*40 + col*4 + t with "
"col = (laneid>>2) + 1, for row = 1..8, 0, 9 in both planes, "
"then (row = col, col = 0) and (row = col, col = 9) in both "
"planes -- 24 v2 loads = 48 f16x2, one b16 per e4m3x2"),
dict(op="dequantise_skip", detail="24 ld.global.v2.u16 from param+24 "
"(s1613-s4384) then 96 cvt.rn.f16x2.e4m3x2"),
dict(op="add", detail="204 add.f16: staged patch + dequantised skip"),
dict(op="rms_norm", detail="mul.f16 squares, 60 add.f16x2 in-lane "
"sums, 24 shfl.sync.bfly.b32, + 2^-13 (mov.f64 0x3F20000000000000, "
"s5030), rsqrt.approx.ftz.f32, times the column gain read at "
"W + ((laneid<<2)&12) + {8448, 8464, 8480, 8496} (4 x "
"ld.global.v2.u16 = 8 f16 per lane = 32 columns). No mean "
"subtraction and no 1/N"),
dict(op="requantise", detail="48 cvt.rn.satfinite.e4m3x2.f16x2 "
"(s5860-s5928) -> 24 b32 = phases 2 and 3's A operand"),
],
# the two registers E5 still reads after the softmax: the skip residual
# sums %rs425 / %rs428 (s4484/s4485), defined once here and read again at
# s10118 as phase 6's C seed
residual_registers=("%rs425", "%rs428"),
unverified="the lane -> (row, col) map of the staging and of the "
"gather was derived by executing the region s539-s1457 as "
"integer PTX for all 32 lanes; the resulting map is "
"self-consistent (exactly 800 distinct cells, one writer "
"each, every read cell written) but has not been checked "
"against NVIDIA's own values.",
)
def _softmax(scale=0.017181396484375):
"""The cubic-exponent softmax E4 and E9 use (identical code, shifted regs)."""
return dict(
clamp=(-0.55615234375, 0.55615234375),
scale=scale,
cube=dict(c0=0.92724609375, note="u = f16(1.375 + m*(c0 - m*m))"),
trick="expval = f16 from bits ((u_bits << 5) & 0x7FE07FE0), i.e. "
"2^(32*frac(u) - 47)",
row_sum="8 row sums (one per row g, g+8 of the four m-tiles) from 12 "
"add.f16x2 partials and 16 shfl.sync.bfly.b32; there is NO row "
"maximum in this epilogue, the max/min f16x2 are the clamp",
normalise="192 rcp.approx.ftz.f32 -> 96 reciprocals, then 192 mul.f16",
)
def _activation():
"""The clamped cubic activation of E12, E14, E16 and E18."""
return dict(
clamp=(-2.0, 2.0),
gain="g = f16(0.5 + y*(0.412109375 - 0.0810546875*|y|)) with "
"y = clamp(x, -2, 2), all in f16",
out="x * g, computed in f32 and rounded to f16",
note="a cubic ramp, not sigmoid/GELU: g(+2) = 1 so out = x above 2, "
"g(-2) = 0 so out = 0 below -2",
requantise="32 cvt.rn.satfinite.e4m3x2.f16x2 -> 16 b32 = the next "
"phase's A",
)
# ---------------------------------------------------------------------------
# the phases
# ---------------------------------------------------------------------------
def p1():
return dict(
index=1,
stmts=(326, 421), lines=(1393, 2061),
gemm=dict(m_tiles=3, n_tiles=16, k=64, k_slices=2, mma=96,
mma_shape="m16n8k32", wmma_steps=4,
note="the tool prints 96x256 K=64; the 48 chains of two steps "
"reproduce 3*16*2 = 96, so N is 128 and the 32 B "
"fragments are 16 n-tiles x 2 k slices"),
mma_order="for mi in 0..2, for p in 0..7, for ks in 0..1, for j in 0..1: "
"a_fragment = 2*mi + ks, n_tile = 2*p + j, b_fragment = "
"4*p + 2*ks + j",
a=dict(source="arg+8 input plane, staged by NVIDIA's prologue into 24 "
"registers (%r756..%r759, %r776..%r779, %r1076..%r1079, "
"%r1096..%r1099, %r1396..%r1399, %r1416..%r1419) in mma first-use "
"order", fragments=6, registers=24, requantise="none: already e4m3"),
b=dict(source="weight image", base=0,
tile="k_slices*(nt>>1) + ks = 2*(nt>>1) + ks",
fragments=32, note="16 n-tiles x 2 k slices; the largest word "
"offset is 7688, so the phase reads [0, 7692)"),
c=dict(kind="per-n-tile f16x2 bias word", count=16, base=8192, stride=16,
addr="W + ((laneid<<2)&12) + 8192 + 16*nt",
regs=["%r1113", "%r1123", "%r1153", "%r1163", "%r1193", "%r1203",
"%r1233", "%r1243", "%r1273", "%r1283", "%r1313", "%r1323",
"%r1353", "%r1363", "%r1393", "%r1403"],
note="the first k step of each chain; the second reads the first "
"step's f16 result"),
d_registers=96,
epilogue=_e1(),
)
def _proj(index, stmts, lines, a_source, a_registers, b_base, c, m_tiles,
n_tiles, epilogue, note=""):
"""One K=32 projection: the shape of phases 2, 3, 6, 7, 8 and 11..19."""
return dict(
index=index, stmts=stmts, lines=lines,
gemm=dict(m_tiles=m_tiles, n_tiles=n_tiles, k=32, k_slices=1,
mma=m_tiles * n_tiles, mma_shape="m16n8k32", wmma_steps=2,
note=note),
mma_order="for mi in 0..%d, for nt in 0..%d: a_fragment = mi, "
"n_tile = nt" % (m_tiles - 1, n_tiles - 1),
a=a_source,
b=dict(source="weight image", base=b_base, tile="(nt>>1)",
fragments=n_tiles,
span="[%d, %d)" % (b_base, b_base + 512 * ((n_tiles + 1) // 2))),
c=c,
d_registers=2 * m_tiles * n_tiles,
epilogue=epilogue,
)
def _zero_c():
return dict(kind="constant zero", regs=["%r13609"],
note="mov.u32 %r13609, 0 at L6552; the mma's C field is "
"{%r13609, %r13609}")
def _norm_tile_a(regs):
return dict(source="E1's normalised tile (the 24 packed b32 E1's last "
"requantise writes)", fragments=6, registers=24, regs=regs,
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E1")
def _bias_c(base, count, stride=16):
return dict(kind="per-n-tile f16x2 bias word", count=count, base=base,
stride=stride,
addr="W + ((laneid<<2)&12) + %d + %d*nt" % (base, stride))
def _acc_c(src, regs, n_tiles, m_tiles):
return dict(kind="the previous phase's f16x2 accumulators",
count=2 * m_tiles * n_tiles, regs=regs, note=src)
def p2():
return _proj(2, (5937, 5960), (10458, 10622),
_norm_tile_a(["%r9539", "%r9579", "%r9619", "%r9659", "%r9699",
"%r9739"]),
None, 8512, _zero_c(), 6, 4,
dict(name="E2", stmts=(5961, 5967),
steps=[dict(op="load", detail="phase 3's B, two "
"ld.weak.global.ca.v4 at W + 16*laneid + "
"{9536, 10048}")]))
def p3():
return _proj(3, (5968, 5991), (10639, 10803),
_norm_tile_a(["%r9539", "%r9579", "%r9619", "%r9659", "%r9699",
"%r9739"]),
None, 9536, _zero_c(), 6, 4,
dict(name="E3", stmts=(5992, 6162),
steps=[
dict(op="load", detail="phase 4's C bias: 24 "
"ld.weak.global.ca.v4 at "
"W + 16*laneid + 10560 + 512*i, i = 0..23 "
"= [10560, 22348)"),
dict(op="requantise", detail="phase 3's D (%r4904.."
"%r5135) -> 80 cvt.rn.satfinite.e4m3x2 in 40 "
"b32 = phase 4's A (16 b32) and B (24 b32)"),
]))
def p4():
return dict(
index=4, stmts=(6163, 6210), lines=(11188, 11520),
gemm=dict(m_tiles=4, n_tiles=12, k=32, k_slices=1, mma=48,
mma_shape="m16n8k32", wmma_steps=2,
note="the 64x96 score block of the first attention window"),
mma_order="for mi in 0..3, for nt in 0..11: a_fragment = mi, "
"n_tile = nt",
a=dict(source="E3's pack of phase 3's D", fragments=4, registers=16,
regs=["%r5433", "%r5553", "%r5673", "%r5793"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E3"),
b=dict(source="E3's pack of phase 2's D (the V of the first window)",
fragments=12, registers=24,
regs=["%r5687", "%r5697", "%r5707", "%r5717", "%r5727", "%r5737",
"%r5747", "%r5757", "%r5767", "%r5777", "%r5787", "%r5797"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E3"),
c=dict(kind="per-(m-tile, n-tile) f16x2 seed pair", count=48,
base=10560, span=(10560, 22348),
word_offset="base + 2048*(i >> 1) + 8*(i & 1) for mma index i, "
"the pair at +0 and +4",
addr="W + ((laneid<<2)&12) + word_offset",
note="the 12288-byte position bias of the first window; four "
"consecutive mma share a 512-byte stride"),
d_registers=96,
epilogue=dict(
name="E4", stmts=(6211, 10053), softmax=_softmax(),
steps=[
dict(op="exponent", detail="96 element groups: "
"m = clamp(f16(x * 0.017181396484375), -0.55615234375, "
"+0.55615234375); u = f16(1.375 + m*(0.92724609375 - m*m)); "
"expval = f16 from bits ((u_bits << 5) & 0x7FE07FE0)"),
dict(op="row_sum", detail="12 add.f16x2 partials per lane, 16 "
"shfl.sync.bfly.b32, 8 row sums; no row maximum"),
dict(op="normalise", detail="192 rcp.approx.ftz.f32 -> 96 "
"reciprocals; 192 mul.f16 -> the 96 probabilities"),
dict(op="transpose", detail="48 movmatrix.sync.trans.aligned."
"m8n8.b16 over phase 2's D (the V matrix)"),
dict(op="requantise", detail="144 cvt.rn.satfinite.e4m3x2 in 48 "
"b32 = phase 5's A (48 b32, 4 m-tiles x 3 k slices) and B "
"(24 b32, 4 n-tiles x 3 k slices)"),
]),
)
def p5():
return dict(
index=5, stmts=(10054, 10101), lines=(21667, 21999),
gemm=dict(m_tiles=4, n_tiles=4, k=96, k_slices=3, mma=48,
mma_shape="m16n8k32", wmma_steps=6,
note="16 chains of three steps; the tool prints 192x96 "
"because it attributes the k slices to M and N"),
mma_order="for mi in 0..3, for nt in 0..3, for ks in 0..2: "
"a_fragment = 3*mi + ks, b_fragment = 3*nt + ks",
a=dict(source="E4's pack of the softmax probabilities (P of the first "
"window)", fragments=12, registers=48,
regs=["%r8639", "%r8659", "%r8679", "%r8759", "%r8779", "%r8799",
"%r8879", "%r8899", "%r8919", "%r8999", "%r9019", "%r9039"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E4"),
b=dict(source="E4's pack of the transposed phase 2's D (the V matrix)",
fragments=12, registers=24,
regs=["%r8933", "%r8943", "%r8953", "%r8963", "%r8973", "%r8983",
"%r8993", "%r9003", "%r9013", "%r9023", "%r9033", "%r9043"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E4"),
c=dict(kind="zero then the chain's own accumulator", count=1,
regs=["%r13609"],
note="ks = 0 reads {%r13609, %r13609}; ks > 0 reads the previous "
"step's D pair"),
d_registers=96,
epilogue=dict(
name="E5", stmts=(10102, 10261),
steps=[
dict(op="load", detail="phase 6's B, two v4 at "
"W + 16*laneid + {22848, 23360}"),
dict(op="table", detail="phase 6's column table: four "
"ld.global.v2.u16 at W + ((laneid<<2)&12) + "
"{39232, 39248, 39264, 39280} = 8 f16 = 32 columns"),
dict(op="add", detail="64 add.f16: table[column] + the E1 "
"residual %rs425 / %rs428 (defined at s4484/s4485, the only "
"definition in the entry), giving phase 6's C operand"),
dict(op="requantise", detail="32 cvt.rn.satfinite.e4m3x2 of "
"phase 5's D -> 16 b32 = phase 6's A"),
]),
)
def p6():
return _proj(
6, (10262, 10277), (22427, 22535),
dict(source="E5's pack of phase 5's D (P*V of the first window)",
fragments=4, registers=16, regs=["%r9121", "%r9161", "%r9201",
"%r9241"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E5"),
None, 22848,
_acc_c("E1's skip residual plus the per-column table at W + 39232",
["%r9097", "%r9107", "%r9117", "%r9127", "%r9137", "%r9147",
"%r9157", "%r9167", "%r9177", "%r9187", "%r9197", "%r9207",
"%r9217", "%r9227", "%r9237", "%r9247"], 4, 4),
4, 4,
dict(name="E6", stmts=(10278, 10284),
steps=[dict(op="load", detail="phase 7's B at "
"W + 16*laneid + {23872, 24384}")]))
def p7():
return _proj(
7, (10285, 10308), (22552, 22716),
_norm_tile_a(["%r9539", "%r9579", "%r9619", "%r9659", "%r9699", "%r9739"]),
None, 23872, _zero_c(), 6, 4,
dict(name="E7", stmts=(10309, 10315),
steps=[dict(op="load", detail="phase 8's B at "
"W + 16*laneid + {24896, 25408}")]))
def p8():
return _proj(
8, (10316, 10339), (22733, 22897),
_norm_tile_a(["%r9539", "%r9579", "%r9619", "%r9659", "%r9699", "%r9739"]),
None, 24896, _zero_c(), 6, 4,
dict(name="E8", stmts=(10340, 10510),
steps=[
dict(op="load", detail="phase 9's C bias: 24 v4 at "
"W + 16*laneid + 25920 + 512*i, i = 0..23 = [25920, 37708)"),
dict(op="requantise", detail="phase 8's D (%r9507..%r9737) -> "
"80 cvt.rn.satfinite.e4m3x2 in 40 b32 = phase 9's A (16) "
"and B (24)"),
]))
def p9():
out = dict(p4())
out.update(index=9, stmts=(10511, 10558), lines=(23282, 23614))
out["gemm"] = dict(p4()["gemm"], note="the 64x96 score block of the second "
"attention window")
out["a"] = dict(p4()["a"], source="E8's pack of phase 8's D",
regs=["%r10036", "%r10156", "%r11036", "%r11156"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E8")
out["b"] = dict(p4()["b"], source="E8's pack of phase 7's D (the V of the "
"second window)",
regs=["%r11287", "%r11397", "%r11507", "%r11617", "%r11727",
"%r11837", "%r11947", "%r12057", "%r12167", "%r12277",
"%r12387", "%r12497"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E8")
out["c"] = dict(p4()["c"], base=25920, span=(25920, 37708),
note="the second window's 12288-byte position bias")
out["epilogue"] = dict(name="E9", stmts=(10559, 14396), softmax=_softmax(),
steps=[
dict(op="exponent", detail="identical code to "
"E4's, registers shifted; 96 element groups"),
dict(op="row_sum", detail="12 add.f16x2 partials, "
"16 shfl.sync.bfly.b32, 8 row sums"),
dict(op="normalise", detail="192 rcp.approx."
"ftz.f32 then 192 mul.f16"),
dict(op="transpose", detail="48 movmatrix.sync."
"trans.aligned.m8n8.b16 over phase 7's D"),
dict(op="requantise", detail="144 "
"cvt.rn.satfinite.e4m3x2 in 48 b32 = "
"phase 10's A (48) and B (24)"),
])
return out
def p10():
out = dict(p5())
out.update(index=10, stmts=(14397, 14444), lines=(33756, 34088))
out["gemm"] = dict(p5()["gemm"], note="P*V of the second attention window")
out["a"] = dict(p5()["a"], source="E9's pack of the second window's "
"probabilities",
regs=["%r13242", "%r13262", "%r13282", "%r13302", "%r13322",
"%r13342", "%r13362", "%r13382", "%r13402", "%r13422",
"%r13442", "%r13462"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E9")
out["b"] = dict(p5()["b"], source="E9's pack of the transposed phase 7's D",
regs=["%r13933", "%r13943", "%r13953", "%r13963", "%r13973",
"%r13983", "%r13993", "%r14003", "%r14013", "%r14023",
"%r14033", "%r14043"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E9")
out["epilogue"] = dict(name="E10", stmts=(14445, 14499), steps=[
dict(op="load", detail="phase 11's B at W + 16*laneid + {38208, 38720}"),
dict(op="requantise", detail="32 cvt.rn.satfinite.e4m3x2 of phase 10's "
"D -> 16 b32 = phase 11's A; no column table"),
], note="phase 11's C is phase 6's accumulator (the residual), so phase 11 "
"computes D11 = A11*B11 + D6")
return out
def p11():
return _proj(
11, (14500, 14515), (34217, 34325),
dict(source="E10's pack of phase 10's D", fragments=4, registers=16,
regs=["%r13723", "%r13763", "%r13803", "%r13843"],
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in E10"),
None, 38208,
_acc_c("phase 6's accumulators (the running residual)",
["%r9089", "%r9099", "%r9109", "%r9119", "%r9129", "%r9139",
"%r9149", "%r9159", "%r9169", "%r9179", "%r9189", "%r9199",
"%r9209", "%r9219", "%r9229", "%r9239"], 4, 4),
4, 4,
dict(name="E11", stmts=(14516, 15402), op="rms_norm",
steps=[
dict(op="squares", detail="192 mul.f16 of phase 11's D"),
dict(op="sum", detail="16 add.f16x2 in-lane, 16 "
"shfl.sync.bfly.b32 + add.f16 across the four lanes"),
dict(op="scale", detail="+ 2^-13 then rsqrt.approx.ftz.f32, 64 "
"cvt.f32.f16 / rsqrt / cvt.rn.f16.f32, rounded to f16"),
dict(op="gain", detail="times the column gain at "
"W + ((laneid<<2)&12) + {39296, 39312, 39328, 39344}; "
"192 mul.f16 (inv*gain, then x*(inv*gain))"),
dict(op="load", detail="phase 12's B at W + 16*laneid + "
"{39360, 39872} and its C bias at "
"W + ((laneid<<2)&12) + {40384, 40400, 40416, 40432}"),
dict(op="requantise", detail="32 cvt.rn.satfinite.e4m3x2 -> "
"16 b32 = phase 12's A"),
],
note="no mean subtraction and no 1/N; the normalised tile N (32 "
"f16x2) is also what E13, E15 and E17 repack, so phases 12, "
"14, 16 and 18 all take the same A operand"))
def _mlp_a(source, regs, note=""):
return dict(source=source, fragments=4, registers=16, regs=regs,
requantise="cvt.rn.satfinite.e4m3x2.f16x2 in the preceding "
"epilogue", note=note)
N_REGS = ["%r15544", "%r15584", "%r15624", "%r15664"] # E13's pack of N
ACT14 = ["%r16674", "%r16714", "%r16754", "%r16794"] # E14's pack
ACT16 = ["%r16876", "%r16916", "%r16956", "%r16996"] # E15's pack of N
ACT17 = ["%r18006", "%r18046", "%r18086", "%r18126"] # E16's pack
ACT18 = ["%r18208", "%r18248", "%r18288", "%r18328"] # E17's pack of N
ACT19 = ["%r19338", "%r19378", "%r19418", "%r19458"] # E18's pack
def p12():
return _proj(
12, (15403, 15418), (36471, 36579),
_mlp_a("E11's pack of the RMS-normed tile N",
["%r14211", "%r14251", "%r14291", "%r14331"]),
None, 39360, _bias_c(40384, 4), 4, 4,
dict(name="E12", stmts=(15419, 16127), activation=_activation(),
steps=[
dict(op="activate", detail="over each of phase 12's 32 f16x2 D "
"registers: y = clamp(x, -2, 2), "
"g = 0.5 + y*(0.412109375 - 0.0810546875*|y|), out = x*g, "
"with 128 cvt.rn.f16.f32, 32 each neg/max/min/abs.f16x2, "
"96 mul.f16x2, 32 sub.f16x2, 32 add.f16x2"),
dict(op="table", detail="64 add.f16 of the column table read at "
"W + ((laneid<<2)&12) + {41472, 41488, 41504, 41520}, "
"written over the *activation inputs*, so phase 13's "
"activation input is D12 + table[column]"),
dict(op="load", detail="phase 13's B at W + 16*laneid + "
"{40448, 40960}"),
dict(op="requantise", detail="32 packs of the activation "
"outputs -> 16 b32 = phase 13's A"),
],
note="the only one of the four branches with a per-column table"))
def p13():
return _proj(
13, (16128, 16143), (38934, 39042),
_mlp_a("E12's pack of the activated (and table-added) tile",
["%r15342", "%r15382", "%r15422", "%r15462"]),
None, 40448,
_acc_c("phase 12's accumulators plus the per-column table at W + 41472, "
"i.e. the values E12 wrote over the activation inputs",
["%r15269", "%r15329", "%r15389", "%r15449", "%r15509", "%r15569",
"%r15629", "%r15689", "%r15749", "%r15809", "%r15869", "%r15929",
"%r15989", "%r16049", "%r16109", "%r16169"], 4, 4),
4, 4,
dict(name="E13", stmts=(16144, 16207), steps=[
dict(op="load", detail="phase 14's B at W + 16*laneid + {41536, "
"42048} and its C bias at W + ((laneid<<2)&12) + "
"{42560, 42576, 42592, 42608}"),
dict(op="requantise", detail="32 packs of the *normed tile* N (not "
"of phase 13's D) -> 16 b32 = phase 14's A"),
]))
def p14():
return _proj(
14, (16208, 16223), (39182, 39290),
_mlp_a("E13's pack of the RMS-normed tile N", N_REGS,
note="the same 16 registers as phase 16's A"),
None, 41536, _bias_c(42560, 4), 4, 4,
dict(name="E14", stmts=(16224, 16823), activation=_activation(),
steps=[
dict(op="activate", detail="the same clamped cubic as E12, over "
"phase 14's own D; no ld.global and no table add"),
dict(op="load", detail="phase 15's B at W + 16*laneid + "
"{42624, 43136}"),
dict(op="requantise", detail="32 packs of the activation "
"outputs -> 16 b32 = phase 15's A"),
]))
def p15():
return _proj(
15, (16824, 16839), (41342, 41450),
_mlp_a("E14's pack of the activated tile", ACT14), None, 42624,
_bias_c(44672, 4), 4, 4,
dict(name="E15", stmts=(16840, 16903), steps=[
dict(op="load", detail="phase 16's B at W + 16*laneid + {43648, "
"44160} and its C bias at W + ((laneid<<2)&12) + "
"{44672, 44688, 44704, 44720}"),
dict(op="requantise", detail="32 packs of the normed tile N -> "
"16 b32 = phase 16's A"),
]))
def p16():
return _proj(
16, (16904, 16919), (41590, 41698),
_mlp_a("E15's pack of the RMS-normed tile N", ACT16), None, 43648,
_bias_c(46784, 4), 4, 4,
dict(name="E16", stmts=(16920, 17519), activation=_activation(),
steps=[
dict(op="activate", detail="the same clamped cubic as E12, over "
"phase 16's own D"),
dict(op="load", detail="phase 17's B at W + 16*laneid + "
"{44736, 45248}"),
dict(op="requantise", detail="32 packs of the activation "
"outputs -> 16 b32 = phase 17's A"),
]))
def p17():
return _proj(
17, (17520, 17535), (43750, 43858),
_mlp_a("E16's pack of the activated tile", ACT17), None, 44736,
_bias_c(42560, 4) if False else dict(
kind="per-n-tile f16x2 bias word", count=4, base=46784, stride=16,
addr="W + ((laneid<<2)&12) + 46784 + 16*nt",
note="phase 17's C is read in E17 as "
"W + ((laneid<<2)&12) + {46784, 46800, 46816, 46832}"),
4, 4,
dict(name="E17", stmts=(17536, 17599), steps=[
dict(op="load", detail="phase 18's B at W + 16*laneid + {45760, "
"46272} and its C bias at W + ((laneid<<2)&12) + "
"{46784, 46800, 46816, 46832}"),
dict(op="requantise", detail="32 packs of the normed tile N -> "
"16 b32 = phase 18's A"),
]))
def p18():
return _proj(
18, (17600, 17615), (43998, 44106),
_mlp_a("E17's pack of the RMS-normed tile N", ACT18), None, 45760,
_bias_c(46784, 4), 4, 4,
dict(name="E18", stmts=(17616, 18215), activation=_activation(),
steps=[
dict(op="activate", detail="the same clamped cubic as E12, over "
"phase 18's own D; this is the last requantisation in "
"the entry"),
dict(op="load", detail="phase 19's B at W + 16*laneid + "
"{46848, 46960}"),
dict(op="requantise", detail="32 packs of the activation "
"outputs -> 16 b32 = phase 19's A"),
]))
def p19():
return _proj(
19, (18216, 18231), (46158, 46266),
_mlp_a("E18's pack of the activated tile", ACT19), None, 46848,
_acc_c("phase 18's accumulators",
["%r19089", "%r19149", "%r19209", "%r19269", "%r19329", "%r19389",
"%r19449", "%r19509"], 4, 2),
4, 4,
dict(name="E19", stmts=(18232, 19160),
note="the final epilogue stays NVIDIA's in the mirror cut; it is "
"described here because the engine's output feeds it",
steps=[
dict(op="row_max", detail="a real row maximum (unlike E4's): a "
"selp chain over laneid&3 picks one of the four D pairs "
"per lane, then four shfl.sync.idx.b32 with "
"srcLane = (laneid&7) | ((((laneid>>3)+d)&3)<<2) for "
"d = 0,1,2,-1, then a selp chain over laneid>>3; max.f16x2 "
"/ max.f16 and sub.f16x2 subtract it"),
dict(op="bounds", detail="ld.param.v2.u32 at param+0 gives "
"(ex, ey); out-of-range lanes branch to $L__BB1_367 and "
"skip the rest"),
dict(op="surface1", detail="4 cvt.rn.satfinite.e4m3x2 and one "
"st.global.v4.u16 at 4*(ex*y + x) into param+48"),
dict(op="surface2", detail="per element v: a "
"1 - 2/(1 + 2^(|v|*log2e)) branch (with a hard 1.0 above "
"9.0109...) and an fma polynomial over v^2 selected by "
"|v| >= 0.6, then sust.p.2d.v4.b32.zero at (x, y) into "
"param+144"),
dict(op="softmax", detail="x - max in f16, 2^((x-max)*log2e) "
"in f32, an f32 row sum, div.approx.ftz.f32 by 1.0, five "
"mul.f16x2"),
dict(op="texture_fuse", detail="18 tex.base.2d.v4.f16.f32 of "
"param+88 at f32 coordinates built from param+136, plus "
"tex.base.2d.v4.f16.s32 of param+96 and param+104 at "
"integer (x, y); folded with fma.rn.f16, guarded by "
"set.nan.f16.f16 / setp.ne.s16 so a NaN sample is "
"skipped, written by sust.b.2d.v4.b16.zero at (8*x, y) "
"into param+112 with a constant 0 fourth channel"),
],
unverified="E19's two evaluation branches are exact in the PTX "
"(every instruction is cited) but which model function "
"they implement is not determinable from it."),
)
def _final():
return None
Reference in new issue View Git Blame Copy Permalink