radio CUDA: reuse one pinned pixel stage for every image

docs/runtime/cuda-copy-sources.md

@@ -20,6 +20,10 @@ The legacy allocator argument and Stage.gpa field remain source compatible. They
 Quarantinable Node and CPU data allocations always use the core's process-lifetime page allocator. Pinned data is
 owned by the driver until its completion and disposal callbacks return. Caller arenas may end after submission.
 
+The RADIO tower owns one pixel Stage, reserved with its scratch for the largest image. Each image waits on the
+previous copy's event before refilling it, so tower uploads neither allocate nor free pinned memory and never compete
+for temporary-upload spares.
+
 The graph wrappers track up to 64 simultaneous captures by driver, context, stream and capture mode across threads.
 Temporary synchronous/default-stream and asynchronous uploaders refuse while a wrapper capture applies. A global
 capture guards all contexts of that driver; other modes guard its saved context. A failed begin removes its scope.

zig/src/families/radio/radio.zig

@@ -75,4 +75,5 @@ test {
 
 test {
     _ = @import("video_test.zig");
+    _ = @import("upload_test.zig");
 }

zig/src/families/radio/tower.zig

@@ -17,6 +17,7 @@ pub const Tower = struct {
     g: vit.Gemms,
     norms: cuda.norms.Norms,
     scratch: cuda.DeviceBuffer,
+    stage: cuda.copy_source.Stage,
     b: Buffers,
     bytes: u64, // on the device: weights, scratch and the Linears' workspace
 
@@ -45,6 +46,9 @@ pub const Tower = struct {
         for (sizes) |sz| total += std.mem.alignForward(usize, sz, 256);
         var scratch = try cuda.DeviceBuffer.alloc(d, total);
         errdefer scratch.free();
+        var stage: cuda.copy_source.Stage = .{};
+        errdefer stage.deinit();
+        _ = try stage.take(d, n * c.patchWidth() * @sizeOf(f32));
         var ptrs: [sizes.len]u64 = undefined;
         var off: usize = 0;
         for (sizes, &ptrs) |sz, *p| {
@@ -60,12 +64,14 @@ pub const Tower = struct {
             .g = g,
             .norms = norms,
             .scratch = scratch,
+            .stage = stage,
             .b = .{ .px = ptrs[0], .h = ptrs[1], .x = ptrs[2], .wide = ptrs[3], .a = ptrs[4], .acc = ptrs[5] },
             .bytes = w.bytes + total + g.workspace.len,
         };
     }
 
     pub fn deinit(t: *Tower) void {
+        t.stage.deinit();
         t.scratch.free();
         t.norms.unload();
         t.g.deinit();
@@ -111,7 +117,9 @@ pub const Tower = struct {
         if (img.pixels.len != values or values > room) return error.ImageShape;
         const hd = c.hidden;
         // the float patches land in the f32 products' room (the patch Linear has no bias), cast to bf16 on the GPU
-        try t.scratch.uploadAsync(b.acc - t.scratch.ptr, std.mem.sliceAsBytes(img.pixels), s.handle);
+        const pixels = std.mem.sliceAsBytes(img.pixels);
+        @memcpy(try t.stage.take(s.d, pixels.len), pixels);
+        try t.stage.copyTo(t.scratch, b.acc - t.scratch.ptr, s);
         try k.cast(s, b.acc, b.px, n * pw);
         // the cls and register rows, then each patch's projection plus its position row
         const rows = n + c.prefix;

zig/src/families/radio/upload_test.zig

@@ -0,0 +1,116 @@
+//! The tower's upload through a fake driver, stopping before device arithmetic.
+const std = @import("std");
+const t = std.testing;
+const cuda = @import("cuda");
+const abi = cuda.abi;
+const source = cuda.copy_source;
+const Tower = @import("tower.zig").Tower;
+
+const Fake = struct {
+    var allocs: usize = 0;
+    var frees: usize = 0;
+    var waits: usize = 0;
+    var flight: ?[*]const u8 = null;
+    var expected: []const u8 = &.{};
+    var held: []const u8 = &.{};
+    var blocks: std.ArrayList([]align(16) u8) = .empty;
+    fn alloc(out: *?*anyopaque, len: usize, _: c_uint) callconv(.c) abi.Result {
+        const b = t.allocator.alignedAlloc(u8, .@"16", len) catch unreachable;
+        blocks.append(t.allocator, b) catch unreachable;
+        allocs += 1;
+        out.* = b.ptr;
+        return 0;
+    }
+    fn free(ptr: ?*anyopaque) callconv(.c) abi.Result {
+        std.debug.assert(flight == null);
+        for (blocks.items, 0..) |b, i| if (@intFromPtr(b.ptr) == @intFromPtr(ptr)) {
+            t.allocator.free(b);
+            _ = blocks.swapRemove(i);
+            frees += 1;
+            return 0;
+        };
+        unreachable;
+    }
+    fn copy(_: u64, ptr: ?*const anyopaque, len: usize, _: abi.Stream) callconv(.c) abi.Result {
+        std.debug.assert(flight == null);
+        const bytes: [*]const u8 = @ptrCast(ptr.?);
+        std.debug.assert(std.mem.eql(u8, expected, bytes[0..len]));
+        flight = bytes;
+        held = expected;
+        return 0;
+    }
+    fn create(out: *abi.Event, _: c_uint) callconv(.c) abi.Result {
+        out.* = @ptrFromInt(1);
+        return 0;
+    }
+    fn record(_: abi.Event, _: abi.Stream) callconv(.c) abi.Result {
+        return 0;
+    }
+    fn query(_: abi.Event) callconv(.c) abi.Result {
+        return if (flight == null) 0 else abi.error_not_ready;
+    }
+    fn wait(_: abi.Event) callconv(.c) abi.Result {
+        if (flight) |ptr| std.debug.assert(std.mem.eql(u8, held, ptr[0..held.len]));
+        flight = null;
+        waits += 1;
+        return 0;
+    }
+    fn destroy(_: abi.Event) callconv(.c) abi.Result {
+        return 0;
+    }
+    fn sync(_: abi.Stream) callconv(.c) abi.Result {
+        return wait(null);
+    }
+    fn launch(_: abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result {
+        return abi.error_not_ready;
+    }
+};
+
+test "RADIO upload reuses its pinned pixels without freeing between images" {
+    Fake.allocs = 0;
+    Fake.frees = 0;
+    Fake.waits = 0;
+    Fake.flight = null;
+    defer Fake.blocks.deinit(t.allocator);
+    var d: cuda.Driver = undefined;
+    d.api.cuMemHostAlloc = Fake.alloc;
+    d.api.cuMemFreeHost = Fake.free;
+    d.api.cuMemcpyHtoDAsync_v2 = Fake.copy;
+    d.api.cuEventCreate = Fake.create;
+    d.api.cuEventRecord = Fake.record;
+    d.api.cuEventQuery = Fake.query;
+    d.api.cuEventSynchronize = Fake.wait;
+    d.api.cuEventDestroy_v2 = Fake.destroy;
+    d.api.cuStreamSynchronize = Fake.sync;
+    d.api.cuLaunchKernel = Fake.launch;
+    source.setCurrent(null);
+    source.setSparesForTest(.{});
+    defer source.setSparesForTest(.{ .count = 0, .bytes = 0 });
+    defer source.dropSpares(&d, null);
+    const stream: cuda.Stream = .{ .d = &d, .handle = null };
+    var tower: Tower = undefined;
+    tower.c = @import("video_test.zig").cfg;
+    tower.b = .{ .px = 0, .h = 0, .x = 0, .wide = 0, .a = 0, .acc = 256 };
+    tower.stage = .{};
+    defer tower.stage.deinit();
+    const len = tower.c.max_patches * 3 * tower.c.patch * tower.c.patch;
+    const images = try t.allocator.alloc(f32, 2 * len);
+    defer t.allocator.free(images);
+    @memset(images[0..len], 0.5);
+    @memset(images[len..], -0.5);
+    defer source.synchronizeStream(&d, null) catch unreachable;
+    tower.scratch = .{ .d = &d, .ptr = 256, .len = len * @sizeOf(f32) };
+    tower.k.cast_fn = .{ .d = &d, .handle = null };
+    _ = try tower.stage.take(&d, len * @sizeOf(f32));
+    for (0..4) |i| {
+        const pixels = images[i % 2 * len ..][0..len];
+        Fake.expected = std.mem.sliceAsBytes(pixels);
+        try t.expectError(error.NotReady, tower.features(stream, .{ .pixels = pixels, .gh = 128, .gw = 104, .tokens = 3328, .at = 0 }));
+        try t.expectEqual(@as(usize, 1), Fake.allocs);
+        try t.expectEqual(@as(usize, 0), Fake.frees);
+    }
+    try t.expectEqual(@as(usize, 3), Fake.waits);
+    try tower.stage.close();
+    try t.expectEqual(@as(usize, 4), Fake.waits);
+    try t.expectEqual(@as(usize, 1), Fake.frees);
+}