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).
This commit is contained in:
woutermans committed
1 parent
798e25611c
commit
0fe3a28cb7
58 files changed
+28330
-5
No files matched your search
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Splice recipe for `cuda_dldn_engine_swin_enc1_kernel` (corpus module 0014).
|
||||
|
||||
Self-contained: `kernels/tex/make_ptx.py` imports this module by path when the kernel
|
||||
name matches and uses what `recipe()` returns. Everything below was derived from the
|
||||
corpus PTX with `kernels/rr/rr_layer_spec.py enc1` and the register traces recorded in
|
||||
`kernels/rr/rrswin_enc1_epilogues.md`:
|
||||
|
||||
* the cut is the mma region, statements s8270-s26139, i.e. file lines 13343-66944
|
||||
(1072 `mma.sync` in 32 phases). `make_ptx.py` requires line 13343 to start an mma
|
||||
and line 66944 to end one, so CUT's second entry is 66945.
|
||||
* the live-in set is the two staged 96x32 e4m3 A operand sets the prologue built
|
||||
(%r11097.. for phases 1/2/8/9, %r11595.. for phases 3/4/10/11), phase 1's C seed
|
||||
register (%r15665, the constant zero of `mov.u32 %r15665, 0`), and the 64 packed
|
||||
f16x2 registers of the dequantised input tile (%r4678..%r4741) that E6 adds to
|
||||
phase 7's per-column table. The tile is passed in the order E6 consumes it:
|
||||
(m-tile 0..3, n-tile 0..7, row half 0..1).
|
||||
* the live-out set is phase 32's eight mma D pairs (16 registers). The two extents
|
||||
%r31837/%r31838 come from `ld.param.v2.u32 [param_0+0]` inside the cut, so the
|
||||
splice re-reads them instead of carrying them across.
|
||||
|
||||
ZLUDA lowers `call.uni`'s return parameters to the callee's LLVM return type and forms
|
||||
no parameter-space object for them (ptx/src/pass/llvm/emit.rs, `emit_call`), so the
|
||||
accumulators cross the cut through LDS instead. This module's entry declares no
|
||||
`.shared` array at all -- 0014 has none -- so the splice declares `d4r_scratch` itself:
|
||||
ZLUDA parses a `.shared` declaration in the middle of a body and sizes the workgroup's
|
||||
LDS from it (verified with d4r_emit). `kernels/rr/rrswin_enc1.hip` writes the sixteen
|
||||
accumulators at `base + 64*laneid + 4*i`; the splice reads them back with ld.shared.
|
||||
"""
|
||||
|
||||
ENTRY = "cuda_dldn_engine_swin_enc1_kernel"
|
||||
FILE = "0014-PREPASS_ENTRYPOINT_NAME.ptx"
|
||||
|
||||
# cut: the mma region, 1-based inclusive file lines
|
||||
CUT = (13343, 66945)
|
||||
|
||||
# A operand set 1: phases 1, 2, 8 and 9 (six 16-row fragments x four b32)
|
||||
AS1 = [
|
||||
"%r11097", "%r11098", "%r11099", "%r11100", "%r11137", "%r11138", "%r11139", "%r11140",
|
||||
"%r11177", "%r11178", "%r11179", "%r11180", "%r11217", "%r11218", "%r11219", "%r11220",
|
||||
"%r11257", "%r11258", "%r11259", "%r11260", "%r11297", "%r11298", "%r11299", "%r11300"
|
||||
]
|
||||
|
||||
# A operand set 2: phases 3, 4, 10 and 11 (the prologue's second, normalised packing)
|
||||
AS2 = [
|
||||
"%r11595", "%r11596", "%r11597", "%r11598", "%r11635", "%r11636", "%r11637", "%r11638",
|
||||
"%r11675", "%r11676", "%r11677", "%r11678", "%r11715", "%r11716", "%r11717", "%r11718",
|
||||
"%r11755", "%r11756", "%r11757", "%r11758", "%r11795", "%r11796", "%r11797", "%r11798"
|
||||
]
|
||||
|
||||
# phase 1's C seed: the constant-zero register %r15665
|
||||
CSEED = ["%r15665"]
|
||||
|
||||
# the dequantised input tile E6 adds to phase 7's column table, indexed
|
||||
# (m-tile 0..3, n-tile 0..7, row half 0..1)
|
||||
RES = [
|
||||
"%r4678", "%r4682", "%r4680", "%r4684", "%r4679", "%r4683", "%r4681", "%r4685",
|
||||
"%r4686", "%r4690", "%r4688", "%r4692", "%r4687", "%r4691", "%r4689", "%r4693",
|
||||
"%r4694", "%r4698", "%r4696", "%r4700", "%r4695", "%r4699", "%r4697", "%r4701",
|
||||
"%r4702", "%r4706", "%r4704", "%r4708", "%r4703", "%r4707", "%r4705", "%r4709",
|
||||
"%r4710", "%r4714", "%r4712", "%r4716", "%r4711", "%r4715", "%r4713", "%r4717",
|
||||
"%r4718", "%r4722", "%r4720", "%r4724", "%r4719", "%r4723", "%r4721", "%r4725",
|
||||
"%r4726", "%r4730", "%r4728", "%r4732", "%r4727", "%r4731", "%r4729", "%r4733",
|
||||
"%r4734", "%r4738", "%r4736", "%r4740", "%r4735", "%r4739", "%r4737", "%r4741"
|
||||
]
|
||||
|
||||
# live-out: phase 32's eight mma D pairs. %r31837/%r31838 (the extents) are re-read
|
||||
# from param_0 by the splice, so only the sixteen accumulators cross the cut.
|
||||
RET_ACCS = [
|
||||
"%r31908", "%r31909", "%r31918", "%r31919", "%r31948", "%r31949", "%r31958", "%r31959",
|
||||
"%r31988", "%r31989", "%r31998", "%r31999", "%r32028", "%r32029", "%r32038", "%r32039"
|
||||
]
|
||||
|
||||
NP = len(AS1) + len(AS2) + len(CSEED) + len(RES)
|
||||
B32 = AS1 + AS2 + CSEED + RES
|
||||
SCRATCH_WORDS = 2048 # 8 KiB of LDS: the phase 5/6 exchange, then the accumulators
|
||||
RET_STRIDE = 4 * len(RET_ACCS)
|
||||
|
||||
|
||||
def extern():
|
||||
"""`.extern .func` for the callee: the b32 live-ins, then the weight image, the
|
||||
parameter block and the LDS scratch address."""
|
||||
names = [".param .b32 d4r_a{}".format(i) for i in range(NP)]
|
||||
names += [".param .b64 d4r_w", ".param .b64 d4r_v", ".param .b32 d4r_s"]
|
||||
return ".extern .func d4r_swin_enc1\n(\n\t" + ", ".join(names) + "\n)\n;\n"
|
||||
|
||||
|
||||
def call():
|
||||
"""The splice that replaces the mma region: hand the live-ins to the native body,
|
||||
then read the accumulators back out of LDS."""
|
||||
stores = "\n".join("\tst.param.b32 [q%d+0], %s;" % (i, r) for i, r in enumerate(B32))
|
||||
hel = ", ".join(["q%d" % i for i in range(NP)] + ["qw", "qv", "qs"])
|
||||
gets = "\n".join("\tld.shared.b32 %s, [%%rsw1+%d];" % (RET_ACCS[i], 4 * i)
|
||||
for i in range(len(RET_ACCS)))
|
||||
return "\n".join([
|
||||
"{",
|
||||
".reg .b32 %rsw<3>;",
|
||||
".reg .b64 %rdw<2>;",
|
||||
".shared .b32 d4r_scratch[%d];" % SCRATCH_WORDS,
|
||||
"mov.u32 %rsw0, d4r_scratch;",
|
||||
".param .b32 q<%d>;" % NP,
|
||||
stores,
|
||||
".param .b64 qw;",
|
||||
"ld.param.u64 %%rdw0, [%s_param_0+40];" % ENTRY,
|
||||
"cvta.to.global.u64 %rdw0, %rdw0;",
|
||||
"st.param.b64 [qw+0], %rdw0;",
|
||||
".param .b64 qv;",
|
||||
"mov.b64 %%rdw1, %s_param_0;" % ENTRY,
|
||||
"st.param.b64 [qv+0], %rdw1;",
|
||||
".param .b32 qs;",
|
||||
"st.param.b32 [qs+0], %rsw0;",
|
||||
"call.uni d4r_swin_enc1, (%s);" % hel,
|
||||
"ld.param.v2.u32 {%%r31837, %%r31838}, [%s_param_0+0];" % ENTRY,
|
||||
"mad.lo.s32 %%rsw1, %%laneid, %d, %%rsw0;" % RET_STRIDE,
|
||||
gets,
|
||||
"}",
|
||||
])
|
||||
|
||||
|
||||
def recipe():
|
||||
return dict(file=FILE, cut_call=CUT, extern=extern(), call=call())
|
||||
Reference in new issue
Block a user