1 files changed, 351 insertions, 22 deletions

diff --git a/kernel/gsp_sema_service.c b/kernel/gsp_sema_service.c
index 930118e..9edb66d 100644
--- a/kernel/gsp_sema_service.c
+++ b/kernel/gsp_sema_service.c
@@ -142,6 +142,10 @@ struct svc_channel {
 	u32	logged_expected;	/* last pool value logged (rate limit) */
 	unsigned long rescue_last;	/* rate limit for the full-RAM rescue scan */
 	u32	rescue_fail;	/* consecutive failed rescues for the entry */
+	u64	counter_pa;	/* guest PA of the u32 GPU completion counter
+				 * (registered via the sema_counter debug param);
+				 * 0 = not registered */
+	u32	counter_last;	/* last value published to counter_pa */
 };
 
 /*
@@ -1069,6 +1073,145 @@ static void svc_pin_store(u64 va, u64 pa, int ap)
 	}
 }
 
+static void svc_add_map(u64 gpu_va, u64 pa);
+
+/*
+ * Targeted resolve: when a VA lookup misses in the discovered map, scan guest
+ * RAM specifically for an rm_mem record containing this GPU VA. Only records
+ * that are still live carry a valid header, so this finds the *current*
+ * translation for a VA the content-scanned map lost.
+ *
+ * The scan is sliced and only one VA is scanned at a time: a full-RAM scan in
+ * one call blocks the service thread for seconds and every GPFIFO pump and
+ * release application queues behind it. A caller that hits while another VA
+ * owns the scan gets a plain miss; its scan starts once the active one
+ * finishes. VAs whose full scan found nothing are remembered and not rescanned
+ * for SVC_TR_FAIL_HOLD.
+ */
+#define SVC_TR_SLICE	(128ULL << 20)	/* guest RAM covered per call */
+#define SVC_TR_CHUNK	(64 * 1024)	/* one guest read */
+#define SVC_TR_FAIL_MAX	16
+#define SVC_TR_FAIL_HOLD	(60 * HZ)
+
+static u64 tr_va;		/* page VA of the active scan, 0 = idle */
+static u64 tr_cursor;		/* next GPA to read */
+static unsigned long tr_start;	/* active scan start (staleness bound) */
+static u64 tr_scans;
+static struct {
+	u64 va;
+	unsigned long at;
+} tr_failed[SVC_TR_FAIL_MAX];
+static u32 tr_fail_next;
+
+static bool svc_tr_recently_failed(u64 page_va)
+{
+	u32 i;
+
+	for (i = 0; i < SVC_TR_FAIL_MAX; i++)
+		if (tr_failed[i].va == page_va &&
+		    time_before(jiffies, tr_failed[i].at + SVC_TR_FAIL_HOLD))
+			return true;
+	return false;
+}
+
+static void svc_tr_note_failure(u64 page_va)
+{
+	tr_failed[tr_fail_next].va = page_va;
+	tr_failed[tr_fail_next].at = jiffies;
+	tr_fail_next = (tr_fail_next + 1) % SVC_TR_FAIL_MAX;
+}
+
+static int svc_targeted_resolve(struct nv_vgpu_device *vgpu, u64 gpu_va,
+				u64 *pa_out)
+{
+	u8 *buf;
+	u64 gpa, end;
+	u64 page_va = gpu_va & ~(u64)(SEMA_PAGE_SIZE - 1);
+	u64 direct_base = svc.direct_base;
+
+	if (!direct_base)
+		return -ENOENT;
+	if (svc_tr_recently_failed(page_va))
+		return -ENOENT;
+
+	if (tr_va && tr_va != page_va)
+		return -EAGAIN;	/* another VA owns the active scan */
+	if (!tr_va) {
+		tr_va = page_va;
+		tr_cursor = 0;
+		tr_start = jiffies;
+	}
+	if (time_after(jiffies, tr_start + 30 * HZ)) {
+		/* The guest stopped giving us its memory: abandon the scan. */
+		svc_tr_note_failure(tr_va);
+		tr_va = 0;
+		return -ENOENT;
+	}
+
+	buf = kmalloc(SVC_TR_CHUNK, GFP_KERNEL);
+	if (!buf)
+		return -ENOMEM;
+
+	end = tr_cursor + SVC_TR_SLICE;
+	if (end > SEMA_PA_MAX)
+		end = SEMA_PA_MAX;
+
+	for (gpa = tr_cursor; gpa < end; gpa += SVC_TR_CHUNK) {
+		u64 chunk = min_t(u64, SVC_TR_CHUNK, end - gpa);
+		u32 off;
+
+		if (svc_read(vgpu, gpa, buf, chunk))
+			continue;
+
+		for (off = 48; off + 16 <= chunk; off += 8) {
+			u64 cpu_va, found_gpu_va, pa;
+			u32 type;
+			u64 mapped_on;
+
+			memcpy(&cpu_va, buf + off, 8);
+			memcpy(&found_gpu_va, buf + off + 8, 8);
+			if (cpu_va < KVA_LO || cpu_va >= KVA_HI)
+				continue;
+			if (found_gpu_va != page_va)
+				continue;
+			/* Validate the rm_mem header. */
+			memcpy(&type, buf + off - 48, 4);
+			memcpy(&mapped_on, buf + off - 40, 8);
+			if ((type != 0 && type != 1) ||
+			    (mapped_on & 0x3) != 0x3)
+				continue;
+			if (cpu_va < direct_base)
+				continue;
+			pa = cpu_va - direct_base;
+			if (pa >= SEMA_PA_MAX || (pa & 0xfff))
+				continue;
+
+			*pa_out = pa | (gpu_va & (SEMA_PAGE_SIZE - 1));
+			svc_pin_store(page_va, pa, SVC_APERTURE_SYS);
+			svc_lat_store(page_va, pa, SVC_APERTURE_SYS);
+			tr_va = 0;
+			tr_scans++;
+			pr_info("sema-svc: targeted resolve va=0x%llx -> pa=0x%llx (gpa=0x%llx off=%u, scan #%llu)\n",
+				page_va, pa, gpa, off, tr_scans);
+			kfree(buf);
+			return 0;
+		}
+	}
+
+	tr_cursor = end;
+	kfree(buf);
+	if (tr_cursor >= SEMA_PA_MAX) {
+		pr_info("sema-svc: targeted resolve va=0x%llx: full RAM scan found nothing (scan #%llu)\n",
+			page_va, tr_scans);
+		svc_tr_note_failure(page_va);
+		tr_va = 0;
+		return -ENOENT;
+	}
+	/* Not found yet; the caller treats it as a miss and the next call
+	 * continues the scan from tr_cursor. */
+	return -EAGAIN;
+}
+
 /*
  * Translate a GPU VA using trusted sources only: a pinned translation, the
  * learned cache, the discovered map, the directories the driver published and
@@ -1111,6 +1254,13 @@ static int svc_lookup_trusted(struct nv_vgpu_device *vgpu, u64 va, u64 *addr,
 			return 0;
 		}
 	}
+	/* Targeted resolve: scan RAM for the specific rm_mem record with this VA. */
+	if (!svc_targeted_resolve(vgpu, va, &pa)) {
+		*addr = pa;
+		*ap = SVC_APERTURE_SYS;
+		svc_lat_store(va, pa, SVC_APERTURE_SYS);
+		return 0;
+	}
 	return -ENOENT;
 }
 
