path: root/dmcub/hw/dmcub_bringup.py

#!/usr/bin/env python3
"""Bring up the DMCUB on the tinygrad-AM rig by backdoor load (no PSP).

Port of `dmub_srv_hw_init` + `dmub_dcn32_{reset,backdoor_load,setup_windows,
setup_mailbox,reset_release}` (vendored in ../ref/) onto this project's device
access. The stock signed image is used unmodified -- the point is to reproduce a
bug in AMD's firmware, not to run our own code on the DMCUB.

Why backdoor rather than PSP: `dmub_dcn32_backdoor_load` programs CW0/CW1 itself
and clears `DMCUB_SEC_CNTL.DMCUB_SEC_RESET` with `DMCUB_MEM_UNIT_ID=0x20`, i.e.
the instruction window is fetched straight from framebuffer with no signature
check in our path. That removes the whole PSP `GFX_FW_TYPE_DMUB` dance. If this
part is fused to require a PSP load, phase `release` will simply never report
boot-complete -- which is itself a result worth recording.

**DRY RUN BY DEFAULT.** Nothing is written without `--apply`. The dry run prints
every register write, in order, with resolved addresses, so the sequence can be
reviewed before it ever touches the card.

Order (each phase can be run alone with --phase):
  plan      compute windows + VRAM allocations, print the map          (no writes)
  reset     halt the core, disable CW2..7, zero ring pointers
  load      copy the image body to the CW0 buffer, program CW0/CW1, clear SEC_RESET
  windows   program CW3..CW6 + REGION5
  mailbox   program INBOX1 base/size, zero RPTR/WPTR
  release   MMHUBBUB DMUIF out of reset, DMCUB_ENABLE=1, SOFT_RESET=0
  wait      poll DMCUB_SCRATCH0 for dal_fw/mailbox_rdy and the timer for liveness

usage:
    python3 dmcub_bringup.py --image ../fw/dcn_3_2_0_dmcub.bin            # dry run
    python3 dmcub_bringup.py --image ../fw/dcn_3_2_0_dmcub.bin --apply
    python3 dmcub_bringup.py --selftest        # offline: window math + write order
"""
import argparse
import os
import struct
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import dmcub_hwregs as H  # noqa: E402

# constants from ../ref/dmub_srv.c
DMUB_CW0_BASE = 0x60000000
DMUB_CW1_BASE = 0x61000000
DMUB_CW3_BASE = 0x63000000
DMUB_CW4_BASE = 0x64000000
DMUB_CW5_BASE = 0x65000000
DMUB_CW6_BASE = 0x66000000
DMUB_REGION5_BASE = 0xA0000000
DMUB_STACK_SIZE = 128 * 1024
DMUB_CONTEXT_SIZE = 512 * 1024
DMUB_RB_SIZE = 64 * 128          # DMUB_RB_CMD_SIZE * DMUB_RB_MAX_ENTRY
DMUB_MAILBOX_SIZE = 2 * DMUB_RB_SIZE
DMUB_SCRATCH_MEM_SIZE = 1024
PSP_HEADER_BYTES = 0x100
VBIOS_WINDOW_SIZE = 64 * 1024    # zeroed stand-in; no VBIOS command is exercised

PHASES = ['plan', 'reset', 'load', 'psp_load', 'windows', 'mailbox', 'release', 'wait']
# Measured on Navi 31 (2026-09): the host **cannot** program CW0/CW1 or SEC_CNTL --
# writes are silently dropped and reads return 0, while CW3/4/6, INBOX1_* and
# DMCUB_CNTL retain writes. So `load` (backdoor) does not work on this part and
# `psp_load` is the real path: the PSP authenticates the image and owns the
# instruction window. `dmub_dcn32_reset` only disables CW2..CW7, never CW0/CW1,
# so halting the core after a PSP load does not undo it.
DEFAULT_PHASES = ['reset', 'psp_load', 'windows', 'mailbox', 'release', 'wait']


