diff --git a/docs/runtime/cuda-copy-sources.md b/docs/runtime/cuda-copy-sources.md index a7629d536c..5499f4b844 100644 --- a/docs/runtime/cuda-copy-sources.md +++ b/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. diff --git a/zig/src/families/radio/radio.zig b/zig/src/families/radio/radio.zig index 509017d18c..f6d5072218 100644 --- a/zig/src/families/radio/radio.zig +++ b/zig/src/families/radio/radio.zig @@ -75,4 +75,5 @@ test { test { _ = @import("video_test.zig"); + _ = @import("upload_test.zig"); } diff --git a/zig/src/families/radio/tower.zig b/zig/src/families/radio/tower.zig index cabd18cc47..dfbba7906d 100644 --- a/zig/src/families/radio/tower.zig +++ b/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; diff --git a/zig/src/families/radio/upload_test.zig b/zig/src/families/radio/upload_test.zig new file mode 100644 index 0000000000..e1861f1e9d --- /dev/null +++ b/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); +}