@@ -1148,6 +1298,14 @@ static int svc_resolve_va(struct nv_vgpu_device *vgpu, u64 va, u64 *addr,
 	 */
 	if (svc_lookup_trusted(vgpu, va, addr, &ap))
 		return -ENOENT;
+	/*
+	 * Pin the resolved page: the discovered map is rebuilt wholesale on a
+	 * rolling cadence and its yield swings with guest activity, which can
+	 * drop a live semaphore entry (observed: probe NO ENTRY for the pool
+	 * page after a rebuild). Pins survive rebuilds, so a page that was
+	 * resolvable stays resolvable for the boot.
+	 */
+	svc_pin_store(va, *addr, ap);
 	*aperture = (ap == SVC_APERTURE_FB) ? 1 : 0;
 	return 0;
 }
@@ -1288,10 +1446,36 @@ static void svc_add_map(u64 gpu_va, u64 pa)
 	}
 }
 
-/* Add one scanned {CPU KVA, GPU VA} pair to the map. */
-static void svc_add_pair(u64 base, const struct raw_pair *p)
+/* Index of a GPU VA page in a sorted map sub-range, -1 if absent. */
+static int svc_map_find_range(u64 page, u32 to)
+{
+	int lo = 0, hi = (int)to - 1;
+
+	while (lo <= hi) {
+		int mid = (lo + hi) / 2;
+
+		if (svc.map[mid].gpu_page == page)
+			return mid;
+		if (svc.map[mid].gpu_page < page)
+			lo = mid + 1;
+		else
+			hi = mid - 1;
+	}
+	return -1;
+}
+
+/*
+ * Merge one scanned pair into the map: a VA already present in the sorted
+ * range gets its translation updated in place (the fresh scan wins), a new VA
+ * is appended. Used by svc_rebuild so that rebuilds extend the map instead of
+ * replacing it wholesale: the scan yield collapses with guest activity
+ * (observed 27811 -> 314 pairs after a context teardown) and replacing the map
+ * dropped live translations that the pump and the release path still needed.
+ */
+static void svc_add_pair_merge(u64 base, const struct raw_pair *p, u32 sorted_n)
 {
-	u64 pa;
+	u64 pa, page;
+	int idx;
 
 	if (p->cpu_va < base)
 		return;
@@ -1300,7 +1484,15 @@ static void svc_add_pair(u64 base, const struct raw_pair *p)
 		return;
 	if (p->gpu_va < SEMA_VA_LO || p->gpu_va >= SEMA_VA_HI)
 		return;
-	svc_add_map(p->gpu_va, pa);
+
+	page = p->gpu_va & ~(u64)(SEMA_PAGE_SIZE - 1);
+	idx = svc_map_find_range(page, sorted_n);
+	if (idx >= 0) {
+		svc.map[idx].pa_page = pa & ~(u64)(SEMA_PAGE_SIZE - 1);
+		svc.pairs_seen++;
+		return;
+	}
+	svc_add_map(page, pa);
 	svc.pairs_seen++;
 }
 
