path: root/kernel/gsp_rpc_client.c

/*
 * GSP RPC Client for vGPU-IO
 * Based on GSP 570.144W firmware decompilation
 * Firmware base: 0x04000000
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/mm.h>
#include <linux/dma-mapping.h>
#include <linux/mutex.h>
#include <linux/spinlock.h>
#include <linux/list.h>
#include <linux/workqueue.h>
#include <linux/delay.h>
#include <asm/io.h>

#include "nvidia-vgpu-vfio.h"
#include "memory.h"
#include "gsp_rpc_client.h"

/* ============================================================
 * FIRMWARE ADDRESSES (from decompilation)
 * ============================================================ */

#define GSP_FW_BASE                 0x04000000UL

/* Primary RPC Dispatcher */
#define GSP_RPC_DISPATCH_PRIMARY    (GSP_FW_BASE + 0x17A1AC0)  /* 0x057a1ac0 */
#define GSP_RPC_JUMP_TABLE          (GSP_FW_BASE + 0x17A3AD8)  /* 0x057a3ad8 - 32 entries */

/* Secondary RPC Dispatcher */
#define GSP_RPC_DISPATCH_SECONDARY  (GSP_FW_BASE + 0x0ED39DC)  /* 0x04ed39dc */

/* Specialized Handlers */
#define GSP_ALLOC_MAIN              (GSP_FW_BASE + 0x17E94EC)  /* 0x057e94ec */
#define GSP_SCHEDULE_MAIN           (GSP_FW_BASE + 0x11EFBF4)  /* 0x051efbf4 */
#define GSP_BAR_TRAP_HANDLER_JT     (GSP_FW_BASE + 0x17675B7)  /* 0x057675b7 */
#define GSP_BAR_TRAP_HANDLER_PRIM   (GSP_FW_BASE + 0x0CB8189)  /* 0x04cb8189 */
#define GSP_PIN_HANDLER_SEC         (GSP_FW_BASE + 0x0ED34B8)  /* 0x04ed34b8 */

/* Jump Table Handlers (from 0x059bfeb0) */
#define GSP_MAP_HANDLER_A           (GSP_FW_BASE + 0x1766AB7)  /* 0x05766ab7 */
#define GSP_MAP_HANDLER_B           (GSP_FW_BASE + 0x1778A37)  /* 0x05778a37 */
#define GSP_UNMAP_HANDLER_A         (GSP_FW_BASE + 0x17669B7)  /* 0x057669b7 */
#define GSP_UNMAP_HANDLER_B         (GSP_FW_BASE + 0x1779937)  /* 0x05779937 */
#define GSP_PIN_HANDLER_A           (GSP_FW_BASE + 0x17664B7)  /* 0x057664b7 */
#define GSP_PIN_HANDLER_B           (GSP_FW_BASE + 0x177A3B7)  /* 0x0577a3b7 */
#define GSP_UNPIN_HANDLER_A         (GSP_FW_BASE + 0x17682B7)  /* 0x057682b7 */
#define GSP_UNPIN_HANDLER_B         (GSP_FW_BASE + 0x1779FB7)  /* 0x05779fb7 */
#define GSP_SCHEDULE_HANDLER_A      (GSP_FW_BASE + 0x1767F37)  /* 0x05767f37 */
#define GSP_SCHEDULE_HANDLER_B      (GSP_FW_BASE + 0x175AEB7)  /* 0x0575aeb7 */
#define GSP_YIELD_HANDLER_A         (GSP_FW_BASE + 0x175AE37)  /* 0x0575ae37 */
#define GSP_YIELD_HANDLER_B         (GSP_FW_BASE + 0x1767337)  /* 0x05767337 */
#define GSP_NOTIFY_HANDLER_A        (GSP_FW_BASE + 0x177C8B7)  /* 0x0577c8b7 */
#define GSP_NOTIFY_HANDLER_B        (GSP_FW_BASE + 0x1767837)  /* 0x05767837 */

/* Entry point */
#define GSP_ENTRY_POINT             (GSP_FW_BASE + 0x0009000)  /* 0x04009000 */

/* ============================================================
 * RPC PROTOCOL DEFINITIONS
 * ============================================================ */

#define GSP_RPC_ALLOC       0x01
#define GSP_RPC_FREE        0x02
#define GSP_RPC_MAP         0x03
#define GSP_RPC_UNMAP       0x04
#define GSP_RPC_PIN         0x05
#define GSP_RPC_UNPIN       0x06
#define GSP_RPC_SCHEDULE    0x07
#define GSP_RPC_YIELD       0x08
#define GSP_RPC_NOTIFY      0x09
#define GSP_RPC_BAR_TRAP    0x0A
#define GSP_RPC_MIGRATE     0x0B
#define GSP_RPC_DEBUG       0xFF

