lib/accy/src/kernel/library/random/fold/threefry/family.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const choir_abi = @import("choir_abi");
2 const random = @import("../../root.zig");
3
4 const base = random.base;
5 const block = random.block;
6 const runtime = @import("runtime.zig");
7 const types = @import("types.zig");
8
9 const std = base.std;
10 const artifact_product = base.artifact_product;
11 const shape = base.shape;
12 const entry = base.entry;
13 const kernel = base.kernel;
14 const tuning = base.tuning;
15 const runtimeExtentArgument = base.runtimeExtentArgument;
16
17 const threefryBlock = block.threefryBlock;
18 const uniformFromBits = block.uniformFromBits;
19 const randomRuntimeExtentBounds = base.randomRuntimeExtentBounds;
20 const randomShapeFamily = base.randomShapeFamily;
21 const randomFoldDerivedLaunch = base.randomFoldDerivedLaunch;
22
23 pub const threefry_fold_family_version = types.threefry_fold_family_version;
24 pub const ThreefryFold = types.ThreefryFold;
25 pub const ThreefryFoldRuntimeFamilyF32 = runtime.ThreefryFoldRuntimeFamilyF32;
26 pub const ThreefryFoldRuntimeFamilyI32 = runtime.ThreefryFoldRuntimeFamilyI32;
27
28 pub fn threefryFoldUniformReference(rounds: u32, output_index: u32, samples: u32, key: [2]u32) f32 {
29 var accumulator: f32 = 0.0;
30 for (0..samples) |sample| {
31 const words = threefryBlock(rounds, .{ output_index, @intCast(sample) }, key);
32 for (words) |word| accumulator += uniformFromBits(word);
33 }
34 return accumulator;
35 }
36
37 pub fn threefryFoldBitsReference(rounds: u32, output_index: u32, samples: u32, key: [2]u32) u32 {
38 var accumulator: u32 = 0;
39 for (0..samples) |sample| {
40 const words = threefryBlock(rounds, .{ output_index, @intCast(sample) }, key);
41 for (words) |word| accumulator ^= word;
42 }
43 return accumulator;
44 }
45
46 pub fn threefryFoldFamilyTarget(allocator: std.mem.Allocator, instance: ThreefryFold) ![]u8 {
47 return std.fmt.allocPrint(
48 allocator,
49 "accy.kernel.random.threefry_fold_family_{d}r_{d}_{s}",
50 .{ instance.rounds, instance.threads, instance.dtype.name() },
51 );
52 }
53
54 pub fn threefryFoldFamilyEntryName(allocator: std.mem.Allocator, instance: ThreefryFold) ![]u8 {
55 return std.fmt.allocPrint(
56 allocator,
57 "accy_kernel_random_threefry_fold_family_{d}r_{d}_{s}",
58 .{ instance.rounds, instance.threads, instance.dtype.name() },
59 );
60 }
61 pub fn threefryFoldRuntimeArguments(instance: ThreefryFold) ![4]choir_abi.ScalarArgument {
62 if (instance.samples == 0) return error.ExtentOverflowsIndexRange;
63 return .{
64 .{ .u32 = try runtimeExtentArgument(instance.count) },
65 .{ .u32 = instance.samples },
66 .{ .u32 = instance.seedLo() },
67 .{ .u32 = instance.seedHi() },
68 };
69 }
70 pub fn threefryFoldShapeProfileDimensions(instance: ThreefryFold) [2]artifact_product.KernelCallShapeProfileDimension {
71 return .{
72 .{ .name = instance.count_axis, .runtime_scalar_argument_index = 0, .bounds = randomRuntimeExtentBounds() },
73 .{ .name = "k", .runtime_scalar_argument_index = 1, .bounds = randomRuntimeExtentBounds() },
74 };
75 }
76 pub fn threefryFoldTuningExtents(instance: ThreefryFold) [2]u64 {
77 return .{ instance.count, instance.samples };
78 }
79
80 pub fn threefryFoldTuningOperation(instance: ThreefryFold) entry.Operation {
81 return .{ .random = .{ .threefry_fold = instance.rounds } };
82 }
83
84 pub fn threefryFoldFamilyTuningKey(
85 backing_allocator: std.mem.Allocator,
86 device_fingerprint: u64,
87 instance: ThreefryFold,
88 ) !tuning.FamilyTuningKey {
89 const family_fingerprint = try threefryFoldFamilyFingerprint(backing_allocator, instance);
90 const extents = threefryFoldTuningExtents(instance);
91 return tuning.FamilyTuningKey.init(
92 device_fingerprint,
93 family_fingerprint,
94 entry.operationFingerprint(threefryFoldTuningOperation(instance)),
95 instance.dtype,
96 threefry_fold_family_version,
97 extents[0..],
98 ) orelse unreachable;
99 }
100
101 pub fn threefryFoldFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: ThreefryFold) !u64 {
102 var family = try randomShapeFamily(backing_allocator, "threefry_fold", instance.count_axis);
103 defer family.deinit();
104 return shape.fingerprint(family);
105 }
106 pub fn createThreefryFoldFamilyArtifact(
107 allocator: std.mem.Allocator,
108 handle: kernel.BackendHandle,
109 instance: ThreefryFold,
110 options: entry.ArtifactOptions,
111 ) !kernel.OwnedKernelCallArtifact {
112 const target = try threefryFoldFamilyTarget(allocator, instance);
113 defer allocator.free(target);
114 const entry_name = try threefryFoldFamilyEntryName(allocator, instance);
115 defer allocator.free(entry_name);
116 const family_fingerprint = options.shape_family_fingerprint orelse try threefryFoldFamilyFingerprint(allocator, instance);
117 const shape_profile_dimensions = threefryFoldShapeProfileDimensions(instance);
118 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
119 .name = "threefry_fold",
120 .fingerprint = family_fingerprint,
121 .dimensions = shape_profile_dimensions[0..],
122 };
123
124 var graph = switch (instance.dtype) {
125 .f32 => try ThreefryFoldRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
126 .i32 => try ThreefryFoldRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),
127 else => return error.UnsupportedDType,
128 };
129 defer graph.deinit();
130 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
131 .target = target,
132 .version = threefry_fold_family_version,
133 .format = options.format,
134 .kernel_plan = options.kernel_plan,
135 .element_count_argument = options.element_count_argument,
136 .shape_family_fingerprint = family_fingerprint,
137 .shape_profile = shape_profile,
138 .launch = options.launch orelse try randomFoldDerivedLaunch(instance.threads),
139 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 4 else options.runtime_scalar_argument_count,
140 .static_arguments = options.static_arguments,
141 });
142 }