lib/accy/src/kernel/library/random/philox/profile.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const choir_abi = @import("choir_abi");
 2 const random = @import("../root.zig");
 3 const types = @import("types.zig");
 4 
 5 const base = random.base;
 6 const block = random.block;
 7 
 8 const std = base.std;
 9 const artifact_product = base.artifact_product;
10 const shape = base.shape;
11 const entry = base.entry;
12 const runtimeExtentArgument = base.runtimeExtentArgument;
13 const randomRuntimeExtentBounds = base.randomRuntimeExtentBounds;
14 const randomShapeFamily = base.randomShapeFamily;
15 const randomDTypeSupported = block.randomDTypeSupported;
16 
17 const Philox = types.Philox;
18 
19 pub fn philoxRuntimeArguments(instance: Philox) ![3]choir_abi.ScalarArgument {
20     return .{
21         .{ .u32 = try runtimeExtentArgument(instance.count) },
22         .{ .u32 = instance.seedLo() },
23         .{ .u32 = instance.seedHi() },
24     };
25 }
26 
27 pub fn philoxShapeProfileDimensions(instance: Philox) [1]artifact_product.KernelCallShapeProfileDimension {
28     return .{
29         .{ .name = instance.count_axis, .runtime_scalar_argument_index = 0, .bounds = randomRuntimeExtentBounds() },
30     };
31 }
32 
33 pub fn philoxFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Philox) !u64 {
34     var family = try philoxShapeFamily(backing_allocator, instance);
35     defer family.deinit();
36     return shape.fingerprint(family);
37 }
38 
39 pub fn philoxShapeFamily(backing_allocator: std.mem.Allocator, instance: Philox) !shape.Family {
40     return randomShapeFamily(backing_allocator, "philox", instance.count_axis);
41 }
42 
43 pub fn philoxFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Philox) !entry.OwnedSpecialization {
44     var owned = entry.OwnedSpecialization.init(backing_allocator);
45     errdefer owned.deinit();
46     const lifetime_allocator = owned.allocator();
47 
48     const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
49     outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.count_axis, instance.count);
50 
51     owned.value = .{
52         .dtype = instance.dtype,
53         .operation = .{ .random = .{ .philox = instance.rounds } },
54         .inputs = &.{},
55         .outputs = outputs,
56         .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "e", instance.generators(), instance.threads),
57     };
58     owned.value.launch = owned.value.schedule.?.launch();
59     var family = try philoxShapeFamily(backing_allocator, instance);
60     errdefer family.deinit();
61     try owned.takeShapeFamily(&family);
62     return owned;
63 }
64 
65 pub fn philoxInstanceFromSpecialization(specialization: entry.Specialization) ?Philox {
66     if (!specialization.scheduleMatchesLaunch()) return null;
67     const operation = specialization.operation orelse return null;
68     const rounds = switch (operation) {
69         .random => |random_operation| switch (random_operation) {
70             .philox => |rounds| rounds,
71             else => return null,
72         },
73         else => return null,
74     };
75     const dtype = specialization.dtype orelse return null;
76     if (!randomDTypeSupported(dtype)) return null;
77     if (specialization.inputs.len != 0 or specialization.outputs.len != 1) return null;
78     if (specialization.reductions.len != 0) return null;
79     const output = specialization.outputs[0];
80     if (output.axes.len != 1) return null;
81     const launch = specialization.launch orelse return null;
82     if (launch.threadgroup[0] == 0) return null;
83     return .{
84         .count = output.axes[0].extent,
85         .rounds = rounds,
86         .dtype = dtype,
87         .threads = launch.threadgroup[0],
88         .count_axis = output.axes[0].name,
89     };
90 }