#!/usr/bin/env python3
"""Test the SDMA SRBM_WRITE privilege gate on silicon, and name it by behaviour.

Static finding (`../AUDIT-SDMA-01.md`): the `SRBM_WRITE` packet -- which writes a
host-named register -- is gated before it acts:

    ldw r5, reg[r0, #0x4a14] ; and #0x200 ; lsr #9         ; bit 9
    ldw r5, reg[r0, #0x5968] ; lsr #31                     ; OR bit 31
    ldw r5, reg[r0, #0x5940] ; lsr #23 ; and #1            ; AND bit 23
    and r5, r6, r5
    cbz r5, _PKT_0xe00_2                                   ; -> error 0x3c/0x32

Attempts to *name* those registers by matching the F32 index space against
`gc_11_0_0_offset.h` failed: the best-fitting constant base scored 25/54 against a null
95th percentile of 24, i.e. not significant (see AUDIT-SDMA-01). So the gate is
identified by **behaviour** instead, which is stronger anyway.

Bit 23 is the giveaway: `SDMA*_QUEUE*_RB_CNTL__RB_PRIV` is mask 0x00800000 = bit 23.
If F32 index 0x5940 is the current queue's RB_CNTL, then clearing `RB_PRIV` must turn
an accepted SRBM_WRITE into a rejected one. tinygrad's AM driver sets `rb_priv=1` when
it configures the queue, so both states are reachable.

Experiment, per RB_PRIV setting:
  1. host zeroes a scratch register and reads it back (control that we can see it)
  2. a NOP packet is submitted and the ring pointer must advance (control that the
     queue actually runs)
  3. an SRBM_WRITE packet targeting that scratch register with a magic value
  4. read the scratch register from the host

Expected if the static reading is right: RB_PRIV=1 -> magic lands; RB_PRIV=0 -> it does
not. Any other outcome falsifies the reading and is reported as such.

Safety: the target is a GRBM scratch register (no side effects). A rejected packet may
halt the queue; AM resets SDMA on the next device open.

usage: PYTHONPATH=/root/psp/tinygrad sdma_srbm_probe.py [--apply]
"""
import argparse
import ctypes
import struct
import sys
import time

SDMA_OP_NOP = 0
SDMA_OP_SRBM_WRITE = 14
MAGIC = 0x5100D00D


