Files
d4r/kernels/tools/swin_mma.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

30 lines
1.0 KiB
Python
Executable File

#!/usr/bin/env python3
"""Dump the mma statements of a swin entry's phase range with their A/B/C/D registers.
Usage: swin_mma.py LAYER LO_STMT [HI_STMT]
LAYER is a rr_layer_spec alias (enc0, ...). One line per mma, with the
statement index, the physical PTX line and the register operands in order.
"""
import os
import re
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "rr"))
import rr_layer_spec as R # noqa: E402
def main():
layer = sys.argv[1] if len(sys.argv) > 1 else "enc0"
lo = int(sys.argv[2]) if len(sys.argv) > 2 else 0
hi = int(sys.argv[3]) if len(sys.argv) > 3 else 10 ** 9
path, entry = R.resolve(layer, None, None)
lines = open(path).read().split("\n")
first, last = next((a, b) for (nm, a, b) in R.find_entries(lines) if nm == entry)
stmts, _ = R.linearise(lines, first, last)
for s in stmts:
if s.is_mma and lo <= s.idx <= hi:
print(f"s{s.idx} L{s.line}: " + " ".join(re.findall(r"%r\d+", s.text)))
if __name__ == "__main__":
main()