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 }