class Probe:
    def __init__(self, adev, queue=0):
        self.adev = adev
        self.q = 'regSDMA0_QUEUE%d' % queue
        self.ring_size = 0x1000
        self.ring_pa = adev.mm.palloc(self.ring_size, zero=True, boot=False)
        self.rptr_pa = adev.mm.palloc(0x1000, zero=True, boot=False)
        self.wptr_pa = adev.mm.palloc(0x1000, zero=True, boot=False)
        self.ring_mc = adev.paddr2mc(self.ring_pa)
        self.wptr = 0

    def reg(self, suffix):
        return self.adev.reg(self.q + suffix)

    def setup(self, rb_priv):
        a = self.adev
        self.reg('_MINOR_PTR_UPDATE').write(0x1)
        a.wreg_pair(self.q + '_RB_RPTR', '', '_HI', 0)
        a.wreg_pair(self.q + '_RB_WPTR', '', '_HI', 0)
        a.wreg_pair(self.q + '_RB_BASE', '', '_HI', self.ring_mc >> 8)
        a.wreg_pair(self.q + '_RB_RPTR_ADDR', '_LO', '_HI', a.paddr2mc(self.rptr_pa))
        a.wreg_pair(self.q + '_RB_WPTR_POLL_ADDR', '_LO', '_HI', a.paddr2mc(self.wptr_pa))
        self.reg('_MINOR_PTR_UPDATE').write(0x0)
        # wptr poll left OFF so we submit by writing RB_WPTR directly
        self.reg('_RB_CNTL').write(rb_vmid=0, rptr_writeback_enable=1,
                                   rptr_writeback_timer=4, rb_enable=1,
                                   rb_priv=rb_priv,
                                   rb_size=(self.ring_size // 4).bit_length() - 1)
        self.reg('_IB_CNTL').update(ib_enable=1)
        self.wptr = 0
        got = self.reg('_RB_CNTL').read_bitfields().get('rb_priv')
        return got

    def submit(self, dwords, timeout_s=1.0):
        """Append dwords to the ring, kick RB_WPTR, wait for RB_RPTR to catch up."""
        off = (self.wptr * 4) % self.ring_size
        blob = struct.pack('<%dI' % len(dwords), *dwords)
        self.adev.vram.view(self.ring_pa + off, len(blob))[:] = blob
        self.adev.gmc.flush_hdp()
        self.wptr += len(dwords)
        self.adev.wreg_pair(self.q + '_RB_WPTR', '', '_HI', self.wptr * 4)
        t0 = time.time()
        while time.time() - t0 < timeout_s:
            r = self.adev.rreg(self.adev.reg(self.q + '_RB_RPTR').addr[0])
            if r == (self.wptr * 4) & 0xffffffff:
                return True, time.time() - t0
            time.sleep(0.002)
        return False, time.time() - t0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--target', default='regGRBM_SCRATCH_REG0')
    args = ap.parse_args()
    if not args.apply:
        print('dry run: pass --apply to touch the device')
        return 0

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    try:
        treg = adev.reg(args.target)
    except KeyError:
        raise SystemExit('no such register: %s' % args.target)
    tidx = treg.addr[0]
    print('target %s -> absolute index 0x%x' % (args.target, tidx))

    p = Probe(adev)
    print('ring paddr 0x%x mc 0x%x' % (p.ring_pa, p.ring_mc))

    # Phase B trick: a queue configured from scratch with RB_PRIV=0 does not even
    # fetch a NOP (the ring fetch itself is privileged), so "no write" there proves
    # nothing. Instead bring the queue up privileged, then clear ONLY RB_PRIV on the
    # live queue and re-submit: the ring stays configured, so if the NOP still retires
    # while SRBM_WRITE stops taking effect, the gate is isolated to RB_PRIV.
    results = {}
    for priv in (1, 0):
        print('\n=== RB_PRIV = %d' % priv)
        if priv == 1:
            got = p.setup(rb_priv=1)
        else:
            p.reg('_RB_CNTL').update(rb_priv=0)     # live flip, queue stays up
            got = p.reg('_RB_CNTL').read_bitfields().get('rb_priv')
        print('  RB_CNTL.rb_priv reads back as %s' % got)
        treg.write(0)
        base = treg.read()
        print('  control: host zeroed %s -> reads 0x%08x' % (args.target, base))

        ok, dt = p.submit([SDMA_OP_NOP])
        print('  control: NOP packet consumed=%s (%.1f ms)' % (ok, 1000 * dt))
        if not ok:
            print('  !! queue is not running; result for this setting is inconclusive')
            results[priv] = None
            continue

        hdr = SDMA_OP_SRBM_WRITE | (0xf << 28)
        magic = MAGIC if priv else (MAGIC ^ 0xffff)
        ok2, dt2 = p.submit([hdr, tidx, magic])
        val = treg.read()
        print('  SRBM_WRITE consumed=%s (%.1f ms) -> %s = 0x%08x %s' % (
            ok2, 1000 * dt2, args.target, val,
            'WROTE' if val == magic else 'not written'))
        results[priv] = (ok2, val == magic)

    print('\n=== result')
    r1, r0 = results.get(1), results.get(0)
    if r1 and r0:
        if r1[1] and not r0[1]:
            print('  CONFIRMED: SRBM_WRITE executes only with RB_PRIV=1.')
            print('  => F32 register index 0x5940 is the queue RB_CNTL, and its bit 23')
            print('     (RB_PRIV) is the privilege gate read at 0x0b07-0x0b0a.')
        elif r1[1] and r0[1]:
            print('  FALSIFIED: the write landed with RB_PRIV=0 too -- the gate is not')
            print('  RB_PRIV, or it is satisfied by the OR-term (0x4a14 bit 9 / 0x5968')
            print('  bit 31) independently. The static reading needs revisiting.')
        elif not r1[1]:
            print('  INCONCLUSIVE: the write did not land even with RB_PRIV=1, so this')
            print('  setup does not exercise the path (wrong register, aperture, or')
            print('  queue state). No conclusion about the gate.')
    else:
        print('  inconclusive: a control failed')
    for rn in ('regSDMA0_STATUS_REG', 'regSDMA0_QUEUE0_CONTEXT_STATUS'):
        try:
            print('%s = 0x%08x' % (rn, adev.reg(rn).read()))
        except KeyError:
            pass
    return 0


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