NVMTLDriver (2.5.3): cursor lag and slow drags during screen capture - IOSurface re-imports and tiny-surface re-uploads (measurements + patch)

diff --git a/plugin/NVMTLDevice.m b/plugin/NVMTLDevice.m
index 90205a1..8783455 100644
--- a/plugin/NVMTLDevice.m
+++ b/plugin/NVMTLDevice.m
@@ -918,16 +918,126 @@ static void nvmtl_vram_tag_set(IOSurfaceRef surf, NSUInteger plane, uint32_t see
     CFDataRef d = CFDataCreate(NULL, (const UInt8 *)&t, sizeof t); CFStringRef k = nvmtl_vram_tag_key(plane);
     IOSurfaceSetValue(surf, k, d); CFRelease(k); CFRelease(d);
 }
+// Shared VRAM imports (surfcache). WindowServer wraps the same few IOSurfaces again and again: a 10 s dtrace during a
+// Discord screen capture saw 1518 -newTextureWithDescriptor:iosurface:plane: calls for about a dozen surfaces, each
+// one an RM import (nvk_AllocateMemory -> nvkmd_nvrm_import_dma_buf) because the per-device surface cache keeps only
+// private images (a VRAM wrapper is dropped with its last texture, after one use). Here the import itself is kept and
+// shared: a texture over a surface's VRAM takes a reference to the surface's import, and the import outlives its
+// textures until the surface is gone (only this cache still retains it), the cache is full, or it is evicted as the
+// least recently used. Each texture still makes its own image over the buffer and does its own seed/tag sync, as
+// before. The VRAM stays the surface's own across page-off because NVAccel parks and hands it back
+// (debug.nvaccel_vm_sticky=1); with sticky off the cache stays off. NVMTL_NO_SURFVRAM_CACHE=1 turns it off.
+typedef struct nvmtl_surfvram {
+    IOSurfaceRef surf; uint32_t id, plane; size_t need;
+    nvk_buffer vb; uint32_t refs; uint64_t used;
+} nvmtl_surfvram;
+#define NVMTL_SURFVRAM_N 256   // a desktop under capture cycles ~70 surfaces per frame (dtrace, 2026-10-22)
+static nvmtl_surfvram g_svc[NVMTL_SURFVRAM_N];
+static pthread_mutex_t g_svc_lock = PTHREAD_MUTEX_INITIALIZER;
+static uint64_t g_svc_clock, g_svc_hits, g_svc_miss, g_svc_evict, g_svc_once;
+#define NVMTL_SURFVRAM_SEEN 512
+static uintptr_t g_svc_seen[NVMTL_SURFVRAM_SEEN]; static unsigned g_svc_seen_at;
+static int nvmtl_surfvram_on(void)
+{
+    static int on = -1;
+    if (on < 0) {
+        const char *e = getenv("NVMTL_NO_SURFVRAM_CACHE");
+        int sticky = 0; size_t n = sizeof sticky;
+        if (sysctlbyname("debug.nvaccel_vm_sticky", &sticky, &n, NULL, 0)) sticky = 0;
+        on = !(e && e[0] == '1') && sticky == 1;
+        nvlog("surfcache: shared IOSurface VRAM imports %s (vm_sticky %d, NVMTL_NO_SURFVRAM_CACHE=%s)", on ? "ON" : "off", sticky, e ?: "unset");
+    }
+    return on;
+}
+// an unused entry (no texture holds it) whose surface only this cache still retains, or any unused one when forced
+static void nvmtl_surfvram_drop_locked(nvmtl_surfvram *e, nvk_buffer *dead, IOSurfaceRef *deadSurf)
+{
+    *dead = e->vb; *deadSurf = e->surf; g_svc_evict++;
+    memset(e, 0, sizeof *e);
+}
+static void nvmtl_surfvram_sweep(int needSlot)
+{
+    nvk_buffer dead[NVMTL_SURFVRAM_N]; IOSurfaceRef deadSurf[NVMTL_SURFVRAM_N]; unsigned nd = 0;
+    pthread_mutex_lock(&g_svc_lock);
+    int freeSlot = 0;
+    for (unsigned i = 0; i < NVMTL_SURFVRAM_N; i++) {
+        nvmtl_surfvram *e = &g_svc[i];
+        if (!e->surf) { freeSlot = 1; continue; }
+        if (!e->refs && CFGetRetainCount(e->surf) <= 1) { nvmtl_surfvram_drop_locked(e, &dead[nd], &deadSurf[nd]); nd++; freeSlot = 1; }
+    }
+    if (needSlot && !freeSlot) {
+        nvmtl_surfvram *lru = NULL;
+        for (unsigned i = 0; i < NVMTL_SURFVRAM_N; i++)
+            if (g_svc[i].surf && !g_svc[i].refs && (!lru || g_svc[i].used < lru->used)) lru = &g_svc[i];
+        if (lru) { nvmtl_surfvram_drop_locked(lru, &dead[nd], &deadSurf[nd]); nd++; }
+    }
+    pthread_mutex_unlock(&g_svc_lock);
+    for (unsigned k = 0; k < nd; k++) { nvmtl_vk_buffer_destroy(&dead[k]); CFRelease(deadSurf[k]); }
+}
+// 0 and *out/*share filled: the texture uses the surface's shared import; -1: import refused (caller goes private)
+static int nvmtl_surfvram_get(IOSurfaceRef surf, uint32_t plane, size_t need, nvk_buffer *out, void **share)
+{
+    *share = NULL;
+    if (!nvmtl_surfvram_on()) return nvmtl_vk_surface_vram(IOSurfaceGetID(surf), plane, need, out);
+    const uint32_t id = IOSurfaceGetID(surf);
+    pthread_mutex_lock(&g_svc_lock);
+    for (unsigned i = 0; i < NVMTL_SURFVRAM_N; i++) {
+        nvmtl_surfvram *e = &g_svc[i];
+        if (e->surf == surf && e->id == id && e->plane == plane && e->need == need) {
+            e->refs++; e->used = ++g_svc_clock; g_svc_hits++;
+            *out = e->vb; out->imported = 1; *share = e;
+            const int sweep = (g_svc_hits & 255) == 0;
+            pthread_mutex_unlock(&g_svc_lock);
+            if (sweep) nvmtl_surfvram_sweep(0);   // let go of surfaces the app has freed even when nothing misses
+            return 0;
+        }
+    }
+    g_svc_miss++;
+    const uint64_t miss = g_svc_miss, hits = g_svc_hits, ev = g_svc_evict;
+    // Admission: a surface is cached from its second import while it is still in the recent ring. Video frames and
+    // other surfaces an app makes once (Firefox playing a video made ~1400 2 MB surfaces in 6 s) pass straight through
+    // and no longer evict the surfaces CoreAnimation re-wraps every frame.
+    const uintptr_t tag = (uintptr_t)surf ^ ((uintptr_t)id << 32) ^ (uintptr_t)need;
+    int admit = 0;
+    for (unsigned i = 0; i < NVMTL_SURFVRAM_SEEN && !admit; i++) admit = g_svc_seen[i] == tag;
+    if (!admit) { g_svc_seen[g_svc_seen_at++ % NVMTL_SURFVRAM_SEEN] = tag; g_svc_once++; }
+    pthread_mutex_unlock(&g_svc_lock);
+    if (!admit) return nvmtl_vk_surface_vram(id, plane, need, out);   // first sight: this texture owns its import
+    if (miss <= 4 || miss % 500 == 0)
+        nvlog("surfcache: %llu hits, %llu imports, %llu evicted", (unsigned long long)hits, (unsigned long long)miss, (unsigned long long)ev);
+    nvk_buffer vb;
+    if (nvmtl_vk_surface_vram(id, plane, need, &vb)) return -1;
+    nvmtl_surfvram_sweep(1);
+    pthread_mutex_lock(&g_svc_lock);
+    nvmtl_surfvram *slot = NULL;
+    for (unsigned i = 0; i < NVMTL_SURFVRAM_N && !slot; i++) if (!g_svc[i].surf) slot = &g_svc[i];
+    if (!slot) { pthread_mutex_unlock(&g_svc_lock); *out = vb; return 0; }   // every entry in use: this texture owns its import
+    CFRetain(surf);
+    slot->surf = surf; slot->id = id; slot->plane = plane; slot->need = need; slot->vb = vb;
+    slot->refs = 1; slot->used = ++g_svc_clock;
+    *out = vb; *share = slot;
+    pthread_mutex_unlock(&g_svc_lock);
+    return 0;
+}
+static void nvmtl_surfvram_put(void *share)
+{
+    if (!share) return;
+    pthread_mutex_lock(&g_svc_lock);
+    nvmtl_surfvram *e = (nvmtl_surfvram *)share;
+    if (e->refs) e->refs--;
+    pthread_mutex_unlock(&g_svc_lock);
+}
+
 static int nvmtl_surface_image_create(NVMTLTexture *t, IOSurfaceRef surf, NSUInteger plane, uint32_t w, uint32_t h,
                                       uint32_t vkf, uint32_t bpp, int a8, int stex)
 {
     if (surf && plane == 0 && IOSurfaceGetPlaneCount(surf) <= 1 && nvmtl_vk_surface_share_on()) {
         size_t row = IOSurfaceGetBytesPerRow(surf), sh = IOSurfaceGetHeight(surf);
         if ((size_t)w * (bpp ? bpp : 4) <= row && h <= sh && row % 32 == 0) {
-            nvk_buffer vb;
-            if (!nvmtl_vk_surface_vram(IOSurfaceGetID(surf), 0, row * h, &vb)) {
+            nvk_buffer vb; void *share = NULL;
+            if (!nvmtl_surfvram_get(surf, 0, row * h, &vb, &share)) {
                 if (!nvmtl_vk_image_create_buffer_alias(w, h, vkf, bpp, a8, stex, 1, &vb, 0, row, &t->_i)) {
-                    t->_vramBuf = vb; t->_vramOn = YES; g_surf_vram++;
+                    t->_vramBuf = vb; t->_vramOn = YES; t->_vramShare = share; g_surf_vram++;
                     uint32_t seed = IOSurfaceGetSeed(surf);
                     int current = nvmtl_vram_tag_current(surf, 0, seed);
                     if (current) { t->_surfSeed = seed; t->_surfSynced = YES; }
@@ -936,7 +1046,8 @@ static int nvmtl_surface_image_create(NVMTLTexture *t, IOSurfaceRef surf, NSUInt
                               w, h, row, current ? ", already current (no upload)" : "", g_surf_vram, g_surf_private);
                     return 0;
                 }
-                nvmtl_vk_buffer_destroy(&vb);
+                if (share) nvmtl_surfvram_put(share);   // the shared import stays for the next wrapper
+                else nvmtl_vk_buffer_destroy(&vb);
             }
         }
     }
@@ -2129,6 +2240,17 @@ - (NSUInteger)gpuCoreCount { uint32_t sm = 0; nvmtl_vk_limits(NULL, NULL, NULL,
 
 #import "NVMTLDevice_contract.m"
 
+// surface-in shadow (see -[NVMTLTexture nvmtlSurfaceIn])
+#define NVMTL_SURFIN_SHADOW_MAX (64u * 1024u)
+// The shadow belongs to the IOSurface object (an associated object), not to a texture: WindowServer wraps the same
+// surface in a new texture on most frames. gen is the surface's VRAM tag generation when the shadow was taken (0 for
+// private images).
+@interface NVMTLSurfShadow : NSObject { @public NSData *_bytes; uint32_t _gen; __weak id _owner; } @end   // _owner: the texture whose private image holds the bytes
+@implementation NVMTLSurfShadow @end
+static char nvmtl_surfin_shadow_key;
+static NVMTLSurfShadow *nvmtl_surfin_shadow(IOSurfaceRef s) { return objc_getAssociatedObject((__bridge id)s, &nvmtl_surfin_shadow_key); }
+static void nvmtl_surfin_shadow_set(IOSurfaceRef s, NVMTLSurfShadow *v) {
+    objc_setAssociatedObject((__bridge id)s, &nvmtl_surfin_shadow_key, v, OBJC_ASSOCIATION_RETAIN); }
 #import "NVMTLObjects.m"
 #import "NVMTLForward.m"
 #import "NVMTLLibraryLoad.m"
diff --git a/plugin/NVMTLObjects.h b/plugin/NVMTLObjects.h
index 0524ba9..545b14b 100644
--- a/plugin/NVMTLObjects.h
+++ b/plugin/NVMTLObjects.h
@@ -45,7 +45,7 @@ static inline void *nvmtl_buf_cpu(NVMTLBuffer *b) { nvmtl_buf_settle(b); return
   uint32_t _viewType, _viewSlices, _relativeLevel, _relativeSlice;
   uint32_t _baseSlice, _baseLevel;
   uint32_t _surfSeed, _surfGen; BOOL _surfSynced, _surfOwned;
-  nvk_buffer _vramBuf; BOOL _vramOn;
+  nvk_buffer _vramBuf; BOOL _vramOn; void *_vramShare;
   BOOL _stex;
   uint64_t _resid;
   NSUInteger _usage;
diff --git a/plugin/NVMTLObjects.m b/plugin/NVMTLObjects.m
index 01f018b..2fb9c88 100644
--- a/plugin/NVMTLObjects.m
+++ b/plugin/NVMTLObjects.m
@@ -1169,7 +1169,10 @@ - (void)dealloc {
     if (!_parent && _bparkable && _backBuf && _i.img) nvmtl_bpark_put(_backBuf, &_i, _pkey);
     if (!_parent) nvmtl_vk_image_destroy(&_i);
     if (_stTwin.img) nvmtl_vk_image_destroy(&_stTwin);
-    if (!_parent && _vramOn) { nvmtl_vk_buffer_destroy(&_vramBuf); _vramOn = NO; }
+    if (!_parent && _vramOn) {
+        if (_vramShare) { nvmtl_surfvram_put(_vramShare); _vramShare = NULL; memset(&_vramBuf, 0, sizeof _vramBuf); }
+        else nvmtl_vk_buffer_destroy(&_vramBuf);
+        _vramOn = NO; }
     if (_surfOwned && _surf) CFRelease(_surf);
     if (_subHeap) { NVMTLHeap *h = _subHeap; NSUInteger c = _subCharge; _subHeap = nil; _subCharge = 0;
                     if (_subPlaced && c) { _subPlaced = NO; [h nvmtlGiveRange:_heapOffset length:c]; }
@@ -1240,8 +1243,26 @@ - (void)nvmtlSurfaceIn {
     if (IOSurfaceLock(r->_surf, kIOSurfaceLockReadOnly, &seed) != kIOReturnSuccess) { nvlog("G13 surface-in: IOSurfaceLock FAILED"); return; }
     void *base = IOSurfaceGetBaseAddressOfPlane(r->_surf, (uint32_t)r->_plane);
     size_t row = IOSurfaceGetBytesPerRowOfPlane(r->_surf, (uint32_t)r->_plane);
+    // Small surfaces: during a screen capture WindowServer re-uploaded one 8x8 surface ~127 times a second at ~1.4 ms
+    // each (host import, map, submit, wait, unmap, free), about 17% of its main thread. The bytes last uploaded are kept;
+    // a new seed over the same bytes is synced without an upload, and a real change is copied through the staging path
+    // (no page import). The shadow is dropped whenever the GPU writes the surface (-nvmtlSurfaceOutDirtied:).
+    const size_t span = (base && row) ? (size_t)r->_i.h * row : 0;
+    const bool small = span && span <= NVMTL_SURFIN_SHADOW_MAX;
+    // shared VRAM can be written by another process's GPU: its tag generation (bumped on every GPU write, in any
+    // process) must be the one recorded when the shadow was taken
+    NVMTLSurfShadow *shadow = small ? nvmtl_surfin_shadow(r->_surf) : nil;
+    // shared VRAM holds the bytes for every texture over the surface; a private image only for the texture uploaded into
+    if (shadow && shadow->_bytes.length == span && (r->_vramOn ? shadow->_gen == nvmtl_vram_gen(r->_surf, r->_plane) : shadow->_owner == r)
+        && !memcmp(shadow->_bytes.bytes, base, span)) {
+        IOSurfaceUnlock(r->_surf, kIOSurfaceLockReadOnly, NULL);
+        r->_surfSeed = seed; r->_surfSynced = YES; r->_surfGen = vgen; g_surfin_skip++;
+        if (r->_vramOn) nvmtl_vram_tag_set(r->_surf, r->_plane, seed);
+        return;
+    }
+    NSData *taken = small ? [NSData dataWithBytes:base length:span] : nil;   // read under the lock; kept only if the upload succeeds
     int rc = -1;
-    if (base && row && nvmtl_surfin_import()) {
+    if (base && row && !small && nvmtl_surfin_import()) {
         const size_t off = (size_t)((uint8_t *)base - (uint8_t *)IOSurfaceGetBaseAddress(r->_surf));
         rc = nvmtl_vk_image_write_host(r->_q, &r->_i, base, row, IOSurfaceGetAllocSize(r->_surf) - off);
     }
@@ -1251,6 +1272,8 @@ - (void)nvmtlSurfaceIn {
     if (rc) { nvlog("G13 surface-in: upload FAILED (%ux%u) — retried at the next use", r->_i.w, r->_i.h); return; }
     r->_surfSeed = seed; r->_surfSynced = YES; r->_surfGen = vgen;
     g_surfin_copy++; if (r->_vramOn) nvmtl_vram_tag_set(r->_surf, r->_plane, seed);
+    if (taken) { NVMTLSurfShadow *v = [NVMTLSurfShadow new]; v->_bytes = taken; v->_gen = r->_vramOn ? nvmtl_vram_gen(r->_surf, r->_plane) : 0; v->_owner = r;
+                 nvmtl_surfin_shadow_set(r->_surf, v); }
     static unsigned n; if (n++ < 6 || n % 2000 == 0) nvlog("G13 surface-in #%u: %ux%u surface %p seed %u uploaded", n, r->_i.w, r->_i.h, (void *)r->_surf, seed);
     }
 }
@@ -1288,6 +1311,7 @@ - (BOOL)nvmtlSurfaceOutDirtied:(BOOL)dirtied {
     NVMTLTexture *r = _parent ? _parent : self;
     if (!r->_surf) return YES;
     @synchronized (nvmtl_surface_sync(r)) {
+    if (nvmtl_surfin_shadow(r->_surf)) nvmtl_surfin_shadow_set(r->_surf, nil);   // the GPU wrote it: the next CPU seed uploads
     if (![r nvmtlSurfaceCopyable]) { nvtrace("SURFOUT surf %u %ux%u NOT copyable", (unsigned)IOSurfaceGetID(r->_surf), r->_i.w, r->_i.h); return NO; }
     nvtrace("SURFOUT surf %u %ux%u vramOn %d t=%llu", (unsigned)IOSurfaceGetID(r->_surf), r->_i.w, r->_i.h, (int)r->_vramOn,
             (unsigned long long)(clock_gettime_nsec_np(CLOCK_UPTIME_RAW) / 1000));