/* ALLOC sub-commands (primary dispatcher jump table indices) */
#define GSP_ALLOC_SUBCMD_DEFAULT    0
#define GSP_ALLOC_SUBCMD_TYPE1      1
#define GSP_ALLOC_SUBCMD_TYPE2      3
#define GSP_ALLOC_SUBCMD_TYPE3      4
#define GSP_ALLOC_SUBCMD_TYPE4      5
#define GSP_ALLOC_SUBCMD_TYPE5      8
#define GSP_ALLOC_SUBCMD_TYPE6      10
#define GSP_ALLOC_SUBCMD_TYPE7      12
#define GSP_ALLOC_SUBCMD_TYPE8      17
#define GSP_ALLOC_SUBCMD_TYPE9      19
#define GSP_ALLOC_SUBCMD_TYPE10     20
#define GSP_ALLOC_SUBCMD_TYPE11     21
#define GSP_ALLOC_SUBCMD_TYPE12     22
#define GSP_ALLOC_SUBCMD_TYPE13     23
#define GSP_ALLOC_SUBCMD_TYPE14     24
#define GSP_ALLOC_SUBCMD_TYPE15     25
#define GSP_ALLOC_SUBCMD_TYPE16     26
#define GSP_ALLOC_SUBCMD_TYPE17     27
#define GSP_ALLOC_SUBCMD_TYPE18     28
#define GSP_ALLOC_SUBCMD_TYPE19     29
#define GSP_ALLOC_SUBCMD_TYPE20     30
#define GSP_ALLOC_SUBCMD_TYPE21     31

/* MIGRATE sub-commands */
#define GSP_MIGRATE_SUBCMD_SAVE     0x04  /* From primary dispatcher branch */
#define GSP_MIGRATE_SUBCMD_RESTORE  0x05

/* PIN sub-commands */
#define GSP_PIN_SUBCMD_DEFAULT      0x05  /* From secondary dispatcher */

/* SCHEDULE sub-commands */
#define GSP_SCHED_SUBCMD_DEFAULT    0x07

/* BAR TRAP sub-commands */
#define GSP_BAR_TRAP_SUBCMD_DEFAULT 0x0A

/* ============================================================
 * RPC MESSAGE STRUCTURES
 * ============================================================ */

/* ============================================================
 * GSP RPC CLIENT CONTEXT
 * ============================================================ */

struct gsp_client {
    struct device *dev;
    void __iomem *gsp_fw_base;    /* Mapped firmware base */
    void __iomem *rpc_doorbell;   /* Doorbell register */
    spinlock_t lock;              /* Submission lock */
    struct mutex rpc_mutex;       /* RPC serialization */
    atomic_t seq_counter;         /* Sequence ID counter */
    struct list_head pending_rpcs; /* Pending RPC list */
    wait_queue_head_t resp_wq;    /* Response wait queue */
    struct workqueue_struct *wq;  /* Workqueue for async */
    bool initialized;
};

struct gsp_pending_rpc {
    struct list_head list;
    struct gsp_rpc_msg msg;
    struct gsp_rpc_response *rsp;
    struct completion done;
    u32 timeout_ms;
};

/* ============================================================
 * LOW-LEVEL RPC SUBMISSION
 * ============================================================ */

static inline u32 gsp_next_seq(struct gsp_client *client)
{
    return atomic_inc_return(&client->seq_counter);
}

static int gsp_rpc_submit_raw(struct gsp_client *client,
                               u32 type, u32 subcmd,
                               u64 *payload, u32 seq_id,
                               struct gsp_rpc_response *rsp,
                               u32 timeout_ms)
{
    struct gsp_pending_rpc *prpc;
    u64 *dispatch_ptr;
    int ret = 0;

    prpc = kzalloc(sizeof(*prpc), GFP_KERNEL);
    if (!prpc)
        return -ENOMEM;

    prpc->msg.type = type;
    prpc->msg.subcmd = subcmd;
    if (payload)
        memcpy(prpc->msg.payload, payload, sizeof(prpc->msg.payload));
    prpc->msg.seq_id = seq_id;
    prpc->rsp = rsp;
    prpc->timeout_ms = timeout_ms;
    init_completion(&prpc->done);

    spin_lock(&client->lock);
    list_add_tail(&prpc->list, &client->pending_rpcs);

    /* Write RPC message to firmware mailbox */
    dispatch_ptr = (u64 *)(client->gsp_fw_base + 0x1000); /* Mailbox offset */
    for (int i = 0; i < 8; i++)
        writeq(((u64 *)&prpc->msg)[i], &dispatch_ptr[i]);

    /* Ring doorbell */
    writel(type, client->rpc_doorbell);
    spin_unlock(&client->lock);

    /* Wait for completion */
    if (!wait_for_completion_timeout(&prpc->done,
                                      msecs_to_jiffies(timeout_ms))) {
        ret = -ETIMEDOUT;
        goto cleanup;
    }

    if (rsp)
        memcpy(rsp, prpc->rsp, sizeof(*rsp));

cleanup:
    spin_lock(&client->lock);
    list_del(&prpc->list);
    spin_unlock(&client->lock);
    kfree(prpc);
    return ret;
}

