lib/accy/src/preparation/kernelization/model/outline.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const records = @import("../../../choir/root.zig").record;
 2 const std = @import("std");
 3 const choir = @import("choir");
 4 
 5 const ir = choir.ir;
 6 
 7 pub const kernel_outline_plan_analysis_name = "accy-choir-kernel-outline-plan";
 8 
 9 pub const KernelOutlineKind = records.kernel.KernelOutlineKind;
10 
11 pub const KernelOutline = struct {
12     id: usize,
13     name: []u8,
14     kind: KernelOutlineKind,
15     work_item_id: usize,
16     root: *ir.Operation,
17     input_slot_ids: []usize,
18     output_slot_id: usize,
19     element_count: u64,
20     op_count: usize,
21 
22     pub fn inputCount(self: KernelOutline) usize {
23         return self.input_slot_ids.len;
24     }
25 
26     fn deinit(self: *KernelOutline, allocator: std.mem.Allocator) void {
27         allocator.free(self.name);
28         allocator.free(self.input_slot_ids);
29         self.* = undefined;
30     }
31 };
32 
33 pub const KernelOutlinePlanAnalysis = struct {
34     allocator: std.mem.Allocator,
35     kernels: std.ArrayListUnmanaged(KernelOutline),
36     work_to_kernel: std.AutoHashMap(usize, usize),
37     total_input_slots: usize = 0,
38     total_scheduled_ops: usize = 0,
39 
40     pub fn init(allocator: std.mem.Allocator) KernelOutlinePlanAnalysis {
41         return .{
42             .allocator = allocator,
43             .kernels = .empty,
44             .work_to_kernel = std.AutoHashMap(usize, usize).init(allocator),
45         };
46     }
47 
48     pub fn deinit(self: *KernelOutlinePlanAnalysis) void {
49         for (self.kernels.items) |*kernel| {
50             kernel.deinit(self.allocator);
51         }
52         self.kernels.deinit(self.allocator);
53         self.work_to_kernel.deinit();
54         self.* = undefined;
55     }
56 
57     pub fn kernelCount(self: KernelOutlinePlanAnalysis) usize {
58         return self.kernels.items.len;
59     }
60 
61     pub fn getKernelForWork(
62         self: *const KernelOutlinePlanAnalysis,
63         work_item_id: usize,
64     ) ?*const KernelOutline {
65         const index = self.work_to_kernel.get(work_item_id) orelse return null;
66         return &self.kernels.items[index];
67     }
68 };