lib/accy/src/choir/record/kernel.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const records = @import("root.zig");
3 const schedule = @import("../../kernel/model/root.zig").core.schedule;
4
5 pub const DotGeneralMmaTile = struct {
6 bm: u32 = 128,
7 bn: u32 = 64,
8 bk: u32 = 16,
9 warps_m: u32 = 2,
10 warps_n: u32 = 2,
11 stages: u32 = 1,
12 splits: u32 = 1,
13
14 pub fn warps(self: DotGeneralMmaTile) u32 {
15 return self.warps_m * self.warps_n;
16 }
17
18 pub fn warpTileM(self: DotGeneralMmaTile) u32 {
19 return self.bm / self.warps_m;
20 }
21
22 pub fn warpTileN(self: DotGeneralMmaTile) u32 {
23 return self.bn / self.warps_n;
24 }
25
26 pub fn threads(self: DotGeneralMmaTile) u32 {
27 return self.warps() * 32;
28 }
29
30 pub fn exact(self: DotGeneralMmaTile, dims: DotGeneralStaticDims) bool {
31 return dims.m % self.bm == 0 and dims.n % self.bn == 0 and dims.k % self.bk == 0;
32 }
33 };
34
35 pub const ReductionAtomicPlan = struct {
36 threads: u32,
37 blocks: u32,
38
39 pub fn stride(self: ReductionAtomicPlan) u64 {
40 return @as(u64, self.threads) * self.blocks;
41 }
42 };
43
44 pub const DotGeneralBlockTile = struct {
45 bm: u32 = 64,
46 bn: u32 = 64,
47 bk: u32 = 16,
48 tm: u32 = 4,
49 tn: u32 = 4,
50 splits: u32 = 1,
51 stages: u32 = 2,
52
53 pub fn threadsX(self: DotGeneralBlockTile) u32 {
54 return self.bn / self.tn;
55 }
56
57 pub fn threadsY(self: DotGeneralBlockTile) u32 {
58 return self.bm / self.tm;
59 }
60
61 pub fn exact(self: DotGeneralBlockTile, dims: DotGeneralStaticDims) bool {
62 return dims.m % self.bm == 0 and dims.n % self.bn == 0 and dims.k % self.bk == 0;
63 }
64 };
65
66 pub const DotGeneralStaticDims = struct {
67 m: u32,
68 n: u32,
69 k: u32,
70 batch: u32 = 1,
71 };
72
73 pub const ReductionWarpRowsPlan = struct {
74 threads: u32,
75 rows: u32,
76 };
77
78 pub const ElementwiseVectorPlan = struct {
79 quads: u64,
80 };
81
82 pub const LoweredKernelBody = union(enum) {
83 generic,
84 dot_block_tile: DotGeneralBlockTile,
85 dot_mma_tile: DotGeneralMmaTile,
86 reduction_single_block: u32,
87 reduction_atomic: ReductionAtomicPlan,
88 reduction_warp_rows: ReductionWarpRowsPlan,
89 row_pipeline: RowPipelinePlan,
90 flash_attention: FlashAttentionPlan,
91 scan: ScanPlan,
92 elementwise_rank2: ElementwiseRank2Plan,
93 elementwise_vector: ElementwiseVectorPlan,
94 };
95
96 pub const ElementwiseRank2Plan = struct {
97 rows: u32,
98 cols: u32,
99 threads_x: u32,
100 threads_y: u32,
101 };
102
103 pub const ScanPlan = struct {
104 blocks: u32,
105 threads: u32,
106 items: u32,
107 };
108
109 pub const FlashAttentionPlan = struct {
110 seq: u32,
111 dim: u32,
112 br: u32,
113 bc: u32,
114 threads_x: u32,
115 threads_y: u32,
116 };
117
118 pub const RowPipelinePlan = struct {
119 threads: u32,
120 rows: u32,
121 cols: u32,
122 warps: u32,
123 };
124
125 pub const GeneratedSchedule = struct {
126 kind: GeneratedScheduleKind,
127 threads: @import("../../kernel/model/root.zig").logical.schedule.Threads,
128 };
129
130 pub const GeneratedScheduleKind = enum {
131 flat,
132 matrix,
133 };
134
135 pub const KernelOutlineKind = enum {
136 elementwise,
137 shape,
138 dot_general,
139 reduction,
140 kernel_call,
141 row_pipeline,
142 iterate,
143 flash_attention,
144 scan,
145 };
146
147 pub const Outline = struct {
148 id: usize,
149 name: []const u8,
150 kind: KernelOutlineKind,
151 work_item_id: usize,
152 root: records.reference.Operation,
153 input_slot_ids: []const usize,
154 output_slot_id: usize,
155 element_count: u64,
156 op_count: usize,
157 };
158
159 pub const Outlines = struct {
160 kernels: []const Outline,
161 total_input_slots: usize,
162 total_scheduled_ops: usize,
163 };
164
165 pub const Lowered = struct {
166 work_item_id: usize,
167 entry_name: []const u8,
168 program: records.program.Record,
169 argument_count: u32,
170 body_fingerprint: u64,
171 dynamic_shared_memory_bytes: u32,
172 schedule: GeneratedSchedule,
173 launch: ?schedule.Launch,
174 output_fill_pattern: ?u32,
175 scratch_fill_pattern: ?u32,
176 body: LoweredKernelBody,
177 };
178
179 pub const Generated = struct {
180 kernels: []const Lowered,
181 };
182
183 pub const Record = struct {
184 outlines: Outlines,
185 generated: Generated,
186 };
187
188 pub fn validate(
189 allocator: std.mem.Allocator,
190 value: Record,
191 memory: records.memory.Record,
192 ) !void {
193 var work = std.AutoHashMapUnmanaged(usize, bool).empty;
194 defer work.deinit(allocator);
195 const count = std.math.cast(u32, value.outlines.kernels.len) orelse
196 return error.InvalidStageRecord;
197 try work.ensureTotalCapacity(allocator, count);
198 var inputs: usize = 0;
199 var operations: usize = 0;
200 for (value.outlines.kernels, 0..) |outline, index| {
201 if (outline.id != index or outline.output_slot_id >= memory.buffers.slots.len) {
202 return error.InvalidStageRecord;
203 }
204 for (outline.input_slot_ids) |slot| {
205 if (slot >= memory.buffers.slots.len) return error.InvalidStageRecord;
206 }
207 const entry = work.getOrPutAssumeCapacity(outline.work_item_id);
208 if (entry.found_existing) return error.InvalidStageRecord;
209 entry.value_ptr.* = false;
210 inputs = std.math.add(usize, inputs, outline.input_slot_ids.len) catch
211 return error.InvalidStageRecord;
212 operations = std.math.add(usize, operations, outline.op_count) catch
213 return error.InvalidStageRecord;
214 }
215 if (inputs != value.outlines.total_input_slots or
216 operations != value.outlines.total_scheduled_ops) return error.InvalidStageRecord;
217 for (value.generated.kernels) |kernel| {
218 const seen = work.getPtr(kernel.work_item_id) orelse return error.InvalidStageRecord;
219 if (seen.*) return error.InvalidStageRecord;
220 seen.* = true;
221 if (kernel.argument_count != kernel.program.params.len) return error.InvalidStageRecord;
222 }
223 }