lib/accy/src/target/abi.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const gpu = @import("gpu");
 3 const choir_abi = @import("choir_abi");
 4 
 5 pub fn argumentCount(format: gpu.ArtifactFormat, count: u32) gpu.BackendError!u32 {
 6     if (!gpu.artifactFormatUsesHostLoopLaunch(format)) return count;
 7     return choir_abi.kernelArgumentCount(count);
 8 }
 9 
10 pub 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)) orelse
18         return error.LaunchArgumentMismatch;
19     return .{ .grid = .{ count, 1, 1 }, .threadgroup = .{ 1, 1, 1 } };
20 }
21 
22 pub 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 }
35 
36 test "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 }