/* ============================================================
 * HIGH-LEVEL RPC API
 * ============================================================ */

/* Memory Allocation */
int gsp_rpc_alloc(struct gsp_client *client,
                   u64 size, u64 alignment, u32 mem_type,
                   u32 flags, u64 addr_hint, u32 pid,
                   u64 *gpu_va, u64 *handle, u64 *actual_size)
{
    if (gsp_rpc_emulation_is_enabled())
        return gsp_rpc_alloc_emul(client, size, alignment, mem_type, flags, addr_hint, pid, gpu_va, handle, actual_size);

    struct gsp_alloc_req req = {0};
    struct gsp_alloc_rsp rsp = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;

    req.size = size;
    req.alignment = alignment;
    req.mem_type = mem_type;
    req.flags = flags;
    req.addr_hint = addr_hint;
    req.pid = pid;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_ALLOC, GSP_ALLOC_SUBCMD_DEFAULT,
                              (u64 *)&req, seq, &response, 5000);
    if (ret)
        return ret;

    if (response.status) {
        pr_err("GSP ALLOC failed: status=%d\n", response.status);
        return -EIO;
    }

    memcpy(&rsp, response.result, sizeof(rsp));
    if (gpu_va) *gpu_va = rsp.gpu_va;
    if (handle) *handle = rsp.handle;
    if (actual_size) *actual_size = rsp.size;

    return 0;
}

int gsp_rpc_free(struct gsp_client *client, u64 handle)
{
    if (gsp_rpc_emulation_is_enabled())
        return gsp_rpc_free_emul(client, handle);

    struct gsp_rpc_response response = {0};
    u64 payload[6] = { handle, 0, 0, 0, 0, 0 };
    u32 seq = gsp_next_seq(client);
    int ret;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_FREE, 0,
                              payload, seq, &response, 2000);
    if (ret)
        return ret;

    return response.status ? -EIO : 0;
}

/* Memory Mapping */
int gsp_rpc_map(struct gsp_client *client,
                 u64 handle, u64 gpu_va, u64 size,
                 u32 prot, u32 flags, u64 offset,
                 u64 *mapped_va)
{
    if (gsp_rpc_emulation_is_enabled())
        return gsp_rpc_map_emul(client, handle, gpu_va, size, prot, flags, offset, mapped_va);

    struct gsp_map_req req = {0};
    struct gsp_map_rsp rsp = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;

    req.handle = handle;
    req.gpu_va = gpu_va;
    req.size = size;
    req.prot = prot;
    req.flags = flags;
    req.offset = offset;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_MAP, 0,
                              (u64 *)&req, seq, &response, 2000);
    if (ret)
        return ret;

    if (response.status)
        return -EIO;

    memcpy(&rsp, response.result, sizeof(rsp));
    if (mapped_va) *mapped_va = rsp.gpu_va;

    return 0;
}

int gsp_rpc_unmap(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size)
{
    if (gsp_rpc_emulation_is_enabled())
        return gsp_rpc_unmap_emul(client, handle, gpu_va, size);

    struct gsp_rpc_response response = {0};
    u64 payload[6] = { handle, gpu_va, size, 0, 0, 0 };
    u32 seq = gsp_next_seq(client);
    int ret;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_UNMAP, 0,
                              payload, seq, &response, 2000);
    if (ret)
        return ret;

    return response.status ? -EIO : 0;
}

/* Page Pinning */
int gsp_rpc_pin(struct gsp_client *client,
                 u64 handle, u64 offset, u64 size,
                 u32 flags, u64 *dma_addr, u32 *num_pages)
{
    if (gsp_rpc_emulation_is_enabled())
        return gsp_rpc_pin_emul(client, handle, offset, size, flags, dma_addr, num_pages);

    struct gsp_pin_req req = {0};
    struct gsp_pin_rsp rsp = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;

    req.handle = handle;
    req.offset = offset;
    req.size = size;
    req.flags = flags;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_PIN, GSP_PIN_SUBCMD_DEFAULT,
                              (u64 *)&req, seq, &response, 5000);
    if (ret)
        return ret;

    if (response.status)
        return -EIO;

    memcpy(&rsp, response.result, sizeof(rsp));
    if (dma_addr) *dma_addr = rsp.dma_addr;
    if (num_pages) *num_pages = rsp.num_pages;

    return 0;
}

