1 files changed, 21 insertions, 29 deletions

diff --git a/kernel/gsp_emulation.c b/kernel/gsp_emulation.c
index fe1e91e..e41f3bd 100644
--- a/kernel/gsp_emulation.c
+++ b/kernel/gsp_emulation.c
@@ -296,13 +296,14 @@ static int gsp_emu_thread_fn(void *data)
             else
                 ctx->processed_rpcs++;
         } else {
-            /* Adaptive polling: busy-wait first, then sleep */
+            /* Adaptive polling: busy-wait first, then sleep (interruptible for kthread_stop) */
             static int poll_count = 0;
             if (poll_count < GSP_EMU_POLL_BUSY_US) {
                 udelay(1);
                 poll_count++;
             } else {
-                usleep_range(GSP_EMU_POLL_SLEEP_MIN_US, GSP_EMU_POLL_SLEEP_MAX_US);
+                /* Use interruptible sleep so kthread_stop can wake us */
+                schedule_timeout_interruptible(usecs_to_jiffies(GSP_EMU_POLL_SLEEP_MAX_US));
                 poll_count = 0;
             }
         }
@@ -310,6 +311,7 @@ static int gsp_emu_thread_fn(void *data)
 
     nv_vgpu_info(ctx->vgpu, "GSP emulation thread stopped (processed=%llu, failed=%llu)",
                  ctx->processed_rpcs, ctx->failed_rpcs);
+    complete(&ctx->thread_stopped);
     return 0;
 }
 
@@ -320,7 +322,6 @@ static int gsp_emu_thread_fn(void *data)
 int gsp_emu_init(struct nv_vgpu_device *vgpu)
 {
     struct gsp_emu_context *ctx;
-    phys_addr_t bar0_phys, doorbell_phys, shared_mem_phys;
     int ret = 0;
 
     nv_vgpu_info(vgpu, "Initializing GSP emulation engine");
@@ -332,22 +333,18 @@ int gsp_emu_init(struct nv_vgpu_device *vgpu)
     ctx->vgpu = vgpu;
     ctx->running = false;
     ctx->initialized = false;
+    ctx->stopping = false;
+    init_completion(&ctx->thread_stopped);
     spin_lock_init(&ctx->lock);
 
-    /* Get BAR0 physical address and map it */
-    bar0_phys = vgpu->bars[0].gpa;
-    if (!bar0_phys) {
-        nv_vgpu_err(vgpu, "BAR0 physical address not available");
+    /* Use already-mapped BAR0 virtual address (mapped by nv_vgpu_bar_init) */
+    if (!vgpu->bars[0].vaddr) {
+        nv_vgpu_err(vgpu, "BAR0 virtual address not available");
         ret = -ENODEV;
         goto err_free_ctx;
     }
 
-    ctx->bar0_vaddr = ioremap(bar0_phys, NV_VGPU_BAR0_SIZE);
-    if (!ctx->bar0_vaddr) {
-        nv_vgpu_err(vgpu, "Failed to map BAR0");
-        ret = -ENOMEM;
-        goto err_free_ctx;
-    }
+    ctx->bar0_vaddr = vgpu->bars[0].vaddr;
 
     /* Locate doorbell register at BAR0+0x110C00 */
     ctx->doorbell_reg = ctx->bar0_vaddr + GSP_EMU_DOORBELL_OFFSET;
@@ -359,19 +356,14 @@ int gsp_emu_init(struct nv_vgpu_device *vgpu)
     if (!ctx->shared_mem_virt) {
         nv_vgpu_err(vgpu, "Failed to allocate shared memory");
         ret = -ENOMEM;
-        goto err_unmap_bar0;
+        goto err_free_ctx;
     }
 
     /* Map shared memory into guest BAR0 at offset 0x1000 */
-    shared_mem_phys = ctx->shared_mem_dma;
-    /* Note: In real implementation, this would be mapped via IOMMU to guest GPA.
-     * For now, we use the host virtual address directly since guest accesses
-     * go through VFIO BAR mapping. */
     ctx->shared_mem = ctx->bar0_vaddr + GSP_EMU_SHARED_MEM_OFFSET;
 
-    /* Initialize RPC client with firmware base and doorbell */
-    doorbell_phys = bar0_phys + GSP_EMU_DOORBELL_OFFSET;
-    ctx->rpc_client = gsp_client_init(&vgpu->dev, bar0_phys, doorbell_phys);
+    /* Initialize RPC client with already-mapped BAR0 virtual address and doorbell */
+    ctx->rpc_client = gsp_client_init_emulation(&vgpu->dev, ctx->bar0_vaddr, ctx->doorbell_reg);
     if (IS_ERR(ctx->rpc_client)) {
         ret = PTR_ERR(ctx->rpc_client);
         nv_vgpu_err(vgpu, "Failed to init GSP RPC client: %d", ret);
@@ -406,8 +398,6 @@ err_fini_client:
 err_free_shared_mem:
     dma_free_wc(&vgpu->pdev->dev, ctx->shared_mem_size,
                 ctx->shared_mem_virt, ctx->shared_mem_dma);
-err_unmap_bar0:
-    iounmap(ctx->bar0_vaddr);
 err_free_ctx:
     kfree(ctx);
     return ret;
@@ -422,9 +412,14 @@ void gsp_emu_cleanup(struct nv_vgpu_device *vgpu)
 
     nv_vgpu_info(vgpu, "Cleaning up GSP emulation engine");
 
-    /* Stop emulation thread */
+    /* Stop emulation thread with timeout */
     if (ctx->thread) {
+        ctx->stopping = true;
         kthread_stop(ctx->thread);
+        /* Wait for thread to actually stop with timeout (5 seconds) */
+        if (!wait_for_completion_timeout(&ctx->thread_stopped, msecs_to_jiffies(5000))) {
+            nv_vgpu_warn(vgpu, "GSP emulation thread did not stop in time, forcing cleanup");
+        }
         ctx->thread = NULL;
     }
     ctx->running = false;
@@ -445,11 +440,8 @@ void gsp_emu_cleanup(struct nv_vgpu_device *vgpu)
         ctx->shared_mem_virt = NULL;
     }
 
-    /* Unmap BAR0 */
-    if (ctx->bar0_vaddr) {
-        iounmap(ctx->bar0_vaddr);
-        ctx->bar0_vaddr = NULL;
-    }
+    /* BAR0 is managed by nv_vgpu_bar_init/cleanup, don't unmap here */
+    ctx->bar0_vaddr = NULL;
 
     vgpu->gsp_comm_ctx = NULL;
     kfree(ctx);