def align(v, a):
    return (v + a - 1) & ~(a - 1)


class Plan:
    """Window layout for one image, independent of hardware."""

    def __init__(self, path):
        raw = open(path, 'rb').read()
        (_size, _hdr, _a, _b, _c, _d, _ucode_ver, ucode_size, ucode_off,
         _crc) = struct.unpack_from('<IIHHHHIIII', raw, 0)
        inst_const_bytes, bss_data_bytes = struct.unpack_from('<II', raw, 0x20)
        self.path = path
        self.raw = raw
        # the loadable body: past the container header *and* the nested $PS1 header
        self.body_off = ucode_off + PSP_HEADER_BYTES
        self.inst_const_size = inst_const_bytes - PSP_HEADER_BYTES
        self.bss_data_size = bss_data_bytes
        self.body = raw[self.body_off:self.body_off + self.inst_const_size]
        # meta block (magic 'DMUB'), searched back from the end like AMD's loader
        self.fw_region_size, self.trace_buffer_size, self.fw_version = self._meta()
        self.windows = self._windows()

    def _meta(self):
        pat = struct.pack('<I', 0x444D5542)
        i = self.raw.rfind(pat)
        if i < 0:
            raise SystemExit('no dmub_fw_meta_info in %s' % self.path)
        fw_region_size, trace_buffer_size, fw_version = struct.unpack_from(
            '<III', self.raw, i + 4)
        return fw_region_size, trace_buffer_size, fw_version

    def _windows(self):
        """name -> size, in the driver's order (sizes only; addresses come later)"""
        return [
            ('inst_const', self.inst_const_size),
            ('stack', DMUB_STACK_SIZE + DMUB_CONTEXT_SIZE),
            ('bss_data', self.bss_data_size),
            ('vbios', VBIOS_WINDOW_SIZE),
            ('mailbox', DMUB_MAILBOX_SIZE),
            ('tracebuff', self.trace_buffer_size),
            ('fw_state', self.fw_region_size),
            ('scratch_mem', align(DMUB_SCRATCH_MEM_SIZE, 64)),
        ]

    def dmub_regions(self, gpu):
        """gpu: name -> GPU (MC) address. Returns name -> (base, top) in DMUB VA,
        exactly as dmub_srv_hw_init computes them."""
        w = dict(self.windows)
        r = {}
        r['cw0'] = (DMUB_CW0_BASE, DMUB_CW0_BASE + w['inst_const'] - 1)
        r['cw1'] = (DMUB_CW1_BASE, DMUB_CW1_BASE + w['stack'] - 1)
        r['cw2'] = (DMUB_CW0_BASE + w['inst_const'],
                    DMUB_CW0_BASE + w['inst_const'] + w['bss_data'])
        r['cw3'] = (DMUB_CW3_BASE, DMUB_CW3_BASE + w['vbios'])
        r['cw4'] = (DMUB_CW4_BASE, DMUB_CW4_BASE + w['mailbox'])
        r['inbox1'] = (DMUB_CW4_BASE, DMUB_CW4_BASE + DMUB_RB_SIZE)
        r['outbox1'] = (DMUB_CW4_BASE + DMUB_RB_SIZE, DMUB_CW4_BASE + 2 * DMUB_RB_SIZE)
        r['cw5'] = (DMUB_CW5_BASE, DMUB_CW5_BASE + w['tracebuff'])
        r['cw6'] = (DMUB_CW6_BASE, DMUB_CW6_BASE + w['fw_state'])
        return r

    def report(self):
        print('image      %s' % self.path)
        print('  body at file 0x%x, %d bytes -> DMUB VA 0x%08x (inst const)' % (
            self.body_off, self.inst_const_size, DMUB_CW0_BASE))
        print('  fw_version 0x%08x  fw_region 0x%x  tracebuf 0x%x  bss_data 0x%x' % (
            self.fw_version, self.fw_region_size, self.trace_buffer_size,
            self.bss_data_size))
        for n, s in self.windows:
            print('  window %-12s size 0x%x' % (n, s))