@@ -1405,6 +1597,7 @@ static void svc_rebuild(struct nv_vgpu_device *vgpu)
 	u32 ncand = 0, npairs = 0, i;
 	u64 best_base = 0;
 	u32 best_count = 0;
+	u32 old_n;
 
 	if (!buf) {
 		buf = kmalloc(SEMA_PAGE_SIZE, GFP_KERNEL);
@@ -1412,7 +1605,19 @@ static void svc_rebuild(struct nv_vgpu_device *vgpu)
 			return;
 	}
 
-	svc.map_n = 0;
+	/*
+	 * Merge across rebuilds: keep the previous entries and let the new
+	 * scan update a VA it re-finds (the fresh translation wins) or append
+	 * VAs it did not know. The scan yield swings with guest activity
+	 * (observed 27811 -> 314 pairs when a context tears down its
+	 * allocations), and replacing the map wholesale dropped live
+	 * translations. When the map is nearly full, compact once from
+	 * scratch to make room for new VAs.
+	 */
+	old_n = svc.map_n;
+	if (old_n > SEMA_MAP_MAX - SEMA_MAP_MAX / 4)
+		old_n = svc.map_n = 0;
+
 	svc_scan_pass(vgpu, buf, cands, &ncand, pairs, &npairs);
 
 	for (i = 0; i < ncand; i++) {
@@ -1428,17 +1633,20 @@ static void svc_rebuild(struct nv_vgpu_device *vgpu)
 		svc.direct_base = best_base;
 
 		/*
-		 * Live rm_mem-style records first (dedup is first-wins), then
-		 * headerless candidates only for VAs the validated set does not
-		 * cover - a freed allocation or an unrelated table holding a
-		 * stale pair for the same VA must not shadow the live record.
+		 * Live rm_mem-style records first (merge updates old entries
+		 * in place, new VAs are appended), then headerless candidates
+		 * only for VAs the validated set does not cover - a freed
+		 * allocation or an unrelated table holding a stale pair for the
+		 * same VA must not shadow the live record.
 		 */
 		for (i = 0; i < npairs; i++) {
 			if (pairs[i].valid_hdr)
-				svc_add_pair(best_base, &pairs[i]);
+				svc_add_pair_merge(best_base, &pairs[i],
+						   old_n);
 		}
+		/* The old range was sorted; the merge appends an unsorted tail. */
+		sort(svc.map, svc.map_n, sizeof(svc.map[0]), svc_map_cmp, NULL);
 		n_valid = svc.map_n;
-		sort(svc.map, n_valid, sizeof(svc.map[0]), svc_map_cmp, NULL);
 		for (i = 0; i < npairs; i++) {
 			u64 dummy;
 
@@ -1447,12 +1655,12 @@ static void svc_rebuild(struct nv_vgpu_device *vgpu)
 			if (!svc_map_lookup_range(pairs[i].gpu_va, 0, n_valid,
 						 &dummy))
 				continue;
-			svc_add_pair(best_base, &pairs[i]);
+			svc_add_pair_merge(best_base, &pairs[i], n_valid);
 		}
 	}
 