int gsp_rpc_unpin(struct gsp_client *client, u64 handle, u64 offset, u64 size)
{
    if (gsp_rpc_emulation_is_enabled())
        return gsp_rpc_unpin_emul(client, handle, offset, size);

    struct gsp_rpc_response response = {0};
    u64 payload[6] = { handle, offset, size, 0, 0, 0 };
    u32 seq = gsp_next_seq(client);
    int ret;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_UNPIN, 0,
                              payload, seq, &response, 2000);
    if (ret)
        return ret;

    return response.status ? -EIO : 0;
}

/* BAR Trap Emulation */
int gsp_rpc_bar_trap(struct gsp_client *client,
                      u64 bar_addr, u32 access_type, u32 size,
                      u64 value, u32 bar_num, u64 *ret_value)
{
    struct gsp_bar_trap_req req = {0};
    struct gsp_bar_trap_rsp rsp = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;

    req.bar_addr = bar_addr;
    req.access_type = access_type;
    req.size = size;
    req.value = value;
    req.bar_num = bar_num;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_BAR_TRAP, GSP_BAR_TRAP_SUBCMD_DEFAULT,
                              (u64 *)&req, seq, &response, 2000);
    if (ret)
        return ret;

    if (response.status)
        return -EIO;

    memcpy(&rsp, response.result, sizeof(rsp));
    if (ret_value) *ret_value = rsp.value;

    return rsp.handled ? 0 : -ENODEV;
}

/* Migration */
int gsp_rpc_migrate_save(struct gsp_client *client,
                          u64 src_gpu_id, u64 dst_gpu_id,
                          u64 memory_bitmap, u64 page_table_ptr,
                          u64 *checkpoint)
{
    struct gsp_migrate_req req = {0};
    struct gsp_migrate_rsp rsp = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;

    req.src_gpu_id = src_gpu_id;
    req.dst_gpu_id = dst_gpu_id;
    req.memory_bitmap = memory_bitmap;
    req.page_table_ptr = page_table_ptr;
    req.state = 0; /* SAVE */
    req.flags = GSP_MIGRATE_SUBCMD_SAVE;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_MIGRATE, GSP_MIGRATE_SUBCMD_SAVE,
                              (u64 *)&req, seq, &response, 30000);
    if (ret)
        return ret;

    if (response.status)
        return -EIO;

    memcpy(&rsp, response.result, sizeof(rsp));
    if (checkpoint) *checkpoint = rsp.checkpoint;

    return 0;
}

int gsp_rpc_migrate_restore(struct gsp_client *client,
                             u64 src_gpu_id, u64 dst_gpu_id,
                             u64 memory_bitmap, u64 page_table_ptr,
                             u64 checkpoint)
{
    struct gsp_migrate_req req = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;

    req.src_gpu_id = src_gpu_id;
    req.dst_gpu_id = dst_gpu_id;
    req.memory_bitmap = memory_bitmap;
    req.page_table_ptr = page_table_ptr;
    req.state = 1; /* RESTORE */
    req.flags = GSP_MIGRATE_SUBCMD_RESTORE;
    req.context = checkpoint;

    ret = gsp_rpc_submit_raw(client, GSP_RPC_MIGRATE, GSP_MIGRATE_SUBCMD_RESTORE,
                              (u64 *)&req, seq, &response, 30000);
    if (ret)
        return ret;

    return response.status ? -EIO : 0;
}

/* ============================================================
 * EMULATION FALLBACKS
 * ============================================================ */

static bool gsp_emulation_enabled = false;

void gsp_rpc_emulation_set(bool enable)
{
    gsp_emulation_enabled = enable;
}

bool gsp_rpc_emulation_is_enabled(void)
{
    return gsp_emulation_enabled;
}

/* VRAM allocation tracking for emulation */
struct gsp_vram_alloc {
    struct list_head list;
    u64 handle;
    u64 gpu_va;
    u64 size;
    u64 alignment;
    u32 mem_type;
    u32 flags;
    u64 addr_hint;
    u32 pid;
    bool mapped;
    u64 mapped_va;
    struct nv_vgpu_simple_mapping *mapping;
};

static LIST_HEAD(gsp_vram_allocs);
static DEFINE_SPINLOCK(gsp_vram_lock);
static atomic64_t gsp_handle_counter = ATOMIC64_INIT(1);

/* Scheduling emulation - local work queue */
struct gsp_sched_work {
    struct list_head list;
    struct work_struct work;
    u64 work_handle;
    u64 callback;
    u64 arg;
    u32 priority;
    u32 flags;
    struct completion done;
};

static LIST_HEAD(gsp_sched_queue);
static DEFINE_SPINLOCK(gsp_sched_lock);
static struct workqueue_struct *gsp_sched_wq = NULL;