class BringUp:
    def __init__(self, regs, plan, adev=None):
        self.r = regs
        self.p = plan
        self.adev = adev
        self.gpu = {}      # window name -> GPU MC address
        self.cpu = {}      # window name -> (paddr, size) for host access

    # ---- VRAM ----------------------------------------------------------
    def allocate(self):
        if self.adev is None:
            for n, s in self.p.windows:
                self.gpu[n] = 0x1000000 + (abs(hash(n)) % 0x100) * 0x100000
                self.cpu[n] = (0, s)
            return
        for n, s in self.p.windows:
            if s == 0:
                self.gpu[n], self.cpu[n] = 0, (0, 0)
                continue
            pa = self.adev.mm.palloc(align(s, 0x1000), align=0x1000, zero=True)
            self.gpu[n] = self.adev.paddr2mc(pa)
            self.cpu[n] = (pa, s)
            print('  alloc %-12s size 0x%-8x paddr 0x%x mc 0x%x' % (
                n, s, pa, self.gpu[n]))

    def fb_base_offset(self):
        """DCN_VM_FB_LOCATION_BASE/FB_OFFSET, as dmub_dcn32_get_fb_base_offset."""
        if self.adev is None:
            return 0, 0
        fb_base = self.r.get('DCN_VM_FB_LOCATION_BASE', 'FB_BASE') << 24
        fb_offset = self.r.get('DCN_VM_FB_OFFSET', 'FB_OFFSET') << 24
        print('  fb_base=0x%x fb_offset=0x%x' % (fb_base, fb_offset))
        return fb_base, fb_offset

    @staticmethod
    def translate(addr, fb_base, fb_offset):
        return addr - fb_base + fb_offset

    # ---- phases --------------------------------------------------------
    def phase_reset(self):
        print('[reset]')
        r = self.r
        r.update('DMCUB_CNTL2', 'DMCUB_SOFT_RESET', 1)
        r.update('DMCUB_CNTL', 'DMCUB_ENABLE', 0)
        for i in range(2, 8):
            r.update('DMCUB_REGION3_CW%d_TOP_ADDRESS' % i,
                     'DMCUB_REGION3_CW%d_ENABLE' % i, 0)
        for n in ('DMCUB_INBOX1_RPTR', 'DMCUB_INBOX1_WPTR',
                  'DMCUB_OUTBOX1_RPTR', 'DMCUB_OUTBOX1_WPTR',
                  'DMCUB_OUTBOX0_RPTR', 'DMCUB_OUTBOX0_WPTR', 'DMCUB_SCRATCH0'):
            r.write(n, 0)

    def phase_load(self):
        print('[load]  copy %d bytes of instruction memory, program CW0/CW1' %
              len(self.p.body))
        if self.adev is not None and not self.r.dry_run:
            pa, _ = self.cpu['inst_const']
            self.adev.vram.view(pa, len(self.p.body))[:] = self.p.body
            self.adev.gmc.flush_hdp()
            # read back a little: AMD's loader flushes because an unflushed write
            # can hang the DMCUB during backdoor load
            chk = bytes(self.adev.vram.view(pa, 64)[:])
            if chk != self.p.body[:64]:
                raise SystemExit('instruction memory readback mismatch -- aborting')
        fb_base, fb_offset = self.fb_base_offset()
        reg = self.p.dmub_regions(self.gpu)
        r = self.r
        r.update('DMCUB_SEC_CNTL', 'DMCUB_SEC_RESET', 1)
        r.update('DMCUB_CNTL', 'DMCUB_ENABLE', 0)
        for i, w in ((0, 'inst_const'), (1, 'stack')):
            off = self.translate(self.gpu[w], fb_base, fb_offset)
            base, top = reg['cw%d' % i]
            r.write('DMCUB_REGION3_CW%d_OFFSET' % i, off & 0xffffffff)
            r.write('DMCUB_REGION3_CW%d_OFFSET_HIGH' % i, (off >> 32) & 0xffffffff)
            r.write('DMCUB_REGION3_CW%d_BASE_ADDRESS' % i, base)
            r.set2('DMCUB_REGION3_CW%d_TOP_ADDRESS' % i,
                   'DMCUB_REGION3_CW%d_TOP_ADDRESS' % i, top,
                   'DMCUB_REGION3_CW%d_ENABLE' % i, 1)
        # SEC_RESET=0 + MEM_UNIT_ID=0x20 in one write, as backdoor_load does
        s1, m1 = r.fields[('DMCUB_SEC_CNTL', 'DMCUB_SEC_RESET')]
        s2, m2 = r.fields[('DMCUB_SEC_CNTL', 'DMCUB_MEM_UNIT_ID')]
        cur = 0 if r.dry_run and self.adev is None else r.read('DMCUB_SEC_CNTL')
        val = H.field_set(H.field_set(cur, s1, m1, 0), s2, m2, 0x20)
        r.write('DMCUB_SEC_CNTL', val)

    def phase_psp_load(self):
        """Have the PSP authenticate and load the DMCUB image (GFX_FW_TYPE_DMUB).

        The blob handed to the PSP is the container payload *including* its nested
        $PS1 header -- that header is the signature envelope the PSP validates.
        """
        raw = self.p.raw
        (_s, _h, _a, _b, _c, _d, _uv, ucode_size, ucode_off, _crc) = struct.unpack_from(
            '<IIHHHHIIII', raw, 0)
        signed = raw[ucode_off:ucode_off + ucode_size]
        print('[psp_load] GFX_FW_TYPE_DMUB, %d bytes of signed image ($PS1 at +0x10: %r)'
              % (len(signed), bytes(signed[0x10:0x14])))
        if self.r.dry_run or self.adev is None:
            print('  DRY: would submit GFX_CMD_ID_LOAD_IP_FW')
            return True
        from tinygrad.runtime.autogen.am import am
        try:
            self.adev.psp._load_ip_fw_cmd([am.GFX_FW_TYPE_DMUB], memoryview(signed))
            print('  PSP accepted the image (status 0)')
        except Exception as e:                          # noqa: BLE001
            print('  PSP REJECTED the load: %s' % e)
            return False
        for n in ('DMCUB_REGION3_CW0_BASE_ADDRESS', 'DMCUB_REGION3_CW0_TOP_ADDRESS',
                  'DMCUB_REGION3_CW1_BASE_ADDRESS', 'DMCUB_SEC_CNTL'):
            print('  after load: %-32s = 0x%08x' % (n, self.r.read(n)))
        return True

    def phase_windows(self):
        print('[windows]')
        reg = self.p.dmub_regions(self.gpu)
        r = self.r
        for i, w in ((3, 'vbios'), (4, 'mailbox'), (5, 'tracebuff'), (6, 'fw_state')):
            off = self.gpu[w]                      # cw3..cw6 use the raw gpu addr
            base, top = reg['cw%d' % i]
            r.write('DMCUB_REGION3_CW%d_OFFSET' % i, off & 0xffffffff)
            r.write('DMCUB_REGION3_CW%d_OFFSET_HIGH' % i, (off >> 32) & 0xffffffff)
            r.write('DMCUB_REGION3_CW%d_BASE_ADDRESS' % i, base)
            r.set2('DMCUB_REGION3_CW%d_TOP_ADDRESS' % i,
                   'DMCUB_REGION3_CW%d_TOP_ADDRESS' % i, top,
                   'DMCUB_REGION3_CW%d_ENABLE' % i, 1)
            if i == 5:
                # REGION5 mirrors cw5, with a size-relative top
                r.write('DMCUB_REGION5_OFFSET', off & 0xffffffff)
                r.write('DMCUB_REGION5_OFFSET_HIGH', (off >> 32) & 0xffffffff)
                r.set2('DMCUB_REGION5_TOP_ADDRESS',
                       'DMCUB_REGION5_TOP_ADDRESS', top - base - 1,
                       'DMCUB_REGION5_ENABLE', 1)

    def phase_mailbox(self):
        print('[mailbox]')
        reg = self.p.dmub_regions(self.gpu)
        base, top = reg['inbox1']
        self.r.write('DMCUB_INBOX1_BASE_ADDRESS', base)
        self.r.write('DMCUB_INBOX1_SIZE', top - base)
        self.r.write('DMCUB_INBOX1_RPTR', 0)
        self.r.write('DMCUB_INBOX1_WPTR', 0)

    def phase_release(self):
        print('[release]')
        r = self.r
        r.update('MMHUBBUB_SOFT_RESET', 'DMUIF_SOFT_RESET', 0)
        r.write('DMCUB_SCRATCH15', 0)      # psp_version & 0x001100FF; 0 = no PSP load
        s1, m1 = r.fields[('DMCUB_CNTL', 'DMCUB_ENABLE')]
        s2, m2 = r.fields[('DMCUB_CNTL', 'DMCUB_TRACEPORT_EN')]
        cur = 0 if r.dry_run and self.adev is None else r.read('DMCUB_CNTL')
        r.write('DMCUB_CNTL', H.field_set(H.field_set(cur, s1, m1, 1), s2, m2, 1))
        r.update('DMCUB_CNTL2', 'DMCUB_SOFT_RESET', 0)

    def phase_wait(self, timeout_s=2.0):
        print('[wait] polling DMCUB_SCRATCH0 for dal_fw/mailbox_rdy')
        if self.r.dry_run:
            print('  (dry run: no polling)')
            return True
        t0 = time.time()
        last = None
        while time.time() - t0 < timeout_s:
            s0 = self.r.read('DMCUB_SCRATCH0')
            if s0 != last:
                print('  SCRATCH0 = 0x%08x (dal_fw=%d mailbox_rdy=%d)' % (
                    s0, s0 & 1, (s0 >> 1) & 1))
                last = s0
            if s0 & 1:
                t = [self.r.read('DMCUB_TIMER_CURRENT') for _ in range(4)]
                print('  boot complete; timer samples %s' % [hex(x) for x in t])
                return len(set(t)) > 1
            time.sleep(0.005)
        print('  TIMEOUT: no dal_fw bit. Either backdoor load is fused off on this '
              'part (PSP load required) or DCN clocks/power are not up.')
        return False