-	/* Stale entries from a previous rebuild (wrong PA for a reused VA)
-	 * must not outlive the map they were derived from. */
+	/* The learned cache is a cache: a stale entry must not shadow the
+	 * merged map's fresh translation for a reused VA. */
 	memset(svc.lat, 0, sizeof(svc.lat));
 
 	svc_map_finalize();
@@ -1739,6 +1947,25 @@ static bool svc_method_stream_plausible(const u8 *buf, u32 len)
 }
 
 /*
+ * Relaxed push-prefix check for rescue candidate validation only. Real pushes
+ * observed in this workload start with methods at non-zero addresses (e.g.
+ * 0x2001001e = count 1, addr 0x78), which the strict check above rejects
+ * because it requires addr 0. Rescue only needs the candidate page to begin
+ * a plausible method stream; work execution still uses the strict check.
+ */
+static bool svc_stream_prefix_ok(const u8 *buf, u32 len)
+{
+	u32 hdr = 0, op, count;
+
+	if (len < 8)
+		return false;
+	memcpy(&hdr, buf, 4);
+	op = (hdr >> 29) & 0x7;
+	count = (hdr >> 16) & 0xfff;
+	return op == 1 && count && count <= 64 && 4 + count * 4 <= len;
+}
+
+/*
  * Decode one page as a method stream and execute the pushes it contains.
  *
  * Push buffers are written back to back and pushes can span page
@@ -1753,7 +1980,12 @@ static void svc_decode_stream(struct nv_vgpu_device *vgpu, const u8 *buf,
 			      u32 ch_h, u32 ch_eng, u64 push_va)
 {
 	struct svc_push_state st;
-	struct svc_memset_op ops[SVC_MAX_MEMSETS];
+	/*
+	 * Static for the same reason as svc_handle_chan_dump: a 32-entry ops
+	 * array is a ~2 KB frame, and this is called from the single service
+	 * thread only.
+	 */
+	static struct svc_memset_op ops[SVC_MAX_MEMSETS];
 	u32 n_ops = 0;
 	u32 ops_runs = 0;
 	u32 off = 0;
@@ -1886,9 +2118,10 @@ static void svc_decode_stream(struct nv_vgpu_device *vgpu, const u8 *buf,
 					svc.pushes_pending++;
 					cur = (op == HOST_SEM_OP_REDUCTION) ?
 						cur + payload : payload;
-					if (!svc_sem_write(vgpu, va, cur))
+					if (!svc_sem_write(vgpu, va, cur)) {
 						svc.releases_applied++;
 						svc.post_hint_until = jiffies + 10 * HZ;
+					}
 				}
 			}
 			off += 4 + count * 4;