static void gsp_sched_worker(struct work_struct *work)
{
    struct gsp_sched_work *sched = container_of(work, struct gsp_sched_work, work);
    
    /* In emulation mode, we just complete successfully */
    complete(&sched->done);
}

/* Find VRAM allocation by handle */
static struct gsp_vram_alloc *gsp_find_vram_alloc(u64 handle)
{
    struct gsp_vram_alloc *alloc;
    spin_lock(&gsp_vram_lock);
    list_for_each_entry(alloc, &gsp_vram_allocs, list) {
        if (alloc->handle == handle) {
            spin_unlock(&gsp_vram_lock);
            return alloc;
        }
    }
    spin_unlock(&gsp_vram_lock);
    return NULL;
}

/* Find VRAM allocation by GPU VA */
static struct gsp_vram_alloc *gsp_find_vram_by_va(u64 gpu_va)
{
    struct gsp_vram_alloc *alloc;
    spin_lock(&gsp_vram_lock);
    list_for_each_entry(alloc, &gsp_vram_allocs, list) {
        if (alloc->gpu_va <= gpu_va && gpu_va < alloc->gpu_va + alloc->size) {
            spin_unlock(&gsp_vram_lock);
            return alloc;
        }
    }
    spin_unlock(&gsp_vram_lock);
    return NULL;
}

/* Memory Allocation - Emulation Implementation */
int gsp_rpc_alloc_emul(struct gsp_client *client,
                               u64 size, u64 alignment, u32 mem_type,
                               u32 flags, u64 addr_hint, u32 pid,
                               u64 *gpu_va, u64 *handle, u64 *actual_size)
{
    struct nv_vgpu_device *vgpu = client->dev->driver_data;
    struct gsp_vram_alloc *alloc;
    u64 va = addr_hint ? addr_hint : 0x100000000ULL; /* Default 4GB */
    int ret;

    alloc = kzalloc(sizeof(*alloc), GFP_KERNEL);
    if (!alloc)
        return -ENOMEM;

    alloc->handle = atomic64_inc_return(&gsp_handle_counter);
    alloc->size = size;
    alloc->alignment = alignment;
    alloc->mem_type = mem_type;
    alloc->flags = flags;
    alloc->addr_hint = addr_hint;
    alloc->pid = pid;
    alloc->mapped = false;

    /* Align VA */
    va = ALIGN(va, alignment);

    /* In emulation, we allocate host memory and map it to GPU VA */
    alloc->mapping = kzalloc(sizeof(*alloc->mapping), GFP_KERNEL);
    if (!alloc->mapping) {
        kfree(alloc);
        return -ENOMEM;
    }

    alloc->mapping->gpa = va;
    alloc->mapping->size = size;
    alloc->mapping->page_count = (size + PAGE_SIZE - 1) / PAGE_SIZE;
    alloc->mapping->pages = vzalloc(alloc->mapping->page_count * sizeof(struct page *));
    if (!alloc->mapping->pages) {
        kfree(alloc->mapping);
        kfree(alloc);
        return -ENOMEM;
    }

    /* Allocate host pages */
    for (int i = 0; i < alloc->mapping->page_count; i++) {
        struct page *page = alloc_page(GFP_KERNEL | GFP_DMA32);
        if (!page) {
            for (int j = 0; j < i; j++)
                __free_page(alloc->mapping->pages[j]);
            vfree(alloc->mapping->pages);
            kfree(alloc->mapping);
            kfree(alloc);
            return -ENOMEM;
        }
        alloc->mapping->pages[i] = page;
    }
    alloc->mapping->pinned = true;

    /* Map to IOVA for DMA */
    ret = dma_map_page(&vgpu->pdev->dev, alloc->mapping->pages[0], 0, size, DMA_BIDIRECTIONAL);
    if (dma_mapping_error(&vgpu->pdev->dev, ret)) {
        for (int i = 0; i < alloc->mapping->page_count; i++)
            __free_page(alloc->mapping->pages[i]);
        vfree(alloc->mapping->pages);
        kfree(alloc->mapping);
        kfree(alloc);
        return -ENOMEM;
    }
    alloc->mapping->iova = ret;
    alloc->mapping->mapped = true;

    alloc->gpu_va = va;
    alloc->mapped_va = va;

    spin_lock(&gsp_vram_lock);
    list_add_tail(&alloc->list, &gsp_vram_allocs);
    spin_unlock(&gsp_vram_lock);

    if (gpu_va) *gpu_va = va;
    if (handle) *handle = alloc->handle;
    if (actual_size) *actual_size = size;

    nv_vgpu_info(vgpu, "EMUL ALLOC: handle=0x%llx gpu_va=0x%llx size=%llu iova=0x%llx",
                 alloc->handle, va, size, alloc->mapping->iova);

    return 0;
}

