lib/accy/src/kernel/library/random/fold/squares/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 squaresBlock = block.squaresBlock;
18 const uniformFromBits = block.uniformFromBits;
19 const randomRuntimeExtentBounds = base.randomRuntimeExtentBounds;
20 const randomShapeFamily = base.randomShapeFamily;
21 const randomFoldDerivedLaunch = base.randomFoldDerivedLaunch;
22
23 pub const SquaresFold = types.SquaresFold;
24 pub const squares_fold_family_version = types.squares_fold_family_version;
25 pub const SquaresFoldRuntimeFamilyF32 = runtime.SquaresFoldRuntimeFamilyF32;
26 pub const SquaresFoldRuntimeFamilyI32 = runtime.SquaresFoldRuntimeFamilyI32;
27
28 pub fn squaresFoldUniformReference(output_index: u32, samples: u32, key: u64) f32 {
29 var accumulator: f32 = 0.0;
30 for (0..samples) |sample| {
31 const counter = @as(u64, output_index) | (@as(u64, @intCast(sample)) << 32);
32 accumulator += uniformFromBits(squaresBlock(counter, key));
33 }
34 return accumulator;
35 }
36
37 pub fn squaresFoldBitsReference(output_index: u32, samples: u32, key: u64) u32 {
38 var accumulator: u32 = 0;
39 for (0..samples) |sample| {
40 const counter = @as(u64, output_index) | (@as(u64, @intCast(sample)) << 32);
41 accumulator ^= squaresBlock(counter, key);
42 }
43 return accumulator;
44 }
45
46 pub fn squaresFoldFamilyTarget(allocator: std.mem.Allocator, instance: SquaresFold) ![]u8 {
47 return std.fmt.allocPrint(
48 allocator,
49 "accy.kernel.random.squares_fold_family_{d}_{s}",
50 .{ instance.threads, instance.dtype.name() },
51 );
52 }
53
54 pub fn squaresFoldFamilyEntryName(allocator: std.mem.Allocator, instance: SquaresFold) ![]u8 {
55 return std.fmt.allocPrint(
56 allocator,
57 "accy_kernel_random_squares_fold_family_{d}_{s}",
58 .{ instance.threads, instance.dtype.name() },
59 );
60 }
61
62 pub fn squaresFoldRuntimeArguments(instance: SquaresFold) ![4]choir_abi.ScalarArgument {
63 if (instance.samples == 0) return error.ExtentOverflowsIndexRange;
64 return .{
65 .{ .u32 = try runtimeExtentArgument(instance.count) },
66 .{ .u32 = instance.samples },
67 .{ .u32 = instance.keyLo() },
68 .{ .u32 = instance.keyHi() },
69 };
70 }
71
72 pub fn squaresFoldShapeProfileDimensions(instance: SquaresFold) [2]artifact_product.KernelCallShapeProfileDimension {
73 return .{
74 .{ .name = instance.count_axis, .runtime_scalar_argument_index = 0, .bounds = randomRuntimeExtentBounds() },
75 .{ .name = "k", .runtime_scalar_argument_index = 1, .bounds = randomRuntimeExtentBounds() },
76 };
77 }
78
79 pub fn squaresFoldTuningExtents(instance: SquaresFold) [2]u64 {
80 return .{ instance.count, instance.samples };
81 }
82
83 pub fn squaresFoldTuningOperation(instance: SquaresFold) entry.Operation {
84 _ = instance;
85 return .{ .random = .squares_fold };
86 }
87
88 pub fn squaresFoldFamilyTuningKey(
89 backing_allocator: std.mem.Allocator,
90 device_fingerprint: u64,
91 instance: SquaresFold,
92 ) !tuning.FamilyTuningKey {
93 const family_fingerprint = try squaresFoldFamilyFingerprint(backing_allocator, instance);
94 const extents = squaresFoldTuningExtents(instance);
95 return tuning.FamilyTuningKey.init(
96 device_fingerprint,
97 family_fingerprint,
98 entry.operationFingerprint(squaresFoldTuningOperation(instance)),
99 instance.dtype,
100 squares_fold_family_version,
101 extents[0..],
102 ) orelse unreachable;
103 }
104
105 pub fn squaresFoldFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: SquaresFold) !u64 {
106 var family = try randomShapeFamily(backing_allocator, "squares_fold", instance.count_axis);
107 defer family.deinit();
108 return shape.fingerprint(family);
109 }
110
111 pub fn createSquaresFoldFamilyArtifact(
112 allocator: std.mem.Allocator,
113 handle: kernel.BackendHandle,
114 instance: SquaresFold,
115 options: entry.ArtifactOptions,
116 ) !kernel.OwnedKernelCallArtifact {
117 const target = try squaresFoldFamilyTarget(allocator, instance);
118 defer allocator.free(target);
119 const entry_name = try squaresFoldFamilyEntryName(allocator, instance);
120 defer allocator.free(entry_name);
121 const family_fingerprint = options.shape_family_fingerprint orelse try squaresFoldFamilyFingerprint(allocator, instance);
122 const shape_profile_dimensions = squaresFoldShapeProfileDimensions(instance);
123 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
124 .name = "squares_fold",
125 .fingerprint = family_fingerprint,
126 .dimensions = shape_profile_dimensions[0..],
127 };
128
129 var graph = switch (instance.dtype) {
130 .f32 => try SquaresFoldRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
131 .i32 => try SquaresFoldRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),
132 else => return error.UnsupportedDType,
133 };
134 defer graph.deinit();
135 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
136 .target = target,
137 .version = squares_fold_family_version,
138 .format = options.format,
139 .kernel_plan = options.kernel_plan,
140 .element_count_argument = options.element_count_argument,
141 .shape_family_fingerprint = family_fingerprint,
142 .shape_profile = shape_profile,
143 .launch = options.launch orelse try randomFoldDerivedLaunch(instance.threads),
144 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 4 else options.runtime_scalar_argument_count,
145 .static_arguments = options.static_arguments,
146 });
147 }