#!/usr/bin/env python3
"""Does the TMR fence key on the access TAG rather than the requester?

Established: the fence blocks the host BAR path and blocks SMU DMA, and was
called "requester-agnostic". But host BAR reads are PA-based and bypass the VM
entirely, and the SMU is one requester. Nobody has tried changing the *tag* on
the access -- and the host driver owns the GPU page tables, which carry
AMDGPU_PTE_TMZ (bit 3), the trusted-memory-zone bit.

Experiment: point a GPU buffer's VA at TMR physical memory, with and without the
TMZ bit, and have a GPU kernel read it.

Controls, in order:
  0. host BAR read of the TMR PA -- must show the documented fence behaviour
  1. remap the VA to a normal VRAM page we own, seeded with a pattern -- the
     kernel must read the pattern back. Validates remap + kernel + TLB flush.
  2. TMR PA without TMZ -- expected blocked
  3. TMR PA with TMZ    -- the question

Progress to /tmp/tmz.txt before every risky step.
"""
import struct, time
from tinygrad import Device, Tensor
from tinygrad.runtime.support.memory import AddrSpace, PageTableTraverseContext
import tinygrad.runtime.autogen.am.am as am

OUT = "/tmp/tmz.txt"


def say(s):
    with open(OUT, "a") as f:
        f.write(s + "\n")
    print(s, flush=True)


d = Device["AMD"]
adev = d.iface.dev_impl
mm = adev.mm
hb = adev.reg("regMP0_SMN_C2PMSG_81")


def alive():
    a = hb.read(); time.sleep(0.2); return hb.read() != a


TMR_PA = 0x900000          # inside the documented TMR range 0x300000..0x1600000
SCRATCH_PA = None
N = 1024          # 4096 bytes: PTE granularity is 0x1000
SZ = N * 4

say("\n######## TMZ / tag experiment ########")
say("psp alive: %s" % alive())

# control 0: host BAR read of the TMR
raw = bytes(adev.vram[TMR_PA:TMR_PA + 32])
say("control 0 -- host BAR read of TMR PA %#x: %s" % (TMR_PA, raw.hex()))

src = (Tensor.zeros(N, dtype="float32").contiguous() + 1.0).contiguous().realize()
dst = (Tensor.zeros(N, dtype="float32").contiguous()).realize()
va = src.uop.buffer._buf
say("src VA %#x size %#x" % (va, SZ))

SCRATCH_PA = mm.palloc(SZ, zero=True, boot=False)
pattern = b"".join(struct.pack("<f", 3.5) for _ in range(N))
adev.vram[SCRATCH_PA:SCRATCH_PA + SZ] = pattern
say("scratch PA %#x seeded with 3.5" % SCRATCH_PA)


def remap(pa, tmz):
    mm.unmap_range(va, SZ)
    mm.map_range(va, SZ, [(pa, SZ)], aspace=AddrSpace.PHYS)
    if tmz:
        ctx = PageTableTraverseContext(adev, mm.root_page_table, va, inspect=True)
        n = 0
        for _off, pt, pte_idx, pte_cnt, _cov in ctx.next(SZ):
            for k in range(pte_cnt):
                pt.entries[pte_idx + k] |= am.AMDGPU_PTE_TMZ
                n += 1
        say("      set TMZ on %d PTEs" % n)
    adev.gmc.flush_tlb(ip='GC', vmid=0)


def gpu_read(tag):
    try:
        out = (src + 0.0).realize().numpy()
        say("   %-28s kernel read -> first4=%s  distinct=%d" % (tag, out[:4], len(set(out.tolist()))))
        return out
    except Exception as e:
        say("   %-28s kernel FAILED: %s %s" % (tag, type(e).__name__, str(e)[:70]))
        return None


say("\ncontrol 1 -- remap VA to our scratch page (no TMZ)")
remap(SCRATCH_PA, False)
c1 = gpu_read("scratch, no TMZ")
say("   expect 3.5 -> remap+kernel path %s" % ("VALID" if c1 is not None and abs(float(c1[0]) - 3.5) < 1e-6 else "NOT VALIDATED"))
say("   psp alive %s" % alive())

if c1 is not None and abs(float(c1[0]) - 3.5) < 1e-6:
    say("\ncontrol 1b -- same scratch page WITH TMZ (does TMZ break a normal read?)")
    remap(SCRATCH_PA, True)
    gpu_read("scratch, TMZ")
    say("   psp alive %s" % alive())

    say("\ntest 2 -- TMR PA, no TMZ")
    remap(TMR_PA, False)
    gpu_read("TMR, no TMZ")
    say("   psp alive %s" % alive())

    say("\ntest 3 -- TMR PA, WITH TMZ")
    remap(TMR_PA, True)
    gpu_read("TMR, TMZ")
    say("   psp alive %s" % alive())
else:
    say("\ncontrol 1 failed -- not proceeding; a broken remap path cannot test a fence")

say("\ndone. psp alive %s" % alive())