/* Memory Free - Emulation Implementation */
int gsp_rpc_free_emul(struct gsp_client *client, u64 handle)
{
    struct nv_vgpu_device *vgpu = client->dev->driver_data;
    struct gsp_vram_alloc *alloc;

    alloc = gsp_find_vram_alloc(handle);
    if (!alloc) {
        nv_vgpu_warn(vgpu, "EMUL FREE: handle 0x%llx not found", handle);
        return -ENOENT;
    }

    spin_lock(&gsp_vram_lock);
    list_del(&alloc->list);
    spin_unlock(&gsp_vram_lock);

    if (alloc->mapping) {
        if (alloc->mapping->mapped) {
            dma_unmap_page(&vgpu->pdev->dev, alloc->mapping->iova, alloc->size, DMA_BIDIRECTIONAL);
        }
        if (alloc->mapping->pinned) {
            for (int i = 0; i < alloc->mapping->page_count; i++) {
                if (alloc->mapping->pages[i])
                    __free_page(alloc->mapping->pages[i]);
            }
        }
        vfree(alloc->mapping->pages);
        kfree(alloc->mapping);
    }

    kfree(alloc);

    nv_vgpu_info(vgpu, "EMUL FREE: handle=0x%llx", handle);
    return 0;
}

/* Memory Map - Emulation Implementation */
int gsp_rpc_map_emul(struct gsp_client *client,
                             u64 handle, u64 gpu_va, u64 size,
                             u32 prot, u32 flags, u64 offset,
                             u64 *mapped_va)
{
    struct nv_vgpu_device *vgpu = client->dev->driver_data;
    struct gsp_vram_alloc *alloc = gsp_find_vram_alloc(handle);

    if (!alloc) {
        nv_vgpu_warn(vgpu, "EMUL MAP: handle 0x%llx not found", handle);
        return -ENOENT;
    }

    if (gpu_va == 0)
        gpu_va = alloc->gpu_va + offset;

    alloc->mapped_va = gpu_va;
    alloc->mapped = true;

    if (mapped_va)
        *mapped_va = gpu_va;

    nv_vgpu_info(vgpu, "EMUL MAP: handle=0x%llx gpu_va=0x%llx size=%llu offset=%llu",
                 handle, gpu_va, size, offset);

    return 0;
}

/* Memory Unmap - Emulation Implementation */
int gsp_rpc_unmap_emul(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size)
{
    struct nv_vgpu_device *vgpu = client->dev->driver_data;
    struct gsp_vram_alloc *alloc = gsp_find_vram_alloc(handle);

    if (!alloc) {
        nv_vgpu_warn(vgpu, "EMUL UNMAP: handle 0x%llx not found", handle);
        return -ENOENT;
    }

    alloc->mapped = false;

    nv_vgpu_info(vgpu, "EMUL UNMAP: handle=0x%llx gpu_va=0x%llx size=%llu",
                 handle, gpu_va, size);

    return 0;
}

/* Page Pin - Emulation Implementation */
int gsp_rpc_pin_emul(struct gsp_client *client,
                             u64 handle, u64 offset, u64 size,
                             u32 flags, u64 *dma_addr, u32 *num_pages)
{
    struct nv_vgpu_device *vgpu = client->dev->driver_data;
    struct gsp_vram_alloc *alloc = gsp_find_vram_alloc(handle);

    if (!alloc || !alloc->mapping) {
        nv_vgpu_warn(vgpu, "EMUL PIN: handle 0x%llx not found or no mapping", handle);
        return -ENOENT;
    }

    if (dma_addr)
        *dma_addr = alloc->mapping->iova + offset;
    if (num_pages)
        *num_pages = (size + PAGE_SIZE - 1) / PAGE_SIZE;

    nv_vgpu_info(vgpu, "EMUL PIN: handle=0x%llx offset=%llu size=%llu dma_addr=0x%llx pages=%u",
                 handle, offset, size, alloc->mapping->iova + offset, *num_pages);

    return 0;
}

/* Page Unpin - Emulation Implementation */
int gsp_rpc_unpin_emul(struct gsp_client *client, u64 handle, u64 offset, u64 size)
{
    struct nv_vgpu_device *vgpu = client->dev->driver_data;
    struct gsp_vram_alloc *alloc = gsp_find_vram_alloc(handle);

    if (!alloc) {
        nv_vgpu_warn(vgpu, "EMUL UNPIN: handle 0x%llx not found", handle);
        return -ENOENT;
    }

    nv_vgpu_info(vgpu, "EMUL UNPIN: handle=0x%llx offset=%llu size=%llu",
                 handle, offset, size);

    return 0;
}

