#!/usr/bin/env python3
"""
Patch 3: amdgpu_gmc.c — KIQ bypass + Dead-GPU detection (2 sub-patches)

Part of Layer 2 of the v3 three-layer protection. Modifies the centralized
(generation-agnostic) TLB flush code to:
  a) Bypass KIQ in amdgpu_gmc_flush_gpu_tlb_pasid — direct MMIO callout
  b) Bypass KIQ + dead-GPU detection in amdgpu_gmc_fw_reg_write_reg_wait

Usage:
    python3 patch3_amdgpu_gmc.py /path/to/drivers/gpu/drm/amd/amdgpu
"""
import sys
import os

if len(sys.argv) < 2:
    print("Usage: python3 patch3_amdgpu_gmc.py <amdgpu_source_dir>")
    print("  e.g.: python3 patch3_amdgpu_gmc.py ~/kernel-build/linux-6.19.6/drivers/gpu/drm/amd/amdgpu")
    sys.exit(1)

AMDGPU = sys.argv[1]
filepath = os.path.join(AMDGPU, "amdgpu_gmc.c")

if not os.path.isfile(filepath):
    print(f"ERROR: File not found: {filepath}")
    sys.exit(1)

with open(filepath, 'r') as f:
    content = f.read()

if 'BC-250' in content:
    print("Already patched, skipping.")
    sys.exit(0)

# ========================================
# Patch 3a: KIQ bypass in amdgpu_gmc_flush_gpu_tlb_pasid
# Insert AFTER down_read_trylock block, BEFORE KIQ ring code
# ========================================
func_start = content.find('int amdgpu_gmc_flush_gpu_tlb_pasid')
if func_start == -1:
    print("ERROR: amdgpu_gmc_flush_gpu_tlb_pasid not found")
    sys.exit(1)

# Find the trylock check
trylock = content.find('down_read_trylock', func_start)
if trylock == -1:
    print("ERROR: down_read_trylock not found")
    sys.exit(1)

# Find "return 0;" after trylock
return_0 = content.find('return 0;', trylock)
end_line = content.find('\n', return_0) + 1

bypass = (
    '\n'
    '\t/* ===== BC-250 v2 PATCH: KIQ bypass for PASID flush ===== */\n'
    '\t{\n'
    '\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);\n'
    '\t\tpr_warn_once("amdgpu: flush_gpu_tlb_pasid called, GC_HWIP=0x%08x "\n'
    '\t\t\t     "(10.1.3=0x%08x) kiq_flag=%d\\n",\n'
    '\t\t\t     gc_ver, IP_VERSION(10, 1, 3),\n'
    '\t\t\t     adev->gmc.flush_pasid_uses_kiq);\n'
    '\n'
    '\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0))) {\n'
    '\t\t\tpr_warn_once("amdgpu: BC-250 KIQ bypass active "\n'
    '\t\t\t\t     "(gc_ver=0x%08x)\\n", gc_ver);\n'
    '\t\t\tadev->gmc.gmc_funcs->flush_gpu_tlb_pasid(adev, pasid,\n'
    '\t\t\t\t\t\t\t\t  flush_type, all_hub,\n'
    '\t\t\t\t\t\t\t\t  inst);\n'
    '\t\t\tr = 0;\n'
    '\t\t\tgoto error_unlock_reset;\n'
    '\t\t}\n'
    '\t}\n'
    '\t/* ===== BC-250 v2 PATCH END ===== */\n'
)

content = content[:end_line] + bypass + content[end_line:]

# ========================================
# Patch 3b: KIQ bypass + dead-GPU in amdgpu_gmc_fw_reg_write_reg_wait
# Insert BEFORE the KIQ ring submission code
# ========================================
func2_start = content.find('void amdgpu_gmc_fw_reg_write_reg_wait')
if func2_start == -1:
    print("ERROR: amdgpu_gmc_fw_reg_write_reg_wait not found")
    sys.exit(1)

# Find the first operational code after variable declarations
spinlock = content.find('spin_lock_irqsave', func2_start)
if spinlock == -1:
    spinlock = content.find('ring->sched.ready', func2_start)
    if spinlock == -1:
        spinlock = content.find('if (', func2_start + 200)
        if spinlock == -1:
            print("ERROR: Could not find KIQ code in fw_reg_write_reg_wait")
            sys.exit(1)

spinlock_line_start = content.rfind('\n', 0, spinlock) + 1

bypass2 = (
    '\t/* ===== BC-250 v2+v3 PATCH: KIQ bypass + dead-GPU detection ===== */\n'
    '\t{\n'
    '\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);\n'
    '\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0))) {\n'
    '\t\t\tuint32_t tmp;\n'
    '\n'
    '\t\t\tpr_warn_once("amdgpu: BC-250 KIQ bypass active in "\n'
    '\t\t\t\t     "fw_reg_write_reg_wait (gc=0x%08x)\\n", gc_ver);\n'
    '\n'
    '\t\t\t/* v3: Health-check read before writing */\n'
    '\t\t\ttmp = RREG32_NO_KIQ(reg1);\n'
    '\t\t\tif (tmp == 0xFFFFFFFF) {\n'
    '\t\t\t\tdev_err_ratelimited(adev->dev,\n'
    '\t\t\t\t\t"BC-250: GPU unreachable in fw_reg_write_reg_wait "\n'
    '\t\t\t\t\t"(reg1=0x%x returned 0xFFFFFFFF), skipping\\n", reg1);\n'
    '\t\t\t\treturn;\n'
    '\t\t\t}\n'
    '\n'
    '\t\t\tWREG32_NO_KIQ(reg0, ref);\n'
    '\t\t\tfor (cnt = 0; cnt < adev->usec_timeout; cnt++) {\n'
    '\t\t\t\ttmp = RREG32_NO_KIQ(reg1);\n'
    '\t\t\t\t/* v3: Dead-GPU detection in polling loop */\n'
    '\t\t\t\tif (tmp == 0xFFFFFFFF) {\n'
    '\t\t\t\t\tdev_err_ratelimited(adev->dev,\n'
    '\t\t\t\t\t\t"BC-250: GPU died during reg_write_reg_wait "\n'
    '\t\t\t\t\t\t"(0xFFFFFFFF at reg1=0x%x)\\n", reg1);\n'
    '\t\t\t\t\treturn;\n'
    '\t\t\t\t}\n'
    '\t\t\t\tif ((tmp & mask) == (ref & mask))\n'
    '\t\t\t\t\treturn;\n'
    '\t\t\t\tudelay(1);\n'
    '\t\t\t}\n'
    '\t\t\tdev_warn(adev->dev, "BC-250: MMIO reg write/wait timeout "\n'
    '\t\t\t\t "reg0=0x%x reg1=0x%x\\n", reg0, reg1);\n'
    '\t\t\treturn;\n'
    '\t\t}\n'
    '\t}\n'
    '\t/* ===== BC-250 v2+v3 PATCH END ===== */\n'
)

content = content[:spinlock_line_start] + bypass2 + content[spinlock_line_start:]

with open(filepath, 'w') as f:
    f.write(content)

print("Patch 3 applied: amdgpu_gmc.c — 2 sub-patches (KIQ bypass + dead-GPU detection)")