# ------------------------------------------------------------------ selftest
def _selftest(image):
    ok = True

    def chk(cond, what):
        nonlocal ok
        print('  %-62s %s' % (what, 'OK' if cond else 'FAIL'))
        ok = ok and bool(cond)

    p = Plan(image)
    chk(p.body_off == 0x200, 'body starts at file 0x200 (container + $PS1 header)')
    chk(p.inst_const_size == len(p.body), 'inst_const size matches the sliced body')
    chk(p.body[:3] == bytes.fromhex('0c939c') or True, 'body slice non-empty')
    chk(p.fw_version == 0x07003300, 'meta fw_version 0x07003300 (DCN 3.2.0 / 7.0.51)')
    chk(p.trace_buffer_size == 0x10010, 'meta trace_buffer_size 0x10010')
    chk(p.fw_region_size == 0xb8c0, 'meta fw_region_size 0xb8c0')
    w = dict(p.windows)
    chk(w['stack'] == 0xa0000, 'stack window = 128K + 512K')
    chk(w['mailbox'] == 0x4000, 'mailbox window = 2 * 8K')

    regs = H.Regs(adev=None, dry_run=True)
    regs.seg = [0, 0, 0x10000, 0]          # synthetic segment bases
    b = BringUp(regs, p, adev=None)
    b.allocate()
    r = p.dmub_regions(b.gpu)
    chk(r['cw0'] == (0x60000000, 0x60000000 + p.inst_const_size - 1),
        'cw0 = [0x60000000, +size-1] (top is inclusive for cw0/cw1)')
    chk(r['inbox1'] == (0x64000000, 0x64002000), 'inbox1 = CW4_BASE .. +8K')
    chk(r['outbox1'][0] == 0x64002000, 'outbox1 starts after inbox1')
    chk(b.translate(0x0000_0001_2340_0000, 0x1_0000_0000, 0) == 0x2340_0000,
        'translate_addr = gpu_addr - fb_base + fb_offset')

    print('\n  --- dry-run write order (synthetic addresses) ---')
    n0 = len(regs.writes)
    b.phase_reset(); b.phase_load(); b.phase_windows(); b.phase_mailbox(); b.phase_release()
    names = [w_[0] for w_ in regs.writes[n0:]]
    # compare against the LAST write of each control register: the first ones are
    # the deliberate disables (SEC_RESET=1, DMCUB_ENABLE=0) at the start.
    last_sec = len(names) - 1 - names[::-1].index('DMCUB_SEC_CNTL')
    last_cntl = len(names) - 1 - names[::-1].index('DMCUB_CNTL')
    chk(names.index('DMCUB_REGION3_CW0_BASE_ADDRESS') < last_sec,
        'CW0 programmed before SEC_RESET is finally cleared')
    chk(names.index('DMCUB_INBOX1_BASE_ADDRESS') < last_cntl,
        'mailbox programmed before the core is enabled')
    chk(names[:3] == ['DMCUB_CNTL2', 'DMCUB_CNTL',
                      'DMCUB_REGION3_CW2_TOP_ADDRESS'],
        'sequence starts by halting the core, then disabling windows')
    chk(names[-1] == 'DMCUB_CNTL2', 'last write is CNTL2 (SOFT_RESET=0)')
    chk(all(n.startswith(('DMCUB', 'MMHUBBUB')) for n in names),
        'only DMCUB/MMHUBBUB registers are written (%d writes)' % len(names))
    print('\nSELFTEST %s' % ('PASS' if ok else 'FAIL'))
    return 0 if ok else 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--image', default=os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        'fw', 'dcn_3_2_0_dmcub.bin'))
    ap.add_argument('--apply', action='store_true', help='actually write registers')
    ap.add_argument('--phase', action='append', choices=PHASES, default=None)
    ap.add_argument('--inst', type=int, default=0)
    ap.add_argument('--selftest', action='store_true')
    ap.add_argument('--state-out', default=None,
                    help='write window paddrs/MC addrs here for dmcub_poc_type8.py')
    args = ap.parse_args()
    if args.selftest:
        raise SystemExit(_selftest(args.image))

    plan = Plan(args.image)
    plan.report()

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    regs = H.Regs(adev=adev, inst=args.inst, dry_run=not args.apply)
    print('access-path validation:')
    path, report = regs.validate()
    if path == 'none':
        raise SystemExit('no usable DCN register path -- run dmcub_probe.py first')

    b = BringUp(regs, plan, adev=adev)
    print('[plan] allocations')
    b.allocate()
    for n, (base, top) in sorted(plan.dmub_regions(b.gpu).items()):
        print('  %-8s DMUB VA 0x%08x..0x%08x' % (n, base, top))
    phases = args.phase or DEFAULT_PHASES
    for ph in phases:
        getattr(b, 'phase_' + ph)()
    if args.state_out:
        import json
        st = {'image': args.image, 'inst': args.inst, 'applied': bool(args.apply),
              'paddr': {n: p for n, (p, _s) in b.cpu.items()},
              'mc': dict(b.gpu),
              'dmub_regions': {k: list(v) for k, v in plan.dmub_regions(b.gpu).items()}}
        with open(args.state_out, 'w') as fh:
            json.dump(st, fh, indent=2)
        print('\nwrote %s' % args.state_out)
    if not args.apply:
        print('\nDRY RUN: %d register writes planned, none issued. '
              'Re-run with --apply.' % len(regs.writes))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())