int gsp_rpc_schedule(struct gsp_client *client,
                      u64 work_handle, u64 callback, u64 arg,
                      u32 priority, u32 flags)
{
    struct gsp_sched_req req = {0};
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;
    
    if (gsp_emulation_enabled) {
        struct gsp_sched_work *sched = kzalloc(sizeof(*sched), GFP_KERNEL);
        if (!sched)
            return -ENOMEM;
        
        INIT_WORK(&sched->work, gsp_sched_worker);
        sched->work_handle = work_handle;
        sched->callback = callback;
        sched->arg = arg;
        sched->priority = priority;
        sched->flags = flags;
        init_completion(&sched->done);
        
        spin_lock(&gsp_sched_lock);
        list_add_tail(&sched->list, &gsp_sched_queue);
        spin_unlock(&gsp_sched_lock);
        
        if (gsp_sched_wq)
            queue_work(gsp_sched_wq, &sched->work);
        
        wait_for_completion(&sched->done);
        
        spin_lock(&gsp_sched_lock);
        list_del(&sched->list);
        spin_unlock(&gsp_sched_lock);
        
        kfree(sched);
        return 0;
    }
    
    req.work_handle = work_handle;
    req.callback = callback;
    req.arg = arg;
    req.priority = priority;
    req.flags = flags;
    
    ret = gsp_rpc_submit_raw(client, GSP_RPC_SCHEDULE, GSP_SCHED_SUBCMD_DEFAULT,
                              (u64 *)&req, seq, &response, 2000);
    if (ret)
        return ret;
    
    return response.status ? -EIO : 0;
}

int gsp_rpc_yield(struct gsp_client *client)
{
    struct gsp_rpc_response response = {0};
    u32 seq = gsp_next_seq(client);
    int ret;
    
    if (gsp_emulation_enabled) {
        /* In emulation, yield is a no-op success */
        return 0;
    }
    
    ret = gsp_rpc_submit_raw(client, GSP_RPC_YIELD, 0,
                              NULL, seq, &response, 1000);
    if (ret)
        return ret;
    
    return response.status ? -EIO : 0;
}

/* Notification emulation - async eventfd support */
struct gsp_notify_cb {
    struct list_head list;
    u32 type;
    void (*callback)(struct gsp_client *, struct gsp_rpc_response *, void *);
    void *ctx;
};

static LIST_HEAD(gsp_notify_callbacks);
static DEFINE_SPINLOCK(gsp_notify_lock);

int gsp_rpc_notify_register(struct gsp_client *client,
                             u32 type,
                             void (*callback)(struct gsp_client *,
                                              struct gsp_rpc_response *,
                                              void *),
                             void *ctx)
{
    struct gsp_notify_cb *nc = kzalloc(sizeof(*nc), GFP_KERNEL);
    if (!nc)
        return -ENOMEM;
    
    nc->type = type;
    nc->callback = callback;
    nc->ctx = ctx;
    
    spin_lock(&gsp_notify_lock);
    list_add_tail(&nc->list, &gsp_notify_callbacks);
    spin_unlock(&gsp_notify_lock);
    
    return 0;
}

void gsp_rpc_notify_unregister(struct gsp_client *client, u32 type)
{
    struct gsp_notify_cb *nc, *tmp;
    
    spin_lock(&gsp_notify_lock);
    list_for_each_entry_safe(nc, tmp, &gsp_notify_callbacks, list) {
        if (nc->type == type) {
            list_del(&nc->list);
            kfree(nc);
        }
    }
    spin_unlock(&gsp_notify_lock);
}

/* Handle async notifications from GSP */
void gsp_handle_notification(struct gsp_client *client, struct gsp_rpc_msg *msg)
{
    struct gsp_notify_cb *nc;
    
    spin_lock(&gsp_notify_lock);
    list_for_each_entry(nc, &gsp_notify_callbacks, list) {
        if (nc->type == msg->type || nc->type == 0) {
            /* Copy message to response format for callback */
            struct gsp_rpc_response rsp = {0};
            rsp.status = 0;
            rsp.type = msg->type;
            rsp.seq_id = msg->seq_id;
            memcpy(rsp.result, msg->payload, sizeof(msg->payload));
            nc->callback(client, &rsp, nc->ctx);
        }
    }
    spin_unlock(&gsp_notify_lock);
}

/* ============================================================
 * CLIENT INITIALIZATION
 * ============================================================ */

struct gsp_client *gsp_client_init(struct device *dev, phys_addr_t fw_base, phys_addr_t doorbell)
{
    struct gsp_client *client;
    resource_size_t fw_size = 0x2000000; /* 32MB */

    client = kzalloc(sizeof(*client), GFP_KERNEL);
    if (!client)
        return ERR_PTR(-ENOMEM);

