lib/accy/src/preparation/kernelization/model/product.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const records = @import("../../../choir/root.zig").record;
2 const std = @import("std");
3 const gpu = @import("gpu");
4 const choir_abi = @import("choir_abi");
5 const choir = @import("choir");
6 const kernel_program = @import("../../../kernel/model/program/root.zig");
7 const kernel_schedule = @import("../../../kernel/model/core/schedule/root.zig");
8
9 const ir = choir.ir;
10 const dialects = choir.dialects;
11 const MemrefDialect = dialects.MemrefDialect;
12 const accounting = choir.passes.pass.work;
13
14 pub const ElementwiseKernel = enum {
15 add,
16 sub,
17 mul,
18 div,
19 min,
20 max,
21 neg,
22 abs,
23 sqrt,
24 exp,
25 log,
26 tanh,
27 sin,
28 cos,
29 tan,
30 floor,
31 round,
32 trunc,
33 pow,
34 atan2,
35 convert,
36 compare,
37 select,
38 };
39
40 pub const ShapeKernel = enum {
41 broadcast_in_dim,
42 iota,
43 reshape,
44 transpose,
45 slice,
46 pad,
47 concatenate,
48 gather,
49 scatter,
50 };
51
52 pub const GeneratedScheduleKind = records.kernel.GeneratedScheduleKind;
53
54 pub const GeneratedSchedule = records.kernel.GeneratedSchedule;
55
56 pub const GeneratedKernelProgram = kernel_program.Program;
57
58 pub const GeneratedKernelSummary = struct {
59 work_item_id: usize,
60 entry_name: []const u8,
61 argument_count: u32,
62 body_fingerprint: u64,
63 dynamic_shared_memory_bytes: u32 = 0,
64 schedule: GeneratedSchedule,
65 launch_geometry: ?choir_abi.LaunchGeometry,
66 };
67
68 pub const GeneratedKernelSummaries = struct {
69 allocator: std.mem.Allocator,
70 items: []GeneratedKernelSummary,
71
72 pub fn deinit(self: *GeneratedKernelSummaries) void {
73 for (self.items) |item| {
74 self.allocator.free(item.entry_name);
75 }
76 self.allocator.free(self.items);
77 self.* = undefined;
78 }
79
80 pub fn len(self: *const GeneratedKernelSummaries) usize {
81 return self.items.len;
82 }
83
84 pub fn summary(
85 self: *const GeneratedKernelSummaries,
86 kernel_index: usize,
87 ) !GeneratedKernelSummary {
88 if (kernel_index >= self.items.len) return error.InvalidIndex;
89 return self.items[kernel_index];
90 }
91
92 pub fn summaryForWork(
93 self: *const GeneratedKernelSummaries,
94 work_item_id: usize,
95 ) !GeneratedKernelSummary {
96 for (self.items) |item| {
97 if (item.work_item_id == work_item_id) return item;
98 }
99 return error.MissingKernelization;
100 }
101 };
102
103 pub const RowPipelinePlan = records.kernel.RowPipelinePlan;
104
105 pub const FlashAttentionPlan = records.kernel.FlashAttentionPlan;
106
107 pub const ScanPlan = records.kernel.ScanPlan;
108
109 pub const ElementwiseRank2Plan = records.kernel.ElementwiseRank2Plan;
110
111 pub const LoweredKernelBody = records.kernel.LoweredKernelBody;
112
113 pub const ElementwiseVectorPlan = records.kernel.ElementwiseVectorPlan;
114
115 pub const ReductionWarpRowsPlan = records.kernel.ReductionWarpRowsPlan;
116
117 pub const LoweredKernel = struct {
118 work_item_id: usize,
119 entry_name: []u8,
120 program: kernel_program.Program,
121 argument_count: u32,
122 body_fingerprint: u64,
123 dynamic_shared_memory_bytes: u32 = 0,
124 schedule: GeneratedSchedule,
125 launch: ?kernel_schedule.Launch = null,
126 output_fill_pattern: ?u32 = null,
127 scratch_fill_pattern: ?u32 = null,
128 body: LoweredKernelBody = .generic,
129
130 pub fn deinit(self: *LoweredKernel, allocator: std.mem.Allocator) void {
131 self.program.deinit();
132 allocator.free(self.entry_name);
133 self.* = undefined;
134 }
135
136 pub fn runtimeScalarArgumentCount(self: *const LoweredKernel) u32 {
137 var count: u32 = 0;
138 for (self.program.params()) |param| {
139 switch (param) {
140 .scalar => count += 1,
141 .buffer => {},
142 }
143 }
144 return count;
145 }
146
147 pub fn requiredDTypes(self: *const LoweredKernel) gpu.DTypeSet {
148 var dtypes: gpu.DTypeSet = .{};
149 for (self.program.params()) |param| {
150 switch (param) {
151 .scalar => |dtype| dtypes.insert(dtype),
152 .buffer => |buffer| dtypes.insert(buffer.dtype),
153 }
154 }
155 return dtypes;
156 }
157
158 pub fn launchGeometry(self: *const LoweredKernel) ?choir_abi.LaunchGeometry {
159 const launch_value = self.launch orelse return null;
160 return .{
161 .grid = launch_value.grid,
162 .threadgroup = launch_value.block,
163 .dynamic_shared_memory_bytes = self.dynamic_shared_memory_bytes,
164 };
165 }
166
167 pub fn summary(self: *const LoweredKernel) GeneratedKernelSummary {
168 return .{
169 .work_item_id = self.work_item_id,
170 .entry_name = self.entry_name,
171 .argument_count = self.argument_count,
172 .body_fingerprint = self.body_fingerprint,
173 .dynamic_shared_memory_bytes = self.dynamic_shared_memory_bytes,
174 .schedule = self.schedule,
175 .launch_geometry = self.launchGeometry(),
176 };
177 }
178 };
179
180 pub const KernelizationAnalysis = struct {
181 allocator: std.mem.Allocator,
182 context: *ir.Context,
183 kernels: std.ArrayListUnmanaged(LoweredKernel),
184 work_to_kernel: std.AutoHashMap(usize, usize),
185
186 /// The kernel stage calls this for its work bound to charge the result object before any kernel
187 /// is generated, so compilation can refuse a pass that would exceed the caller's limits. That
188 /// bound is the costs a pass declares before it runs. The bound covers the result object, its
189 /// shared compiler context, and the list of kernels and the map from a scheduled unit of work
190 /// to kernel, sized for `kernel_count` kernels. That context is the object that owns the
191 /// operations and values of the generated kernels. Generated programs, names, generation
192 /// scratch and the analyses it depends on are charged separately.
193 pub fn baseStorageBound(limits: ir.Context.Limits, kernel_count: u64) !u64 {
194 const context = ir.Context.Capacity.derive(limits) catch return error.WorkOverflow;
195 var bytes: u64 = @sizeOf(KernelizationAnalysis) + @alignOf(KernelizationAnalysis);
196 bytes = try accounting.add(bytes, @sizeOf(ir.Context) + @alignOf(ir.Context));
197 bytes = try accounting.add(bytes, context.storage_bytes);
198 bytes = try accounting.add(bytes, context.storage_alignment.toByteUnits());
199 const kernels = try accounting.arrayListGrowth(LoweredKernel, kernel_count);
200 bytes = try accounting.add(bytes, kernels);
201 bytes = try accounting.add(bytes, try accounting.hashMapGrowth(usize, usize, kernel_count));
202 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
203 return bytes;
204 }
205
206 pub fn reserveKernelCapacity(self: *KernelizationAnalysis, count: usize) !void {
207 const map_count = std.math.cast(u32, count) orelse return error.InvalidArtifact;
208 try self.kernels.ensureTotalCapacity(self.allocator, count);
209 try self.work_to_kernel.ensureTotalCapacity(map_count);
210 }
211
212 pub fn init(allocator: std.mem.Allocator, limits: ir.Context.Limits) !KernelizationAnalysis {
213 const context = try ir.Context.create(allocator, limits);
214 errdefer {
215 context.deinit(allocator);
216 allocator.destroy(context);
217 }
218 dialects.registerChoirDialect(context) catch |err| {
219 return if (err == error.OutOfMemory) error.WorkExhausted else err;
220 };
221 return .{
222 .allocator = allocator,
223 .context = context,
224 .kernels = .empty,
225 .work_to_kernel = std.AutoHashMap(usize, usize).init(allocator),
226 };
227 }
228
229 pub fn deinit(self: *KernelizationAnalysis) void {
230 for (self.kernels.items) |*kernel| {
231 kernel.deinit(self.allocator);
232 }
233 self.kernels.deinit(self.allocator);
234 self.work_to_kernel.deinit();
235 self.context.deinit(self.allocator);
236 self.allocator.destroy(self.context);
237 self.* = undefined;
238 }
239
240 pub fn kernelCount(self: KernelizationAnalysis) usize {
241 return self.kernels.items.len;
242 }
243
244 pub fn getForWork(
245 self: *const KernelizationAnalysis,
246 work_item_id: usize,
247 ) ?*const LoweredKernel {
248 const index = self.work_to_kernel.get(work_item_id) orelse return null;
249 return &self.kernels.items[index];
250 }
251
252 pub fn kernelSummary(
253 self: *const KernelizationAnalysis,
254 kernel_index: usize,
255 ) !GeneratedKernelSummary {
256 const kernel_value = try self.kernelAt(kernel_index);
257 return kernel_value.summary();
258 }
259
260 pub fn kernelProgram(
261 self: *const KernelizationAnalysis,
262 kernel_index: usize,
263 ) !*const GeneratedKernelProgram {
264 const kernel_value = try self.kernelAt(kernel_index);
265 return &kernel_value.program;
266 }
267
268 pub fn kernelSummaryForWork(
269 self: *const KernelizationAnalysis,
270 work_item_id: usize,
271 ) !GeneratedKernelSummary {
272 const kernel_value = self.getForWork(work_item_id) orelse return error.MissingKernelization;
273 return kernel_value.summary();
274 }
275
276 pub fn kernelProgramForWork(
277 self: *const KernelizationAnalysis,
278 work_item_id: usize,
279 ) !*const GeneratedKernelProgram {
280 const kernel_value = self.getForWork(work_item_id) orelse return error.MissingKernelization;
281 return &kernel_value.program;
282 }
283
284 pub fn copyKernelSummaries(
285 self: *const KernelizationAnalysis,
286 result_allocator: std.mem.Allocator,
287 ) !GeneratedKernelSummaries {
288 const items = try result_allocator.alloc(GeneratedKernelSummary, self.kernels.items.len);
289 var copied: usize = 0;
290 errdefer {
291 for (items[0..copied]) |summary| {
292 result_allocator.free(summary.entry_name);
293 }
294 result_allocator.free(items);
295 }
296
297 for (items, self.kernels.items) |*item, *kernel_value| {
298 item.* = kernel_value.summary();
299 item.entry_name = try result_allocator.dupe(u8, kernel_value.entry_name);
300 copied += 1;
301 }
302
303 return .{
304 .allocator = result_allocator,
305 .items = items,
306 };
307 }
308
309 fn kernelAt(
310 self: *const KernelizationAnalysis,
311 kernel_index: usize,
312 ) !*const LoweredKernel {
313 if (kernel_index >= self.kernels.items.len) return error.InvalidIndex;
314 return &self.kernels.items[kernel_index];
315 }
316 };
317
318 pub const DotGeneralStaticDims = records.kernel.DotGeneralStaticDims;
319
320 pub const DotGeneralDescription = struct {
321 input_dtype: choir_abi.DType,
322 output_dtype: choir_abi.DType,
323 dims: DotGeneralStaticDims,
324 };
325
326 pub const DotGeneralBlockTile = records.kernel.DotGeneralBlockTile;
327
328 pub const reduction_single_block_threads: u32 = 256;
329 const reduction_single_block_min_extent: u32 = 256;
330 pub const reduction_warp_rows_threads: u32 = 256;
331 const reduction_warp_rows_min_cols: u32 = 128;
332
333 pub fn reductionWarpRowsThreads(dims: ReductionStaticDims) ?u32 {
334 if (dims.input_rank != 2 or dims.axis != 1) return null;
335 if (dims.cols < reduction_warp_rows_min_cols) return null;
336 return reduction_warp_rows_threads;
337 }
338
339 pub fn reductionSingleBlockThreads(dims: ReductionStaticDims) ?u32 {
340 if (dims.output_element_count != 1) return null;
341 if (dims.input_element_count < reduction_single_block_min_extent) return null;
342 return reduction_single_block_threads;
343 }
344
345 pub const ReductionAtomicPlan = records.kernel.ReductionAtomicPlan;
346
347 const reduction_atomic_min_extent: u32 = 4096;
348 const reduction_atomic_max_blocks: u32 = 512;
349 const reduction_atomic_target_chain: u32 = 32;
350
351 pub fn reductionAtomicPlanFor(
352 dims: ReductionStaticDims,
353 kind: ReductionKind,
354 input_dtype: choir_abi.DType,
355 init_is_constant: bool,
356 ) ?ReductionAtomicPlan {
357 if (dims.output_element_count != 1) return null;
358 if (dims.input_element_count < reduction_atomic_min_extent) return null;
359 if (kind != .sum) return null;
360 if (!init_is_constant) return null;
361 switch (input_dtype) {
362 .f32, .i32, .u32 => {},
363 else => return null,
364 }
365 const threads = reduction_single_block_threads;
366 const per_block = @as(u64, threads) * reduction_atomic_target_chain;
367 const wanted = (dims.input_element_count + per_block - 1) / per_block;
368 const blocks: u32 = @intCast(@min(@max(wanted, 1), reduction_atomic_max_blocks));
369 return .{ .threads = threads, .blocks = blocks };
370 }
371
372 pub const DotGeneralMmaTile = records.kernel.DotGeneralMmaTile;
373
374 const mma_pipeline_max_blocks = 512;
375
376 pub fn dotGeneralMmaTileFor(desc: DotGeneralDescription) ?DotGeneralMmaTile {
377 if (desc.input_dtype != .f32 or desc.output_dtype != .f32) return null;
378 var tile = DotGeneralMmaTile{};
379 if (!tile.exact(desc.dims)) return null;
380 const blocks = @as(u64, desc.dims.batch) * (desc.dims.m / tile.bm) * (desc.dims.n / tile.bn);
381 if (blocks < mma_pipeline_max_blocks) {
382 tile.stages = 2;
383 }
384 if (desc.dims.batch == 1 and blocks < 32) {
385 while (tile.splits < split_k_max) {
386 const doubled = tile.splits * 2;
387 if (blocks * tile.splits >= 128) break;
388 if (desc.dims.k % doubled != 0) break;
389 const chunk = desc.dims.k / doubled;
390 if (chunk % (2 * @as(u64, tile.bk)) != 0 or chunk < 4 * @as(u64, tile.bk)) break;
391 tile.splits = doubled;
392 }
393 }
394 return tile;
395 }
396
397 const block_tile_min_blocks = 64;
398 const wide_block_tile_min_blocks = 128;
399
400 const wide_block_tile = DotGeneralBlockTile{ .bm = 128, .bn = 128, .bk = 16, .tm = 8, .tn = 8 };
401
402 const split_k_target_blocks = 256;
403 const split_k_max = 8;
404
405 pub fn dotGeneralBlockTileFor(desc: DotGeneralDescription) ?DotGeneralBlockTile {
406 if (desc.input_dtype != .f32 or desc.output_dtype != .f32) return null;
407 if (desc.dims.k >= 128 and blockTileFits(wide_block_tile, desc.dims, wide_block_tile_min_blocks)) return wide_block_tile;
408 var tile = DotGeneralBlockTile{ .stages = 4 };
409 if (!blockTileFits(tile, desc.dims, block_tile_min_blocks)) return null;
410 if (desc.dims.batch == 1) {
411 const tiles = (desc.dims.m / tile.bm) * (desc.dims.n / tile.bn);
412 while (tile.splits < split_k_max) {
413 const doubled = tile.splits * 2;
414 if (tiles * tile.splits >= split_k_target_blocks) break;
415 if (desc.dims.k % doubled != 0) break;
416 const chunk = desc.dims.k / doubled;
417 if (chunk % tile.bk != 0 or chunk < 2 * tile.bk) break;
418 tile.splits = doubled;
419 }
420 }
421 return tile;
422 }
423
424 fn blockTileFits(tile: DotGeneralBlockTile, dims: DotGeneralStaticDims, min_blocks: u64) bool {
425 if (dims.m < tile.bm or dims.n < tile.bn) return false;
426 if (dims.k < tile.bk) return false;
427 const threads = @as(u64, tile.threadsX()) * tile.threadsY();
428 if ((@as(u64, tile.bm) * tile.bk) % threads != 0) return false;
429 if ((@as(u64, tile.bk) * tile.bn) % threads != 0) return false;
430 const tiles_m = (dims.m + tile.bm - 1) / tile.bm;
431 const tiles_n = (dims.n + tile.bn - 1) / tile.bn;
432 const blocks = @as(u64, dims.batch) * tiles_m * tiles_n;
433 return blocks >= min_blocks;
434 }
435
436 pub const ReductionStaticDims = struct {
437 input_rank: u8,
438 axis: u8,
439 input_element_count: u32,
440 output_element_count: u32,
441 rows: u32 = 1,
442 cols: u32 = 1,
443 inner: u32 = 1,
444 };
445
446 pub const ReductionKind = enum {
447 sum,
448 max,
449 min,
450 };
451
452 pub const ReductionInitValue = union(enum) {
453 f32: f32,
454 i32: i32,
455 u32: u32,
456 f16: f16,
457 };
458
459 pub const ReductionInit = union(enum) {
460 constant: ReductionInitValue,
461 input_buffer,
462 };
463
464 pub const ReductionDescription = struct {
465 kind: ReductionKind,
466 input_dtype: choir_abi.DType,
467 output_dtype: choir_abi.DType,
468 dims: ReductionStaticDims,
469 init: ReductionInit,
470 };
471
472 pub fn addKernel(analysis: *KernelizationAnalysis, kernel: LoweredKernel) !void {
473 const index = analysis.kernels.items.len;
474 try analysis.work_to_kernel.put(kernel.work_item_id, index);
475 errdefer _ = analysis.work_to_kernel.remove(kernel.work_item_id);
476 try analysis.kernels.append(analysis.allocator, kernel);
477 }
478
479 pub fn dynamicSharedMemoryBytes(module: *ir.Operation) gpu.BackendError!u32 {
480 var state = DynamicSharedMemoryState{};
481 _ = module.walk(.{ .order = .pre_order }, &state, DynamicSharedMemoryState.visit) catch |err| switch (err) {
482 error.OutOfMemory => return error.OutOfMemory,
483 error.InvalidArtifact => return error.InvalidArtifact,
484 else => return error.InvalidArtifact,
485 };
486 return state.bytes;
487 }
488
489 const DynamicSharedMemoryState = struct {
490 bytes: u32 = 0,
491
492 fn visit(self: *@This(), op: *ir.Operation) gpu.BackendError!ir.Operation.WalkResult {
493 if (!std.mem.eql(u8, op.name.name, MemrefDialect.AllocOp.operation_name)) return .advance;
494 const alloc = MemrefDialect.AllocOp{ .op = op };
495 if (alloc.getDynamicSize() == null) return .advance;
496 const result = alloc.getResult();
497 const params = dynamicSharedMemrefParams(result.type) orelse return error.InvalidArtifact;
498 if (params.addr_space != .shared) return .advance;
499 const size = params.size orelse return error.InvalidArtifact;
500 const element_bytes = try dynamicSharedElementByteSize(params.element_type_name);
501 const bytes = std.math.mul(u64, size, element_bytes) catch return error.InvalidArtifact;
502 const end = std.math.add(u64, try dynamicSharedByteOffset(op), bytes) catch return error.InvalidArtifact;
503 self.bytes = @max(self.bytes, std.math.cast(u32, end) orelse return error.InvalidArtifact);
504 return .advance;
505 }
506 };
507
508 fn dynamicSharedMemrefParams(typ: ir.Type) ?MemrefDialect.MemrefParams {
509 const name = typ.getDialectTypeName() orelse return null;
510 if (!std.mem.eql(u8, name, MemrefDialect.name)) return null;
511 return MemrefDialect.parseMemrefParams(typ.getDialectParamKey() orelse return null);
512 }
513
514 fn dynamicSharedByteOffset(op: *ir.Operation) gpu.BackendError!u32 {
515 const attr_name = choir.dialects.gpu.attr_names.dynamic_shared_byte_offset;
516 const attr = op.getAttrAs(ir.Attribute.IntegerAttr, attr_name) orelse return 0;
517 const value = attr.getValue();
518 if (value < 0) return error.InvalidArtifact;
519 return std.math.cast(u32, value) orelse return error.InvalidArtifact;
520 }
521
522 fn dynamicSharedElementByteSize(type_name: []const u8) gpu.BackendError!u32 {
523 const kind = dialects.arith.scalarKindFromTypeName(type_name) orelse return error.InvalidArtifact;
524 return switch (kind) {
525 .bool => 1,
526 .i8, .u8 => 1,
527 .i16, .u16 => 2,
528 .i64, .u64 => 8,
529 .f64 => 8,
530 .f16 => 2,
531 .bf16 => 2,
532 .index, .i32, .u32, .f32 => 4,
533 };
534 }
535
536 test "kernelization model counts generated dynamic shared memory bytes" {
537 const allocator = std.testing.allocator;
538
539 var builder = try kernel_program.Builder.init(
540 allocator,
541 kernel_program.Builder.Limits.testing,
542 "accy_dynamic_shared_product",
543 &.{},
544 );
545 errdefer builder.deinit();
546
547 _ = try builder.dynamicSharedBuffer(.f32, 4, 0);
548 _ = try builder.dynamicSharedBuffer(.f32, 4, 64);
549 _ = try builder.sharedBuffer(.f32, 256);
550 try builder.return_();
551
552 var program = try builder.finish();
553 defer program.deinit();
554
555 const entry_name = try allocator.dupe(u8, "accy_dynamic_shared_product");
556 defer allocator.free(entry_name);
557
558 try std.testing.expectEqual(@as(u32, 80), try dynamicSharedMemoryBytes(program.kernelModule()));
559 try std.testing.expectEqual(@as(u32, 80), (LoweredKernel{
560 .work_item_id = 0,
561 .entry_name = entry_name,
562 .program = program,
563 .argument_count = 0,
564 .body_fingerprint = 0,
565 .dynamic_shared_memory_bytes = 80,
566 .schedule = .{
567 .kind = .flat,
568 .threads = .{ .x = 64 },
569 },
570 .launch = .{
571 .grid = .{ 1, 1, 1 },
572 .block = .{ 64, 1, 1 },
573 },
574 }).summary().dynamic_shared_memory_bytes);
575 }
576
577 test "kernelization analysis storage bounds its real shared Context and containers" {
578 for ([_]usize{ 0, 1, 7, 8, 63, 64, 257 }) |count| {
579 try checkAnalysisStorage(ir.Context.Limits.testing, count);
580 }
581 try checkAnalysisStorage(ir.Context.Limits.standard, 1);
582 }
583
584 fn checkAnalysisStorage(limits: ir.Context.Limits, count: usize) !void {
585 const bound = try KernelizationAnalysis.baseStorageBound(limits, count);
586 const bytes = try std.testing.allocator.alloc(u8, @intCast(bound));
587 defer std.testing.allocator.free(bytes);
588 var buffer = std.heap.FixedBufferAllocator.init(bytes);
589 const base = buffer.allocator();
590 const vtable: std.mem.Allocator.VTable = .{
591 .alloc = base.vtable.alloc,
592 .resize = std.mem.Allocator.noResize,
593 .remap = std.mem.Allocator.noRemap,
594 .free = std.mem.Allocator.noFree,
595 };
596 const allocator: std.mem.Allocator = .{ .ptr = base.ptr, .vtable = &vtable };
597 const analysis = try allocator.create(KernelizationAnalysis);
598 defer allocator.destroy(analysis);
599 analysis.* = try KernelizationAnalysis.init(allocator, limits);
600 defer analysis.deinit();
601 try analysis.reserveKernelCapacity(count);
602 try std.testing.expect(analysis.kernels.capacity >= count);
603 try std.testing.expect(analysis.work_to_kernel.capacity() >= count);
604 const reserved = buffer.end_index;
605 for (0..count) |index| try analysis.work_to_kernel.put(index, index);
606 try analysis.reserveKernelCapacity(count);
607 try std.testing.expectEqual(reserved, buffer.end_index);
608 try std.testing.expectEqual(count, analysis.work_to_kernel.count());
609 try std.testing.expect(buffer.end_index <= bound);
610 try std.testing.expectEqualDeep(limits, analysis.context.capacity.asLimits());
611 try std.testing.expect(analysis.context.exhaustedSegment() == null);
612 }
613
614 test "kernelization analysis storage rejects unrepresentable bounds" {
615 const standard = ir.Context.Limits.standard;
616 try std.testing.expectError(
617 error.WorkOverflow,
618 KernelizationAnalysis.baseStorageBound(standard, std.math.maxInt(u64)),
619 );
620 try std.testing.expectError(
621 error.WorkOverflow,
622 KernelizationAnalysis.baseStorageBound(standard, std.math.maxInt(u32)),
623 );
624 var limits = standard;
625 limits.operations.storage_bytes = std.math.maxInt(usize);
626 try std.testing.expectError(
627 error.WorkOverflow,
628 KernelizationAnalysis.baseStorageBound(limits, 0),
629 );
630 }
631
632 test "kernelization analysis storage cleans up failed container reservations" {
633 try std.testing.checkAllAllocationFailures(std.testing.allocator, testAnalysisReservation, .{});
634 }
635
636 fn testAnalysisReservation(allocator: std.mem.Allocator) !void {
637 var analysis = try KernelizationAnalysis.init(allocator, ir.Context.Limits.testing);
638 defer analysis.deinit();
639 try analysis.reserveKernelCapacity(257);
640 try std.testing.expect(analysis.kernels.capacity >= 257);
641 try std.testing.expect(analysis.work_to_kernel.capacity() >= 257);
642 }
643
644 test "kernelization model owns its Context through constructor failure" {
645 try std.testing.checkAllAllocationFailures(std.testing.allocator, testAnalysisContext, .{});
646 var limits = ir.Context.Limits.standard;
647 limits.configuration.table_bytes = 0;
648 try std.testing.expectError(
649 error.WorkExhausted,
650 KernelizationAnalysis.init(std.testing.allocator, limits),
651 );
652 limits = ir.Context.Limits.standard;
653 limits.operations.storage_bytes = std.math.maxInt(usize);
654 try std.testing.expectError(
655 error.CapacityOverflow,
656 KernelizationAnalysis.init(std.testing.allocator, limits),
657 );
658 }
659
660 fn testAnalysisContext(allocator: std.mem.Allocator) !void {
661 var analysis = try KernelizationAnalysis.init(allocator, ir.Context.Limits.standard);
662 defer analysis.deinit();
663 try std.testing.expectEqual(@as(usize, 0), analysis.kernelCount());
664 try std.testing.expectEqualDeep(
665 ir.Context.Limits.standard,
666 analysis.context.capacity.asLimits(),
667 );
668 try std.testing.expectEqual(
669 @as(?ir.Context.Segment, null),
670 analysis.context.exhaustedSegment(),
671 );
672 }