#!/usr/bin/env python3
"""
READ-ONLY check: what VRAM size does a Navi31 (RX 7900 XTX) report over a
USB<->PCIe enclosure (ASM2464PD-class bridge)?

WHAT THIS DOES: reads four registers and prints them.
WHAT THIS DOES NOT DO: no writes, no flash access, no reset, no power cycling,
no firmware loading. It cannot change the state of your card.

REQUIREMENTS: `tinygrad` importable (a checkout on PYTHONPATH or installed),
plus whatever permissions you normally use for tinygrad AMD-over-USB.

RUN:
    python3 vram_check.py

Then paste the whole output back. Thanks!
"""
import struct, sys

try:
    from tinygrad.runtime.support.usb import USB3
    from tinygrad.runtime.support.system import USBPCIDevice
except Exception as e:
    sys.exit("could not import tinygrad (%s: %s)\n"
             "make sure a tinygrad checkout is on PYTHONPATH" % (type(e).__name__, e))

# Navi31 / MP0 13.0.0 specific offsets.
MEMSIZE_DWORD = 0xde3     # mmRCC_CONFIG_MEMSIZE, dword index (byte 0x378C).
                          # Fixed RCC register; readable before IP discovery,
                          # which is exactly why tinygrad uses it.
TSU_CMD       = 0x582CC   # = MP0 C2PMSG_115, PSP TSU mailbox command register
TSU_DATA      = 0x582D0   # = MP0 C2PMSG_116, its data register
BRIDGE_IDS    = ((0x3801, 0x0001), (0xADD1, 0x0001))

devs = []
for vid, pid in BRIDGE_IDS:
    devs += list(USB3.list_devices(vid, pid))
print("USB bridges found : %d" % len(devs))
if not devs:
    sys.exit("no bridge found - is the enclosure connected and powered on?")

gpu  = USBPCIDevice("AM", *devs[0])
bar5 = gpu.map_bar(5, fmt='B')
rd32 = lambda off: struct.unpack('<I', bytes(bar5[off:off+4]))[0]

memsize = rd32(MEMSIZE_DWORD * 4)
print()
print("mmRCC_CONFIG_MEMSIZE     = %#010x   -> VRAM = %d MB" % (memsize, memsize))
print("TSU mailbox (C2PMSG_115) = %#010x   (bit31 set = PSP task ready)" % rd32(TSU_CMD))
print("TSU data    (C2PMSG_116) = %#010x" % rd32(TSU_DATA))
try:
    print("VRAM[0:16]               = %s" % bytes(gpu.map_bar(0)[0:16]).hex())
except Exception as e:
    print("VRAM[0:16]               = <unreadable: %s>" % e)

print()
if memsize == 0:
    print(">>> MEMSIZE reads 0: the memory controller reports no VRAM,")
    print(">>> which means memory training has not completed.")
else:
    print(">>> MEMSIZE OK: %d MB (%.1f GB). Memory training completed." % (memsize, memsize / 1024))

print("""
-----------------------------------------------------------
Please paste this whole output back.

The key number is mmRCC_CONFIG_MEMSIZE:
  0x6000 (24576 MB) = healthy, memory trained
  0x0               = memory not trained
-----------------------------------------------------------""")