    client->dev = dev;
    spin_lock_init(&client->lock);
    mutex_init(&client->rpc_mutex);
    INIT_LIST_HEAD(&client->pending_rpcs);
    init_waitqueue_head(&client->resp_wq);
    atomic_set(&client->seq_counter, 0);

    client->gsp_fw_base = ioremap(fw_base, fw_size);
    if (!client->gsp_fw_base) {
        kfree(client);
        return ERR_PTR(-ENOMEM);
    }

    client->rpc_doorbell = ioremap(doorbell, 0x1000);
    if (!client->rpc_doorbell) {
        iounmap(client->gsp_fw_base);
        kfree(client);
        return ERR_PTR(-ENOMEM);
    }

    client->wq = alloc_workqueue("gsp_rpc", WQ_UNBOUND, 0);
    if (!client->wq) {
        iounmap(client->rpc_doorbell);
        iounmap(client->gsp_fw_base);
        kfree(client);
        return ERR_PTR(-ENOMEM);
    }

    /* Init emulation workqueue for scheduling */
    gsp_sched_wq = alloc_workqueue("gsp_sched", WQ_UNBOUND, 0);
    if (!gsp_sched_wq) {
        destroy_workqueue(client->wq);
        iounmap(client->rpc_doorbell);
        iounmap(client->gsp_fw_base);
        kfree(client);
        return ERR_PTR(-ENOMEM);
    }

    client->initialized = true;
    pr_info("GSP RPC client initialized: fw_base=%pa, doorbell=%pa\n", &fw_base, &doorbell);

    return client;
}

void gsp_client_fini(struct gsp_client *client)
{
    if (!client)
        return;

    if (client->wq)
        destroy_workqueue(client->wq);

    if (gsp_sched_wq) {
        destroy_workqueue(gsp_sched_wq);
        gsp_sched_wq = NULL;
    }

    if (client->rpc_doorbell)
        iounmap(client->rpc_doorbell);

    if (client->gsp_fw_base)
        iounmap(client->gsp_fw_base);

    kfree(client);
}

/* ============================================================
 * WPR MANAGEMENT (from firmware strings)
 * ============================================================ */

/* WPR mutex operations (from firmware: VPR_WPR_WRITE_MUTEX*) */
int gsp_wpr_write_mutex_acquire(struct gsp_client *client, u32 wpr_id)
{
    /* Calls firmware function at VPR_WPR_WRITE_MUTEX addresses */
    /* Implementation would use RPC_DEBUG or dedicated WPR RPC */
    (void)client;
    (void)wpr_id;
    return 0;
}

int gsp_wpr_write_mutex_release(struct gsp_client *client, u32 wpr_id)
{
    (void)client;
    (void)wpr_id;
    return 0;
}

int gsp_wpr_scratch_mutex_acquire(struct gsp_client *client, u32 wpr_id)
{
    /* From firmware: VPR_WPR_SCRATCH_MUTEX */
    (void)client;
    (void)wpr_id;
    return 0;
}

int gsp_wpr_scratch_mutex_release(struct gsp_client *client, u32 wpr_id)
{
    (void)client;
    (void)wpr_id;
    return 0;
}

/* ============================================================
 * MODULE EXPORTS
 * ============================================================ */

EXPORT_SYMBOL(gsp_client_init);
EXPORT_SYMBOL(gsp_client_fini);
EXPORT_SYMBOL(gsp_rpc_alloc);
EXPORT_SYMBOL(gsp_rpc_free);
EXPORT_SYMBOL(gsp_rpc_map);
EXPORT_SYMBOL(gsp_rpc_unmap);
EXPORT_SYMBOL(gsp_rpc_pin);
EXPORT_SYMBOL(gsp_rpc_unpin);
EXPORT_SYMBOL(gsp_rpc_schedule);
EXPORT_SYMBOL(gsp_rpc_yield);
EXPORT_SYMBOL(gsp_rpc_bar_trap);
EXPORT_SYMBOL(gsp_rpc_migrate_save);
EXPORT_SYMBOL(gsp_rpc_migrate_restore);
EXPORT_SYMBOL(gsp_wpr_write_mutex_acquire);
EXPORT_SYMBOL(gsp_wpr_write_mutex_release);
EXPORT_SYMBOL(gsp_wpr_scratch_mutex_acquire);
EXPORT_SYMBOL(gsp_wpr_scratch_mutex_release);
EXPORT_SYMBOL(gsp_rpc_emulation_set);
EXPORT_SYMBOL(gsp_rpc_emulation_is_enabled);
EXPORT_SYMBOL(gsp_rpc_notify_register);
EXPORT_SYMBOL(gsp_rpc_notify_unregister);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("vGPU-IO Team");
MODULE_DESCRIPTION("GSP 570.144W RPC Client for vGPU-IO");