lib/accy/src/kernel/library/random/base.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const choir_abi = @import("choir_abi");
 2 pub const std = @import("std");
 3 
 4 pub const artifact_product = @import("../../../artifact/model/root.zig");
 5 pub const shape = @import("../../../choir/shape/root.zig");
 6 pub const library = @import("../root.zig");
 7 pub const kernel = @import("../../root.zig");
 8 
 9 pub const entry = library.entry;
10 pub const extent_mod = library.extent;
11 pub const geometry_mod = library.geometry;
12 pub const tuning = library.tuning;
13 
14 pub const DType = choir_abi.DType;
15 pub const indexExtent = extent_mod.indexExtent;
16 pub const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
17 pub const random_thread_caps = geometry_mod.ThreadCaps1D{};
18 
19 pub fn randomRuntimeExtentBounds() shape.Bounds {
20     return .{ .min = 1, .max = extent_mod.runtime_extent_max };
21 }
22 
23 pub fn randomDerivedLaunch(threads: u32, lanes: u64) !artifact_product.KernelCallLaunch {
24     if (threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
25     const divisor = std.math.cast(u32, @as(u64, threads) * lanes) orelse
26         return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
27     return .{ .derived = .{
28         .grid = .{
29             .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = divisor } },
30             .{ .fixed = 1 },
31             .{ .fixed = 1 },
32         },
33         .threadgroup = .{ threads, 1, 1 },
34     } };
35 }
36 
37 pub fn randomFoldDerivedLaunch(threads: u32) !artifact_product.KernelCallLaunch {
38     if (threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
39     return .{ .derived = .{
40         .grid = .{
41             .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = threads } },
42             .{ .fixed = 1 },
43             .{ .fixed = 1 },
44         },
45         .threadgroup = .{ threads, 1, 1 },
46     } };
47 }
48 
49 pub fn randomShapeFamily(backing_allocator: std.mem.Allocator, name: []const u8, count_axis: []const u8) !shape.Family {
50     var builder = try shape.Builder.init(backing_allocator, name);
51     errdefer builder.deinit();
52 
53     const count = try builder.symbol(count_axis);
54     const count_expr = try builder.symbolExpression(count);
55     _ = try builder.tensor("out", &.{count_expr});
56     try builder.assumeBounds(count_expr, randomRuntimeExtentBounds());
57 
58     return builder.finish();
59 }
60 
61 pub fn ceilDivComptime(comptime numerator: u64, comptime denominator: u32) u32 {
62     return @intCast(numerator / denominator + @as(u64, @intFromBool(numerator % denominator != 0)));
63 }