tiny.accy.executable.binding

1const std = @import("std");2const gpu = @import("gpu");3const accy_root = @import("../root.zig");4const artifact_product = @import("../artifact/root.zig");5const plan_mod = @import("plan.zig");67const BackendArtifactPlan = artifact_product.BackendArtifactPlan;8const BackendMemoryPlan = artifact_product.plan.BackendMemoryPlan;9const PlannedSlot = artifact_product.PlannedSlot;1011pub const SlotBinding = struct {12    slot_id: usize,13    binding: gpu.BufferBinding,14};1516pub const ElementCountBufferBinding = struct {17    kernel_id: usize,18    binding: gpu.BufferBinding,19};2021const BackendPlanBuffer = struct {22    allocation_id: usize,23    handle: gpu.BufferHandle,24};2526const OutputCopy = struct {27    output_index: usize,28    slot_id: usize,29};3031const PreparedLaunchBindingsState = struct {32    allocator: std.mem.Allocator,33    handle: gpu.BackendHandle,34    live_buffers: std.ArrayList(gpu.BufferHandle) = .empty,35    slot_bindings: std.ArrayList(SlotBinding) = .empty,36    plan_buffers: std.ArrayList(BackendPlanBuffer) = .empty,37    count_bindings: std.ArrayList(ElementCountBufferBinding) = .empty,38    output_copies: std.ArrayList(OutputCopy) = .empty,39    borrowed_count: usize = 0,4041    fn create(42        allocator: std.mem.Allocator,43        handle: gpu.BackendHandle,44    ) !*PreparedLaunchBindingsState {45        const state = try allocator.create(PreparedLaunchBindingsState);46        state.* = .{ .allocator = allocator, .handle = handle };47        return state;48    }4950    fn destroy(self: *PreparedLaunchBindingsState) void {51        const allocator = self.allocator;52        std.debug.assert(self.borrowed_count <= self.live_buffers.items.len);53        for (self.live_buffers.items) |buffer| self.handle.destroyObject(buffer.id);54        self.output_copies.deinit(allocator);55        self.count_bindings.deinit(allocator);56        self.plan_buffers.deinit(allocator);57        self.slot_bindings.deinit(allocator);58        self.live_buffers.deinit(allocator);59        self.* = undefined;60        allocator.destroy(self);61    }62};6364pub const PreparedLaunchBindings = opaque {65    fn stateConst(self: *const PreparedLaunchBindings) *const PreparedLaunchBindingsState {66        return @ptrCast(@alignCast(self));67    }6869    fn stateMut(self: *PreparedLaunchBindings) *PreparedLaunchBindingsState {70        return @ptrCast(@alignCast(self));71    }7273    pub fn deinit(self: *PreparedLaunchBindings) void {74        self.stateMut().destroy();75    }76};7778pub fn liveBufferCount(bindings: *const PreparedLaunchBindings) usize {79    const state = bindings.stateConst();80    return state.live_buffers.items.len;81}8283pub fn slotBindings(bindings: *const PreparedLaunchBindings) []const SlotBinding {84    const state = bindings.stateConst();85    return state.slot_bindings.items;86}8788pub fn elementCountBindings(bindings: *const PreparedLaunchBindings) []const ElementCountBufferBinding {89    const state = bindings.stateConst();90    return state.count_bindings.items;91}9293pub fn elementCountBindingForKernel(94    bindings: *const PreparedLaunchBindings,95    kernel_id: usize,96) ?gpu.BufferBinding {97    const state = bindings.stateConst();98    for (state.count_bindings.items) |binding| {99        if (binding.kernel_id == kernel_id) return binding.binding;100    }101    return null;102}103104/// Returns how many buffers `prepareBorrowed` imported from caller memory, so a105/// caller or a test can see how many buffers of one run use that memory in106/// place.107pub fn borrowedBufferCount(bindings: *const PreparedLaunchBindings) usize {108    const state = bindings.stateConst();109    std.debug.assert(state.borrowed_count <= state.live_buffers.items.len);110    return state.borrowed_count;111}112113pub fn prepare(114    allocator: std.mem.Allocator,115    handle: gpu.BackendHandle,116    artifact_plan: *const BackendArtifactPlan,117    inputs: []const []const u8,118) !*PreparedLaunchBindings {119    if (artifact_plan.input_slot_ids.len != inputs.len) return error.InvalidArtifact;120    const bindings = try PreparedLaunchBindingsState.create(allocator, handle);121    errdefer bindings.destroy();122    const min_buffer_alignment = try backendMinBufferAlignment(handle);123124    for (artifact_plan.input_slot_ids, inputs) |slot_id, input| {125        const slot = artifact_plan.slotById(slot_id) orelse return error.InvalidArtifact;126        if (slot.byte_size == null or input.len != slot.byte_size.?) return error.InvalidArtifact;127        try bindCopiedInput(bindings, slot.*, input, min_buffer_alignment);128    }129130    var memory_plan = try artifact_product.createBackendMemoryPlan(allocator, artifact_plan);131    defer memory_plan.deinit();132    try bindPlannedSlots(bindings, artifact_plan, &memory_plan, min_buffer_alignment);133    try bindElementCounts(bindings, artifact_plan, min_buffer_alignment);134    return @ptrCast(bindings);135}136137/// Binds the caller's `inputs` and `outputs` to planned buffers of one run, and138/// imports each slice in place when the plans show that no kernel will write139/// through it by surprise, so the caller runs a compiled program directly on140/// its own input and output slices, avoiding a copy when the plans allow it. An141/// input is imported read-only when no kernel output and no loop carry writes142/// its slot, and the call copies the input into a new buffer otherwise. An143/// output is copied back later by `completeBorrowed` when its slot appears144/// earlier in the output list, has no planned allocation, is an input, is a145/// constant, or is a loop carry. Every other output is imported as the buffer146/// of its planned allocation, and two outputs that share one allocation give147/// `error.InvalidArtifact`. Input and output counts that differ from the plan148/// give `error.InvalidArtifact`, and an output that overlaps an input or an149/// earlier output gives `error.InvalidBuffer`. Each slice must be exactly its150/// slot's size, and each imported slice must meet the slot's alignment, else151/// the call returns `error.InvalidBuffer`. Imported slices are used in place by152/// the backend, so the caller keeps every slice alive and unmoved until the153/// returned bindings are released with `deinit`.154pub fn prepareBorrowed(155    allocator: std.mem.Allocator,156    handle: gpu.BackendHandle,157    artifact_plan: *const BackendArtifactPlan,158    launch_plan: plan_mod.LaunchGraphPlan,159    inputs: []const []const u8,160    outputs: []const []u8,161) !*PreparedLaunchBindings {162    if (artifact_plan.input_slot_ids.len != inputs.len) return error.InvalidArtifact;163    if (artifact_plan.output_slot_ids.len != outputs.len) return error.InvalidArtifact;164    try expectDisjointOutputs(inputs, outputs);165    const bindings = try PreparedLaunchBindingsState.create(allocator, handle);166    errdefer bindings.destroy();167    const min_buffer_alignment = try backendMinBufferAlignment(handle);168169    for (artifact_plan.input_slot_ids, inputs) |slot_id, input| {170        const slot = artifact_plan.slotById(slot_id) orelse return error.InvalidArtifact;171        if (try slotByteSize(slot.*) != input.len) return error.InvalidBuffer;172        if (slotIsWritten(artifact_plan, launch_plan, slot_id)) {173            try bindCopiedInput(bindings, slot.*, input, min_buffer_alignment);174            continue;175        }176        const buffer = try importSlotBuffer(177            bindings,178            slot.*,179            @constCast(input),180            min_buffer_alignment,181        );182        try appendSlotBindingWithAccess(183            allocator,184            &bindings.slot_bindings,185            slot_id,186            buffer,187            .read_only,188        );189    }190191    var memory_plan = try artifact_product.createBackendMemoryPlan(allocator, artifact_plan);192    defer memory_plan.deinit();193    for (outputs, 0..) |output, index| {194        try bindBorrowedOutput(195            bindings,196            artifact_plan,197            launch_plan,198            &memory_plan,199            index,200            output,201            min_buffer_alignment,202        );203    }204    try bindPlannedSlots(bindings, artifact_plan, &memory_plan, min_buffer_alignment);205    try bindElementCounts(bindings, artifact_plan, min_buffer_alignment);206    return @ptrCast(bindings);207}208209/// Reads each output that `prepareBorrowed` set aside for copying from its210/// device buffer into the caller's slice, so the caller fills outputs that211/// could not be written in place. The caller calls this function only after the212/// launch has completed, as the session does. An output index beyond `outputs`213/// gives `error.InvalidBuffer`.214pub fn completeBorrowed(215    bindings: *const PreparedLaunchBindings,216    outputs: []const []u8,217) gpu.BackendError!void {218    const state = bindings.stateConst();219    for (state.output_copies.items) |copy| {220        if (copy.output_index >= outputs.len) return error.InvalidBuffer;221        const binding = try bindingForSlot(state.slot_bindings.items, copy.slot_id, .read_only);222        try state.handle.readBuffer(.{223            .handle = binding.handle,224            .bytes = outputs[copy.output_index],225        });226    }227}228229fn bindCopiedInput(230    state: *PreparedLaunchBindingsState,231    slot: PlannedSlot,232    input: []const u8,233    min_buffer_alignment: u32,234) !void {235    const buffer = try allocateBackendSlotBuffer(236        state.allocator,237        state.handle,238        &state.live_buffers,239        slot,240        min_buffer_alignment,241    );242    try state.handle.writeBuffer(.{ .handle = buffer, .bytes = input });243    try appendSlotBinding(state.allocator, &state.slot_bindings, slot.slot_id, buffer);244}245246fn bindBorrowedOutput(247    state: *PreparedLaunchBindingsState,248    artifact_plan: *const BackendArtifactPlan,249    launch_plan: plan_mod.LaunchGraphPlan,250    memory_plan: *const BackendMemoryPlan,251    index: usize,252    output: []u8,253    min_buffer_alignment: u32,254) !void {255    const slot_id = artifact_plan.output_slot_ids[index];256    const slot = artifact_plan.slotById(slot_id) orelse return error.InvalidArtifact;257    if (try slotByteSize(slot.*) != output.len) return error.InvalidBuffer;258    const earlier_outputs = artifact_plan.output_slot_ids[0..index];259    const repeated = std.mem.indexOfScalar(usize, earlier_outputs, slot_id) != null;260    const assignment = memory_plan.assignmentForSlot(slot_id);261    if (repeated or assignment == null or slot.role.input or slot.role.constant or262        slotIsCarried(launch_plan, slot_id))263    {264        try state.output_copies.append(state.allocator, .{265            .output_index = index,266            .slot_id = slot_id,267        });268        return;269    }270    const allocation_id = assignment.?.allocation_id;271    const existing = planBufferHandleFor(state.plan_buffers.items, allocation_id);272    if (existing != null) return error.InvalidArtifact;273    const buffer = try importSlotBuffer(state, slot.*, output, min_buffer_alignment);274    try state.plan_buffers.append(state.allocator, .{275        .allocation_id = allocation_id,276        .handle = buffer,277    });278}279280fn importSlotBuffer(281    state: *PreparedLaunchBindingsState,282    slot: PlannedSlot,283    bytes: []u8,284    min_buffer_alignment: u32,285) !gpu.BufferHandle {286    const buffer = try state.handle.importBuffer(.{287        .bytes = bytes,288        .alignment = try backendAllocationAlignment(slot.alignment, min_buffer_alignment),289        .dtype = slot.dtype,290        .element_count = slot.element_count,291    });292    errdefer state.handle.destroyObject(buffer.id);293    try state.live_buffers.append(state.allocator, buffer);294    state.borrowed_count += 1;295    return buffer;296}297298fn slotByteSize(slot: PlannedSlot) !usize {299    const byte_size = slot.byte_size orelse return error.UnsupportedOperation;300    return std.math.cast(usize, byte_size) orelse error.InvalidArtifact;301}302303fn slotIsWritten(304    artifact_plan: *const BackendArtifactPlan,305    launch_plan: plan_mod.LaunchGraphPlan,306    slot_id: usize,307) bool {308    for (artifact_plan.kernels.items) |kernel| {309        if (kernel.output_slot_id == slot_id) return true;310    }311    return slotIsCarried(launch_plan, slot_id);312}313314fn slotIsCarried(launch_plan: plan_mod.LaunchGraphPlan, slot_id: usize) bool {315    for (launch_plan.loops) |loop| {316        for (loop.carries) |carry| {317            if (carry.initial_slot_id == slot_id or carry.input_slot_id == slot_id) return true;318            if (carry.output_slot_id == slot_id or carry.final_slot_id == slot_id) return true;319        }320    }321    return false;322}323324fn expectDisjointOutputs(325    inputs: []const []const u8,326    outputs: []const []u8,327) gpu.BackendError!void {328    for (outputs, 0..) |output, index| {329        for (inputs) |input| {330            if (bytesOverlap(output, input)) return error.InvalidBuffer;331        }332        for (outputs[0..index]) |earlier| {333            if (bytesOverlap(output, earlier)) return error.InvalidBuffer;334        }335    }336}337338fn bytesOverlap(a: []const u8, b: []const u8) bool {339    if (a.len == 0 or b.len == 0) return false;340    const a_start = @intFromPtr(a.ptr);341    const b_start = @intFromPtr(b.ptr);342    return a_start < b_start + b.len and b_start < a_start + a.len;343}344345fn bindPlannedSlots(346    state: *PreparedLaunchBindingsState,347    artifact_plan: *const BackendArtifactPlan,348    memory_plan: *const BackendMemoryPlan,349    min_buffer_alignment: u32,350) !void {351    const allocator = state.allocator;352    for (artifact_plan.slots) |slot| {353        if (slot.role.input) continue;354        if (slot.role.constant) {355            const used = slot.role.output or plannedSlotHasKernelUse(artifact_plan, slot.slot_id);356            if (!used) continue;357            try bindConstantSlot(state, slot, min_buffer_alignment);358            continue;359        }360        if (!slot.hasStaticSize()) return error.UnsupportedOperation;361        const assignment = memory_plan.assignmentForSlot(slot.slot_id) orelse continue;362        const planned = planBufferHandleFor(state.plan_buffers.items, assignment.allocation_id);363        const slot_handle = planned orelse blk: {364            const allocated = try allocateBackendSlotBuffer(365                allocator,366                state.handle,367                &state.live_buffers,368                slot,369                min_buffer_alignment,370            );371            try state.plan_buffers.append(allocator, .{372                .allocation_id = assignment.allocation_id,373                .handle = allocated,374            });375            break :blk allocated;376        };377        try appendSlotBinding(allocator, &state.slot_bindings, slot.slot_id, slot_handle);378    }379}380381fn bindConstantSlot(382    state: *PreparedLaunchBindingsState,383    slot: PlannedSlot,384    min_buffer_alignment: u32,385) !void {386    const bytes = slot.constantBytes() orelse return error.InvalidArtifact;387    if (bytes.len != try slotByteSize(slot)) return error.InvalidArtifact;388    const buffer = try allocateBackendSlotBuffer(389        state.allocator,390        state.handle,391        &state.live_buffers,392        slot,393        min_buffer_alignment,394    );395    try state.handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });396    try appendSlotBindingWithAccess(397        state.allocator,398        &state.slot_bindings,399        slot.slot_id,400        buffer,401        .read_only,402    );403}404405fn bindElementCounts(406    state: *PreparedLaunchBindingsState,407    artifact_plan: *const BackendArtifactPlan,408    min_buffer_alignment: u32,409) !void {410    for (artifact_plan.kernels.items) |kernel| {411        if (kernel.element_count_argument != .device_buffer_u32) continue;412        const count_handle = try allocateBackendElementCountBuffer(413            state.allocator,414            state.handle,415            &state.live_buffers,416            min_buffer_alignment,417        );418        if (kernel.element_count > std.math.maxInt(u32)) return error.LaunchArgumentMismatch;419        var count_value: u32 = @intCast(kernel.element_count);420        try state.handle.writeBuffer(.{421            .handle = count_handle,422            .bytes = std.mem.asBytes(&count_value),423        });424        try state.count_bindings.append(state.allocator, .{425            .kernel_id = kernel.kernel_id,426            .binding = executableBackendBinding(count_handle, .read_write),427        });428    }429}430431fn allocateBackendSlotBuffer(432    allocator: std.mem.Allocator,433    handle: gpu.BackendHandle,434    list: *std.ArrayList(gpu.BufferHandle),435    slot: artifact_product.PlannedSlot,436    min_buffer_alignment: u32,437) !gpu.BufferHandle {438    const byte_size_u64 = slot.byte_size orelse return error.UnsupportedOperation;439    const byte_size = std.math.cast(usize, byte_size_u64) orelse return error.InvalidArtifact;440    const buffer = try handle.allocateBuffer(.{441        .byte_size = byte_size,442        .alignment = try backendAllocationAlignment(slot.alignment, min_buffer_alignment),443        .dtype = slot.dtype,444        .element_count = slot.element_count,445    });446    errdefer handle.destroyObject(buffer.id);447    try list.append(allocator, buffer);448    return buffer;449}450451fn allocateBackendElementCountBuffer(452    allocator: std.mem.Allocator,453    handle: gpu.BackendHandle,454    list: *std.ArrayList(gpu.BufferHandle),455    min_buffer_alignment: u32,456) !gpu.BufferHandle {457    const buffer = try handle.allocateBuffer(.{458        .byte_size = @sizeOf(u32),459        .alignment = try backendAllocationAlignment(64, min_buffer_alignment),460        .dtype = .i32,461        .element_count = 1,462    });463    errdefer handle.destroyObject(buffer.id);464    try list.append(allocator, buffer);465    return buffer;466}467468fn backendMinBufferAlignment(handle: gpu.BackendHandle) !u32 {469    const caps = try handle.queryCapabilities();470    return caps.memory.min_buffer_alignment;471}472473fn backendAllocationAlignment(preferred: u64, min_buffer_alignment: u32) !u32 {474    const aligned = @max(preferred, @as(u64, min_buffer_alignment));475    return std.math.cast(u32, aligned) orelse return error.InvalidArtifact;476}477478fn appendSlotBinding(479    allocator: std.mem.Allocator,480    slot_bindings: *std.ArrayList(SlotBinding),481    slot_id: usize,482    handle: gpu.BufferHandle,483) !void {484    try appendSlotBindingWithAccess(allocator, slot_bindings, slot_id, handle, .read_write);485}486487fn appendSlotBindingWithAccess(488    allocator: std.mem.Allocator,489    slot_bindings: *std.ArrayList(SlotBinding),490    slot_id: usize,491    handle: gpu.BufferHandle,492    access: gpu.BufferAccess,493) !void {494    try slot_bindings.append(allocator, .{495        .slot_id = slot_id,496        .binding = executableBackendBinding(handle, access),497    });498}499500fn executableBackendBinding(handle: gpu.BufferHandle, access: gpu.BufferAccess) gpu.BufferBinding {501    return .{502        .handle = handle,503        .access = access,504        .ownership = handle.ownership,505        .byte_size = handle.byte_size,506    };507}508509fn planBufferHandleFor(plan_buffers: []const BackendPlanBuffer, allocation_id: usize) ?gpu.BufferHandle {510    for (plan_buffers) |plan_buffer| {511        if (plan_buffer.allocation_id == allocation_id) return plan_buffer.handle;512    }513    return null;514}515516fn plannedSlotHasKernelUse(517    artifact_plan: *const artifact_product.BackendArtifactPlan,518    slot_id: usize,519) bool {520    for (artifact_plan.kernels.items) |kernel| {521        if (kernel.output_slot_id == slot_id) return true;522        for (kernel.input_slot_ids) |input_slot_id| {523            if (input_slot_id == slot_id) return true;524        }525    }526    return false;527}528529pub fn bindingForSlot(530    bindings: []const SlotBinding,531    slot_id: usize,532    access: gpu.BufferAccess,533) gpu.BackendError!gpu.BufferBinding {534    for (bindings) |slot_binding| {535        if (slot_binding.slot_id != slot_id) continue;536        return .{537            .handle = slot_binding.binding.handle,538            .access = access,539            .ownership = slot_binding.binding.ownership,540            .byte_size = slot_binding.binding.byte_size,541        };542    }543    return error.InvalidBuffer;544}545546pub fn countBinding(binding: gpu.BufferBinding) gpu.BufferBinding {547    return .{548        .handle = binding.handle,549        .access = .read_only,550        .ownership = binding.ownership,551        .byte_size = binding.byte_size,552    };553}554555pub fn elementCountBufferForKernel(556    bindings: []const ElementCountBufferBinding,557    kernel_id: usize,558) ?gpu.BufferBinding {559    for (bindings) |binding| {560        if (binding.kernel_id == kernel_id) return binding.binding;561    }562    return null;563}564565test "borrowed bindings run a cpu fragment on caller memory without copies" {566    try @import("../fixture/root.zig").requireNativeCpuArtifacts();567    const fixture = @import("fixture.zig");568    const fragment_mod = @import("fragment.zig");569    const invocation_mod = @import("invocation.zig");570    const allocator = std.testing.allocator;571572    var state = gpu.cpu.State.init(allocator);573    defer state.deinit();574    const handle = state.handle();575    const module = try fixture.addSemanticModule(allocator, "borrowed_bindings_add");576    const compiled = try fragment_mod.compileFragmentFromSemanticModule(577        allocator,578        handle,579        module,580        .{ .artifact_format = .cpu_object },581    );582    var fragment = try fragment_mod.loadFragment(allocator, handle, compiled, .{583        .artifact_format = .cpu_object,584    });585    defer fragment.deinit();586    const artifact_plan = compiled.artifactPlan();587    const launch_plan = compiled.launchPlan();588589    const lhs = try allocator.alignedAlloc(f32, .@"64", 8);590    defer allocator.free(lhs);591    const rhs = try allocator.alignedAlloc(f32, .@"64", 8);592    defer allocator.free(rhs);593    const sum = try allocator.alignedAlloc(f32, .@"64", 8);594    defer allocator.free(sum);595    for (lhs, rhs, 0..) |*left, *right, index| {596        left.* = @floatFromInt(index);597        right.* = 0.25 * @as(f32, @floatFromInt(index));598    }599    const inputs = [_][]const u8{ std.mem.sliceAsBytes(lhs), std.mem.sliceAsBytes(rhs) };600    const outputs = [_][]u8{std.mem.sliceAsBytes(sum)};601602    const bindings = try prepareBorrowed(603        allocator,604        handle,605        artifact_plan,606        launch_plan,607        &inputs,608        &outputs,609    );610    defer bindings.deinit();611    try std.testing.expectEqual(@as(usize, 3), borrowedBufferCount(bindings));612    try fragment.submitInvocationWithOptions(allocator, bindings, .{});613    try fragment.completeInvocationWithOptions(.{});614    try completeBorrowed(bindings, &outputs);615616    var copied: [8]f32 = undefined;617    const copied_outputs = [_][]u8{std.mem.sliceAsBytes(&copied)};618    try invocation_mod.run(fragment, allocator, allocator, &inputs, &copied_outputs);619    try std.testing.expectEqualSlices(u8, copied_outputs[0], outputs[0]);620621    const overlapping = [_][]u8{std.mem.sliceAsBytes(lhs)};622    try std.testing.expectError(623        error.InvalidBuffer,624        prepareBorrowed(allocator, handle, artifact_plan, launch_plan, &inputs, &overlapping),625    );626    const short = [_][]u8{std.mem.sliceAsBytes(sum)[0..4]};627    try std.testing.expectError(628        error.InvalidBuffer,629        prepareBorrowed(allocator, handle, artifact_plan, launch_plan, &inputs, &short),630    );631}

1pub const binding = @import("binding.zig");