lib/accy/src/kernel/library/catalog/family/random/philox.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 const descriptor_mod = @import("../../root.zig");
 4 const random = @import("../../../root.zig").random;
 5 const match_mod = @import("../../match/root.zig");
 6 const query_mod = descriptor_mod;
 7 
 8 const candidate_mod = @import("candidate.zig");
 9 
10 const Descriptor = descriptor_mod.Descriptor;
11 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
12 const RandomCandidateDescriptors = candidate_mod.RandomCandidateDescriptors;
13 const RandomQuery = query_mod.RandomQuery;
14 
15 pub fn select(backing_allocator: std.mem.Allocator, query: RandomQuery) !?OwnedDescriptor {
16     return descriptorForInstance(backing_allocator, familyInstance(query) orelse return null, query);
17 }
18 
19 pub fn appendCandidates(
20     backing_allocator: std.mem.Allocator,
21     query: RandomQuery,
22     result: *RandomCandidateDescriptors,
23 ) !void {
24     const thread_candidates = random.philoxThreadCandidatesForCount(query.count);
25     for (thread_candidates.slice()) |threads| {
26         const instance = random.Philox{
27             .count = query.count,
28             .rounds = match_mod.randomEffectiveRounds(query),
29             .dtype = query.dtype,
30             .threads = threads,
31         };
32         if (try descriptorForInstance(backing_allocator, instance, query)) |descriptor| {
33             result.items[result.count] = descriptor;
34             result.count += 1;
35         }
36     }
37 }
38 
39 fn descriptorForInstance(
40     backing_allocator: std.mem.Allocator,
41     instance: random.Philox,
42     query: RandomQuery,
43 ) !?OwnedDescriptor {
44     var specialization = try random.philoxFamilySpecialization(backing_allocator, instance);
45     errdefer specialization.deinit();
46     const lifetime_allocator = specialization.allocator();
47     const descriptor = Descriptor{
48         .name = try random.philoxFamilyEntryName(lifetime_allocator, instance),
49         .metadata = .{
50             .target = try random.philoxFamilyTarget(lifetime_allocator, instance),
51             .version = random.philox_family_version,
52             .layer = .logical,
53             .category = .random,
54             .specialization = specialization.value,
55         },
56     };
57     if (!match_mod.randomDescriptorMatches(descriptor, query)) {
58         specialization.deinit();
59         return null;
60     }
61     return .{ .descriptor = descriptor, .specialization = specialization };
62 }
63 
64 fn familyInstance(query: RandomQuery) ?random.Philox {
65     var instance = random.Philox{
66         .count = query.count,
67         .rounds = match_mod.randomEffectiveRounds(query),
68         .dtype = query.dtype,
69         .threads = random.philoxThreadsForCount(query.count),
70     };
71     if (query.schedule) |requested| {
72         switch (requested) {
73             .thread_blocks => |threads| {
74                 if (threads == 0) return null;
75                 if (@as(u64, threads) > instance.generators()) return null;
76                 instance.threads = threads;
77             },
78         }
79     }
80     return instance;
81 }