tiny.accy.target.abi

1const std = @import("std");2const gpu = @import("gpu");3const choir_abi = @import("choir_abi");45pub fn argumentCount(format: gpu.ArtifactFormat, count: u32) gpu.BackendError!u32 {6    if (!gpu.artifactFormatUsesHostLoopLaunch(format)) return count;7    return choir_abi.kernelArgumentCount(count);8}910pub fn launchGeometry(11    format: gpu.ArtifactFormat,12    element_count: u64,13    geometry: ?choir_abi.LaunchGeometry,14) gpu.BackendError!?choir_abi.LaunchGeometry {15    if (!gpu.artifactFormatUsesHostLoopLaunch(format)) return geometry;16    if (geometry) |value| return value;17    const count = std.math.cast(u32, @max(element_count, 1)) orelse18        return error.LaunchArgumentMismatch;19    return .{ .grid = .{ count, 1, 1 }, .threadgroup = .{ 1, 1, 1 } };20}2122pub fn staticArguments(23    allocator: std.mem.Allocator,24    format: gpu.ArtifactFormat,25    element_count: u64,26    geometry: ?choir_abi.LaunchGeometry,27) gpu.BackendError![]choir_abi.ScalarArgument {28    if (!gpu.artifactFormatUsesHostLoopLaunch(format)) return &.{};29    return choir_abi.launchShapeArguments(30        allocator,31        element_count,32        (try launchGeometry(format, element_count, geometry)).?,33    );34}3536test "target ABI preserves host loop and device launch contracts" {37    const allocator = std.testing.allocator;38    try std.testing.expectEqual(@as(u32, 3), try argumentCount(.cuda_ptx, 3));39    try std.testing.expectEqual(@as(u32, 10), try argumentCount(.cpu_object, 3));40    const arguments = try staticArguments(allocator, .cpu_object, 5, null);41    defer allocator.free(arguments);42    const expected = [_]u32{ 5, 5, 1, 1, 1, 1, 1 };43    try std.testing.expectEqual(expected.len, arguments.len);44    for (arguments, expected) |argument, value| try std.testing.expectEqual(value, argument.u32);45    const device = try staticArguments(allocator, .cuda_ptx, 5, null);46    defer allocator.free(device);47    try std.testing.expectEqual(@as(usize, 0), device.len);48    try std.testing.expectEqual(null, try launchGeometry(.cuda_ptx, 5, null));49    try std.testing.expectError(error.InvalidArtifact, argumentCount(50        .cpu_object,51        std.math.maxInt(u32),52    ));53    try std.testing.expectError(error.LaunchArgumentMismatch, launchGeometry(54        .cpu_object,55        std.math.maxInt(u64),56        null,57    ));58}

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