/* wpr2_reset_kmod - kernel module that clears GPU WPR2 registers via BAR0 ioremap.
 *
 * When nvidia.ko is unloaded without a clean GSP teardown, GSP firmware stays
 * resident in WPR2 (a write-protected FB region described by
 * PFB_PRI_MMU_WPR2_ADDR_LO/HI). Every subsequent RmInitAdapter then fails with
 * "unexpected WPR2 already up" and FLR does not help (GSP survives warm resets).
 *
 * This module ioremaps GPU BAR0 and pokes WPR2_ADDR_LO/HI to zero so the next
 * driver load takes the normal (clean) GSP boot path. Equivalent to the effect
 * of FWSEC-SB + Booter Unload teardown, minus running the ucodes.
 *
 * Usage: sudo insmod wpr2_reset_kmod.ko [pci_bdf=0000:01:00.0]
 *        result printed via pr_info (check dmesg)
 */
#include <linux/module.h>
#include <linux/pci.h>
#include <linux/io.h>

#define NV_PFB_PRI_MMU_WPR2_ADDR_LO 0x1fa824
#define NV_PFB_PRI_MMU_WPR2_ADDR_HI 0x1fa828

static char *pci_bdf = "0000:01:00.0";
module_param(pci_bdf, charp, 0444);
MODULE_PARM_DESC(pci_bdf, "PCI BDF of the GPU (default 0000:01:00.0)");

static int __init wpr2_reset_init(void)
{
    struct pci_dev *pdev = NULL;
    unsigned int bus, dev, fn;
    void __iomem *bar0;
    u32 lo_before, hi_before, lo_after, hi_after;

    if (sscanf(pci_bdf, "0000:%x:%x.%u", &bus, &dev, &fn) != 3) {
        pr_err("wpr2_reset: bad pci_bdf %s\n", pci_bdf);
        return -EINVAL;
    }
    pdev = pci_get_domain_bus_and_slot(0, bus, PCI_DEVFN(dev, fn));
    if (!pdev) {
        pr_err("wpr2_reset: device %s not found\n", pci_bdf);
        return -ENODEV;
    }
    if (!pci_resource_start(pdev, 0)) {
        pr_err("wpr2_reset: BAR0 not assigned\n");
        pci_dev_put(pdev);
        return -ENODEV;
    }

    if (pci_enable_device(pdev)) {
        pr_err("wpr2_reset: pci_enable_device failed\n");
        pci_dev_put(pdev);
        return -EIO;
    }
    pci_set_master(pdev);
    bar0 = ioremap_wc(pci_resource_start(pdev, 0), pci_resource_len(pdev, 0));
    if (!bar0) {
        pr_err("wpr2_reset: ioremap failed\n");
        pci_dev_put(pdev);
        return -ENOMEM;
    }

    lo_before = readl(bar0 + NV_PFB_PRI_MMU_WPR2_ADDR_LO);
    hi_before = readl(bar0 + NV_PFB_PRI_MMU_WPR2_ADDR_HI);
    pr_info("wpr2_reset: before LO=0x%08x HI=0x%08x\n", lo_before, hi_before);

    writel(0, bar0 + NV_PFB_PRI_MMU_WPR2_ADDR_LO);
    writel(0, bar0 + NV_PFB_PRI_MMU_WPR2_ADDR_HI);

    lo_after = readl(bar0 + NV_PFB_PRI_MMU_WPR2_ADDR_LO);
    hi_after = readl(bar0 + NV_PFB_PRI_MMU_WPR2_ADDR_HI);
    pr_info("wpr2_reset: after  LO=0x%08x HI=0x%08x\n", lo_after, hi_after);

    if (hi_after == 0)
        pr_info("wpr2_reset: WPR2 cleared, clean GSP boot now possible\n");
    else
        pr_info("wpr2_reset: write dropped (PLM?) - WPR2 still up\n");

    iounmap(bar0);
    pci_dev_put(pdev);
    return hi_after ? -EINVAL : 0;
}

static void __exit wpr2_reset_exit(void)
{
    pr_info("wpr2_reset: unloaded\n");
}

module_init(wpr2_reset_init);
module_exit(wpr2_reset_exit);
MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Clear GPU WPR2 registers to unstick GSP boot");