@@ -1968,9 +2201,10 @@ static void svc_decode_stream(struct nv_vgpu_device *vgpu, const u8 *buf,
 						break;
 					svc.pushes_pending++;
 					if (!svc_sem_write(vgpu, va,
-							   st.rs_payload))
+							   st.rs_payload)) {
 						svc.releases_applied++;
 						svc.post_hint_until = jiffies + 10 * HZ;
+					}
 					break;
 				}
 				default:
@@ -2088,7 +2322,7 @@ static void svc_decode_stream(struct nv_vgpu_device *vgpu, const u8 *buf,
 
 			if (DMA_SEMAPHORE_TYPE(dma) ==
 			    SEM_RELEASE_ONE_WORD) {
-				u64 sem_va, sem_pa = 0;
+				u64 sem_va;
 				u32 cur = 0;
 				int pending = 0;
 
@@ -2133,9 +2367,10 @@ static void svc_decode_stream(struct nv_vgpu_device *vgpu, const u8 *buf,
 						cur += st.sem_payload;
 					else
 						cur = st.sem_payload;
-					if (!svc_sem_write(vgpu, sem_va, cur))
+					if (!svc_sem_write(vgpu, sem_va, cur)) {
 						svc.releases_applied++;
 						svc.post_hint_until = jiffies + 10 * HZ;
+					}
 				}
 				/* release ends the push */
 				n_ops = 0;
@@ -3200,7 +3435,7 @@ static int svc_rescue_push_page(struct nv_vgpu_device *vgpu,
 		}
 		if (!svc_read_va_trusted(vgpu, entry_va, probe,
 					 sizeof(probe)) &&
-		    svc_method_stream_plausible(probe, sizeof(probe))) {
+		    svc_stream_prefix_ok(probe, sizeof(probe))) {
 			for (j = 0; j < pin_len; j += SEMA_PAGE_SIZE)
 				svc_mark_hot(cand[i] + j);
 			pr_info("sema-svc: content rescue ch=0x%x rank=%u slot=%u entry_va=0x%llx len=%u -> pa=0x%llx (cand %u/%u)\n",
@@ -3249,6 +3484,33 @@ static void svc_init_expected_payload(struct nv_vgpu_device *vgpu,
 	}
 }
 
+/*
+ * Publish the channel's GPU completion counter the way the real GPU would:
+ * store the channel's current GPFIFO put position into the counter's guest
+ * PA. The client's tracker polls this u32 and completes when the value
+ * reaches its target, so a channel whose pushes we cannot read (GR
+ * context-init pushes) still advances the client. Registered by hand through
+ * the sema_counter debug parameter until in-kernel derivation lands.
+ */
+static void svc_counter_publish(struct nv_vgpu_device *vgpu,
+				struct svc_channel *ch, u32 put)
+{
+	u32 val;
+
+	if (!ch->counter_pa)
+		return;
+
+	val = put;
+	if (val == ch->counter_last)
+		return;
+	if (svc_write_ap(vgpu, ch->counter_pa, SVC_APERTURE_SYS,
+			 &val, sizeof(val)))
+		return;
+	ch->counter_last = val;
+	pr_info_ratelimited("sema-chan: counter publish h=0x%x pa=0x%llx val=%u\n",
+			    ch->h_object, ch->counter_pa, val);
+}
+
 static void svc_pump_channels(struct nv_vgpu_device *vgpu)
 {
 	/* Only ever called from the service thread. */
@@ -3465,6 +3727,11 @@ static void svc_pump_channels(struct nv_vgpu_device *vgpu)
 							ch->last_get = (idx + 1) %
 								       ch->gpfifo_entries;
 							svc.gpfifo_skipped++;
+							/* The entry carried a release we
+							 * could not apply; run the fast
+							 * content sweep so the client's
+							 * wait still clears. */
+							stalled = true;
 							kfree(pb);
 							continue;
 						} else {
@@ -3485,6 +3752,9 @@ static void svc_pump_channels(struct nv_vgpu_device *vgpu)
 				     ch->userd_as, &g, sizeof(g));
 		}
 
+		/* Advance the client-visible completion counter. */
+		svc_counter_publish(vgpu, ch, put);
+
 		spin_lock_irqsave(&svc.chan_lock, flags);
 		for (k = 0; k < SVC_MAX_CHANNELS; k++) {
 			if (svc.chan[k].valid &&
@@ -3497,6 +3767,20 @@ static void svc_pump_channels(struct nv_vgpu_device *vgpu)
 				svc.chan[k].notified = ch->notified;
 				svc.chan[k].last_get = ch->last_get;
 				svc.chan[k].submissions = ch->submissions;
+				svc.chan[k].counter_last = ch->counter_last;
+				/*
+				 * Persist the per-entry rescue state and the
+				 * expected payload: the pump works on a snapshot
+				 * copy, so without the write-back the stale-entry
+				 * skip (rescue_fail >= 3) and the 5 s rescue rate
+				 * limit (rescue_last) can never fire and the
+				 * payload expectation is reinitialized every
+				 * pass.
+				 */
+				svc.chan[k].expected_sem_payload =
+					ch->expected_sem_payload;
+				svc.chan[k].rescue_last = ch->rescue_last;
+				svc.chan[k].rescue_fail = ch->rescue_fail;
 				break;
 			}
 		}
@@ -3683,7 +3967,14 @@ static bool chan_dump_pending;
 
 static void svc_handle_chan_dump(struct nv_vgpu_device *vgpu)
 {
-	struct svc_channel local[SVC_MAX_CHANNELS];
+	/*
+	 * Snapshot buffer is static: 128 channels x ~145 bytes is an 18 KB
+	 * stack frame, which overflows the 16 KB kernel stack in the
+	 * function prologue and panics the host. Only the single service
+	 * thread calls this function, so static is safe (same pattern as
+	 * svc_pump_channels).
+	 */
+	static struct svc_channel local[SVC_MAX_CHANNELS];
 	unsigned long flags;
 	int i, n = 0;
 
@@ -3833,6 +4124,43 @@ int gsp_sema_service_add_map(u64 gpu_va, u64 pa, u64 len, u32 ap)
 	return 0;
 }
 
+/*
+ * Register a channel's GPU completion counter (guest PA of the u32 the GPU
+ * would advance as pushes complete). Called from the sema_counter debug
+ * parameter; plain state update, so it runs in the caller's context.
+ */
+int gsp_sema_service_counter_add(u32 h_object, u64 pa)
+{
+	unsigned long flags;
+	int i, ret = -ENOENT;
+
+	if (!svc.task)
+		return -ENODEV;
+	/* The counter is a u32: require a 4-byte aligned guest PA. */
+	if (!h_object || !pa || (pa & 0x3) || pa >= SEMA_PA_MAX)
+		return -EINVAL;
+
+	spin_lock_irqsave(&svc.chan_lock, flags);
+	for (i = 0; i < SVC_MAX_CHANNELS; i++) {
+		if (svc.chan[i].valid &&
+		    svc.chan[i].h_object == h_object) {
+			svc.chan[i].counter_pa = pa;
+			svc.chan[i].counter_last = 0;
+			ret = 0;
+			break;
+		}
+	}
+	spin_unlock_irqrestore(&svc.chan_lock, flags);
+
+	if (!ret)
+		pr_info("sema-svc: channel 0x%x completion counter registered pa=0x%llx\n",
+			h_object, pa);
+	else
+		pr_warn_ratelimited("sema-svc: counter add: channel 0x%x not found\n",
+				    h_object);
+	return ret;
+}
+
 static int svc_thread(void *arg)
 {
 	struct nv_vgpu_device *vgpu;
@@ -3865,7 +4193,8 @@ static int svc_thread(void *arg)
 		svc_handle_chan_dump(vgpu);
 
 		if (!svc.direct_base ||
-		    time_after(jiffies, svc.last_build + 10 * HZ))
+		    time_after(jiffies,
+			       (unsigned long)(svc.last_build + 10 * HZ)))
 			svc_rebuild(vgpu);
 
 		/*