lib/accy/src/preparation/kernelization/lowering/elementwise.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const alloc_fixed = @import("alloc_fixed");
5 const generated_abi = @import("abi.zig");
6 const generated_builder = @import("builder.zig");
7 const common = @import("common.zig");
8 const generated_guard = @import("guard.zig");
9 const generated_name = @import("name.zig");
10 const range_analysis = @import("range.zig");
11 const generated_schedule = @import("schedule.zig");
12
13 const ir = common.ir;
14 const dialect_mod = common.dialect_mod;
15 const kernel_root = common.kernel_root;
16 const bufferization = common.bufferization;
17 const kernelization_model = @import("../model/root.zig");
18 const schedule_planning = common.schedule_planning;
19 const shape_analysis = common.shape_analysis;
20 const i64_attr_list_stack_capacity = common.i64_attr_list_stack_capacity;
21 const externalInputIndex = common.externalInputIndex;
22 const readI64ListAttrBounded = common.readI64ListAttrBounded;
23 const mapKernelBuildError = common.mapKernelBuildError;
24 const isName = common.isName;
25
26 const ElementwiseKernel = kernelization_model.ElementwiseKernel;
27 const bufferSlotById = common.bufferSlotById;
28 const LoweredKernel = kernelization_model.LoweredKernel;
29
30 /// The kernel stage calls this once per try at emitting a kernel under one candidate schedule to
31 /// charge element-by-element scratch, so compilation can refuse a pass that would exceed the
32 /// caller's limits. That scratch belongs to the costs a pass declares before it runs, its work
33 /// bound. The bound covers the scratch kept while choosing a kernel body and while emitting either
34 /// body, including the vector fallback, the second body this file emits when the per-element form
35 /// does not apply. The bound is computed from the operation, value and operand counts of the
36 /// source. Kernel names, parameter lists, and the storage of the generated builder and the object
37 /// that owns the operations and values of the generated kernels are charged separately.
38 pub fn scratchStorageBound(source: @import("choir").passes.pass.work.Census) !u64 {
39 const accounting = @import("choir").passes.pass.work;
40 const per_op = @sizeOf(*ir.Operation) + @sizeOf(ElementwiseKernel) +
41 @sizeOf(choir_abi.DType) + 3 * @sizeOf(kernel_root.Value);
42 const per_value = @sizeOf(choir_abi.DType) + @sizeOf(kernel_root.Value);
43 var bytes = try accounting.add(
44 try accounting.multiply(source.operations, per_op),
45 try accounting.multiply(source.values, per_value),
46 );
47 bytes = try accounting.add(bytes, 8 * @alignOf(kernel_root.Value));
48 const walk = try accounting.add(
49 try accounting.hashMapGrowth(*ir.Operation, void, source.operations),
50 try accounting.arrayListGrowth(*ir.Value, source.operands),
51 );
52 bytes = try accounting.add(bytes, try accounting.multiply(2, walk));
53 bytes = try accounting.add(bytes, try accounting.arrayListGrowth(VectorInput, source.values));
54 const memo = try accounting.hashMapGrowth(*ir.Value, kernel_root.Value, source.values);
55 bytes = try accounting.add(bytes, try accounting.multiply(6, memo));
56 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
57 return bytes;
58 }
59
60 pub fn lower(
61 allocator: std.mem.Allocator,
62 ir_ctx: *ir.Context,
63 outline: kernelization_model.KernelOutline,
64 work: schedule_planning.ScheduleWorkItem,
65 buffer_plan: *const bufferization.BufferPlanAnalysis,
66 shape_plan: *const shape_analysis.ShapeLayoutAnalysis,
67 format: ?gpu.ArtifactFormat,
68 ) common.LoweringError!LoweredKernel {
69 if (work.kind != .elementwise_single and work.kind != .elementwise_fusion) {
70 return error.UnsupportedOperation;
71 }
72 if (work.ops.len == 0) return error.UnsupportedOperation;
73 if (outline.inputCount() > std.math.maxInt(u32)) return error.InvalidArtifact;
74 const emit_ops = allocator.alloc(*ir.Operation, work.ops.len) catch return error.OutOfMemory;
75 defer allocator.free(emit_ops);
76 var emit_count: usize = 0;
77 for (work.ops) |op| {
78 if (isName(op.name.name, dialect_mod.AccyDialect.SliceOp.operation_name)) continue;
79 if (isName(op.name.name, dialect_mod.AccyDialect.PadOp.operation_name)) continue;
80 if (isName(op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name)) continue;
81 emit_ops[emit_count] = op;
82 emit_count += 1;
83 }
84 if (emit_count == 0) return error.UnsupportedOperation;
85 const emitted_ops = emit_ops[0..emit_count];
86 if (emitted_ops[emit_count - 1] != work.ops[work.ops.len - 1]) return error.UnsupportedOperation;
87 const kinds = allocator.alloc(ElementwiseKernel, emit_count) catch return error.OutOfMemory;
88 defer allocator.free(kinds);
89 for (emitted_ops, 0..) |op, index| {
90 const kind = kernelForOperation(op) orelse return error.UnsupportedOperation;
91 kinds[index] = kind;
92 if (op.getNumOperands() != expectedInputCount(kind)) return error.UnsupportedOperation;
93 }
94 const result_dtypes = try collectResultDTypes(allocator, emitted_ops, kinds, shape_plan);
95 defer allocator.free(result_dtypes);
96
97 const output_slot = bufferSlotById(buffer_plan, outline.output_slot_id) orelse return error.InvalidArtifact;
98 if (output_slot.dtype != work.dtype) return error.InvalidArtifact;
99
100 const input_dtypes = try collectInputDTypes(allocator, outline, buffer_plan);
101 defer allocator.free(input_dtypes);
102
103 var abi = try generated_abi.flatTyped(allocator, output_slot.dtype, input_dtypes);
104 defer abi.deinit(allocator);
105
106 const contains_gather = try closureContainsGather(allocator, emitted_ops);
107 const rank2 = if (contains_gather) try rank2ElementwiseForResult(work.root, work.element_count) else null;
108
109 if (vectorLoweringEligible(format, work, output_slot, contains_gather)) {
110 if (try lowerVector(allocator, ir_ctx, outline, work, buffer_plan, abi, emitted_ops, kinds, result_dtypes)) |kernel| {
111 return kernel;
112 }
113 }
114
115 const entry_name = try generated_name.elementwise(allocator, kinds, work.id);
116 errdefer allocator.free(entry_name);
117
118 if (rank2) |rank2_value| {
119 var lowered = try generated_builder.withLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrixThreads(rank2_value.cols_threads, rank2_value.rows_threads), .{
120 .abi = abi,
121 .allocator = allocator,
122 .work = work,
123 .ops = emitted_ops,
124 .kinds = kinds,
125 .result_dtypes = result_dtypes,
126 .outline = outline,
127 .buffer_plan = buffer_plan,
128 .rank2 = rank2_value,
129 }, emitRank2ElementwiseBody);
130 lowered.body = .{ .elementwise_rank2 = try rank2ElementwisePlan(rank2_value) };
131 return lowered;
132 }
133
134 return generated_builder.withLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
135 .abi = abi,
136 .allocator = allocator,
137 .work = work,
138 .ops = emitted_ops,
139 .kinds = kinds,
140 .result_dtypes = result_dtypes,
141 .outline = outline,
142 .buffer_plan = buffer_plan,
143 }, emitElementwiseBody);
144 }
145
146 fn vectorLoweringEligible(
147 format: ?gpu.ArtifactFormat,
148 work: schedule_planning.ScheduleWorkItem,
149 output_slot: *const bufferization.BufferSlot,
150 contains_gather: bool,
151 ) bool {
152 if (format != .cuda_ptx) return false;
153 if (contains_gather) return false;
154 if (work.dtype != .f32 or output_slot.dtype != .f32) return false;
155 if (work.element_count < 4 or work.element_count % 4 != 0) return false;
156 return true;
157 }
158
159 const VectorInput = struct {
160 value: *ir.Value,
161 input_index: usize,
162 };
163
164 fn lowerVector(
165 allocator: std.mem.Allocator,
166 ir_ctx: *ir.Context,
167 outline: kernelization_model.KernelOutline,
168 work: schedule_planning.ScheduleWorkItem,
169 buffer_plan: *const bufferization.BufferPlanAnalysis,
170 abi: generated_abi.Flat,
171 ops: []const *ir.Operation,
172 kinds: []const ElementwiseKernel,
173 result_dtypes: []const choir_abi.DType,
174 ) common.LoweringError!?LoweredKernel {
175 const vector_inputs = try collectVectorInputs(allocator, ops, outline, buffer_plan);
176 defer allocator.free(vector_inputs);
177
178 const entry_name = try generated_name.elementwiseVector(allocator, kinds, work.id);
179 var lowered = generated_builder.withLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
180 .abi = abi,
181 .allocator = allocator,
182 .work = work,
183 .ops = ops,
184 .kinds = kinds,
185 .result_dtypes = result_dtypes,
186 .outline = outline,
187 .buffer_plan = buffer_plan,
188 .vector_inputs = vector_inputs,
189 }, emitVectorElementwiseBody) catch |err| switch (err) {
190 error.UnsupportedOperation, error.CapabilityMismatch => {
191 allocator.free(entry_name);
192 return null;
193 },
194 else => {
195 allocator.free(entry_name);
196 return err;
197 },
198 };
199 lowered.body = .{ .elementwise_vector = .{ .quads = work.element_count / 4 } };
200 return lowered;
201 }
202
203 const OperandWalk = struct {
204 allocator: std.mem.Allocator,
205 pending: std.ArrayListUnmanaged(*ir.Value) = .empty,
206 expanded: std.AutoHashMap(*ir.Operation, void),
207
208 fn init(allocator: std.mem.Allocator) OperandWalk {
209 return .{
210 .allocator = allocator,
211 .expanded = std.AutoHashMap(*ir.Operation, void).init(allocator),
212 };
213 }
214
215 fn deinit(self: *OperandWalk) void {
216 self.pending.deinit(self.allocator);
217 self.expanded.deinit();
218 }
219
220 fn expand(self: *OperandWalk, op: *ir.Operation) error{OutOfMemory}!void {
221 const entry = try self.expanded.getOrPut(op);
222 if (entry.found_existing) return;
223 const operands = op.getOperandValues();
224 var remaining = operands.len;
225 while (remaining != 0) {
226 remaining -= 1;
227 try self.pending.append(self.allocator, operands[remaining]);
228 }
229 }
230
231 fn next(self: *OperandWalk) ?*ir.Value {
232 return self.pending.pop();
233 }
234 };
235
236 fn collectVectorInputs(
237 allocator: std.mem.Allocator,
238 ops: []const *ir.Operation,
239 outline: kernelization_model.KernelOutline,
240 buffer_plan: *const bufferization.BufferPlanAnalysis,
241 ) common.LoweringError![]VectorInput {
242 var collected: std.ArrayListUnmanaged(VectorInput) = .empty;
243 errdefer collected.deinit(allocator);
244 var walk = OperandWalk.init(allocator);
245 defer walk.deinit();
246 var remaining = ops.len;
247 while (remaining != 0) {
248 remaining -= 1;
249 try walk.expand(ops[remaining]);
250 }
251 while (walk.next()) |value| {
252 try collectVectorInputValue(&walk, &collected, value, ops, outline, buffer_plan);
253 }
254 return collected.toOwnedSlice(allocator);
255 }
256
257 fn collectVectorInputValue(
258 walk: *OperandWalk,
259 collected: *std.ArrayListUnmanaged(VectorInput),
260 value: *ir.Value,
261 ops: []const *ir.Operation,
262 outline: kernelization_model.KernelOutline,
263 buffer_plan: *const bufferization.BufferPlanAnalysis,
264 ) common.LoweringError!void {
265 for (collected.items) |existing| {
266 if (existing.value == value) return;
267 }
268 if (value.getDefiningOp()) |def_any| {
269 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
270 for (ops) |op| {
271 if (op == def_op) return;
272 }
273 if (isName(def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) return;
274 if (isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) {
275 return;
276 }
277 if (buffer_plan.getSlot(value) == null) {
278 if (kernelForOperation(def_op) != null) try walk.expand(def_op);
279 return;
280 }
281 }
282 const slot = buffer_plan.getSlot(value) orelse return;
283 if (slot.dtype != .f32) return;
284 const input_index = externalInputIndex(outline, slot.id) orelse return;
285 try collected.append(walk.allocator, .{ .value = value, .input_index = input_index });
286 }
287
288 fn emit_vector_elementwise_quad(
289 guarded: anytype,
290 base: kernel_root.Value,
291 body_ctx: anytype,
292 ) !void {
293 const loaded_vectors = try body_ctx.allocator.alloc(
294 kernel_root.Value,
295 body_ctx.vector_inputs.len,
296 );
297 defer body_ctx.allocator.free(loaded_vectors);
298 for (body_ctx.vector_inputs, 0..) |vector_input, input_slot| {
299 loaded_vectors[input_slot] = try guarded.loadVector(
300 body_ctx.abi.input(guarded, vector_input.input_index),
301 base,
302 4,
303 );
304 }
305
306 const op_results = try body_ctx.allocator.alloc(kernel_root.Value, body_ctx.ops.len);
307 defer body_ctx.allocator.free(op_results);
308
309 var lane_results: [4]kernel_root.Value = undefined;
310 for (0..4) |lane| {
311 var memo_storage: ?ValueMemo = ValueMemo.init(body_ctx.allocator);
312 defer if (memo_storage) |*memo| memo.deinit();
313 if (memo_storage) |*memo| {
314 for (body_ctx.vector_inputs, 0..) |vector_input, input_slot| {
315 const lane_value = try guarded.extractLane(
316 loaded_vectors[input_slot],
317 @intCast(lane),
318 .f32,
319 );
320 memo.put(vector_input.value, lane_value) catch return error.OutOfMemory;
321 }
322 }
323 const flat_index = try guarded.add(
324 base,
325 try guarded.constantIndex(@as(i64, @intCast(lane))),
326 );
327 lane_results[lane] = try emitChainAtPosition(
328 guarded,
329 &memo_storage,
330 flat_index,
331 .{},
332 op_results,
333 body_ctx,
334 );
335 }
336
337 const packed_result = try guarded.packVector(lane_results);
338 try guarded.storeIndex(packed_result, body_ctx.out, base);
339 }
340
341 fn emitVectorElementwiseBody(logical: anytype, ctx: anytype) !void {
342 const logical_index = try logical.index1D("i", ctx.work.element_count / 4);
343 try generated_guard.countQuadDo(logical, logical_index, ctx.abi.count(logical), .{
344 .allocator = ctx.allocator,
345 .work = ctx.work,
346 .ops = ctx.ops,
347 .kinds = ctx.kinds,
348 .result_dtypes = ctx.result_dtypes,
349 .outline = ctx.outline,
350 .buffer_plan = ctx.buffer_plan,
351 .abi = ctx.abi,
352 .out = ctx.abi.output(logical),
353 .vector_inputs = ctx.vector_inputs,
354 }, emit_vector_elementwise_quad);
355 }
356
357 fn emitChainAtPosition(
358 builder: anytype,
359 memo: anytype,
360 flat_index: kernel_root.Value,
361 facts: IndexFacts,
362 op_results: []kernel_root.Value,
363 body_ctx: anytype,
364 ) !kernel_root.Value {
365 for (body_ctx.ops, 0..) |op, op_index| {
366 const kind = body_ctx.kinds[op_index];
367 var inputs: [3]kernel_root.Value = undefined;
368 const operands = op.getOperandValues();
369 for (operands, 0..) |operand, operand_index| {
370 inputs[operand_index] = try elementwiseOperandValue(
371 builder,
372 memo,
373 flat_index,
374 facts,
375 operand,
376 body_ctx.ops,
377 op_index,
378 op_results,
379 body_ctx.outline,
380 body_ctx.buffer_plan,
381 body_ctx.abi,
382 );
383 }
384 op_results[op_index] = try emitElementwiseValue(builder, kind, inputs[0..operands.len], body_ctx.result_dtypes[op_index], op);
385 }
386 return elementwiseOutputValue(body_ctx.work, body_ctx.ops, op_results) orelse error.InvalidArtifact;
387 }
388
389 fn closureContainsGather(
390 allocator: std.mem.Allocator,
391 ops: []const *ir.Operation,
392 ) error{OutOfMemory}!bool {
393 var walk = OperandWalk.init(allocator);
394 defer walk.deinit();
395 for (ops) |op| {
396 if (isName(op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name)) return true;
397 try walk.expand(op);
398 }
399 while (walk.next()) |value| {
400 const def_any = value.getDefiningOp() orelse continue;
401 const op: *ir.Operation = @ptrCast(@alignCast(def_any));
402 if (isName(op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name)) return true;
403 try walk.expand(op);
404 }
405 return false;
406 }
407
408 const rank2_tile_width = 32;
409 const rank2_tile_height = 8;
410
411 const Rank2Elementwise = struct {
412 rows: i64,
413 cols: i64,
414 total: i64,
415 rows_threads: u32 = rank2_tile_height,
416 cols_threads: u32 = rank2_tile_width,
417 };
418
419 const Rank2IndexFacts = struct {
420 rows: i64,
421 cols: i64,
422 row: kernel_root.Value,
423 col: kernel_root.Value,
424
425 fn coordinate(self: Rank2IndexFacts, shape: []const i64, dim_index: usize) ?kernel_root.Value {
426 if (shape.len != 2) return null;
427 if (shape[0] != self.rows or shape[1] != self.cols) return null;
428 return switch (dim_index) {
429 0 => self.row,
430 1 => self.col,
431 else => null,
432 };
433 }
434 };
435
436 const IndexFacts = struct {
437 rank2: ?Rank2IndexFacts = null,
438 };
439
440 fn rank2ElementwiseForResult(
441 root: *ir.Operation,
442 element_count: u64,
443 ) common.LoweringError!?Rank2Elementwise {
444 const result = root.getResult(0) orelse return error.InvalidArtifact;
445 var dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
446 const dims = try decodedDims(result, dims_buffer[0..]);
447 return rank2ElementwiseForDims(dims, element_count);
448 }
449
450 fn rank2ElementwiseForDims(dims: []const i64, element_count: u64) ?Rank2Elementwise {
451 if (dims.len != 2) return null;
452 const rows = dims[0];
453 const cols = dims[1];
454 if (rows <= 0 or cols <= 0) return null;
455 if (@rem(rows, rank2_tile_height) != 0 or @rem(cols, rank2_tile_width) != 0) return null;
456 const total = std.math.mul(i64, rows, cols) catch return null;
457 const count = std.math.cast(i64, element_count) orelse return null;
458 if (total != count) return null;
459 if (total > std.math.maxInt(i32)) return null;
460 return .{
461 .rows = rows,
462 .cols = cols,
463 .total = total,
464 };
465 }
466
467 fn rank2ElementwisePlan(
468 rank2: Rank2Elementwise,
469 ) common.LoweringError!kernelization_model.ElementwiseRank2Plan {
470 return .{
471 .rows = std.math.cast(u32, rank2.rows) orelse return error.CapabilityMismatch,
472 .cols = std.math.cast(u32, rank2.cols) orelse return error.CapabilityMismatch,
473 .threads_x = rank2.cols_threads,
474 .threads_y = rank2.rows_threads,
475 };
476 }
477
478 fn collectResultDTypes(
479 allocator: std.mem.Allocator,
480 ops: []const *ir.Operation,
481 kinds: []const ElementwiseKernel,
482 shape_plan: *const shape_analysis.ShapeLayoutAnalysis,
483 ) common.LoweringError![]choir_abi.DType {
484 const result_dtypes = allocator.alloc(choir_abi.DType, ops.len) catch return error.OutOfMemory;
485 errdefer allocator.free(result_dtypes);
486
487 for (ops, kinds, 0..) |op, kind, index| {
488 const result = op.getResult(0) orelse return error.InvalidArtifact;
489 const result_dtype = try valueDType(shape_plan, result);
490 try validateOperationDTypes(op, kind, result_dtype, shape_plan);
491 result_dtypes[index] = result_dtype;
492 }
493
494 return result_dtypes;
495 }
496
497 fn collectInputDTypes(
498 allocator: std.mem.Allocator,
499 outline: kernelization_model.KernelOutline,
500 buffer_plan: *const bufferization.BufferPlanAnalysis,
501 ) common.LoweringError![]choir_abi.DType {
502 const input_dtypes = allocator.alloc(choir_abi.DType, outline.inputCount()) catch return error.OutOfMemory;
503 errdefer allocator.free(input_dtypes);
504
505 for (outline.input_slot_ids, 0..) |slot_id, index| {
506 const input_slot = bufferSlotById(buffer_plan, slot_id) orelse return error.InvalidArtifact;
507 if (!executableBufferDType(input_slot.dtype)) return error.CapabilityMismatch;
508 input_dtypes[index] = input_slot.dtype;
509 }
510
511 return input_dtypes;
512 }
513
514 fn validateOperationDTypes(
515 op: *ir.Operation,
516 kind: ElementwiseKernel,
517 result_dtype: choir_abi.DType,
518 shape_plan: *const shape_analysis.ShapeLayoutAnalysis,
519 ) common.LoweringError!void {
520 const operands = op.getOperandValues();
521 switch (kind) {
522 .convert => {
523 const input_dtype = try valueDType(shape_plan, operands[0]);
524 if (!executableDataDType(input_dtype)) return error.CapabilityMismatch;
525 if (!executableDataDType(result_dtype)) return error.CapabilityMismatch;
526 },
527 .compare => {
528 if (result_dtype != .i1) return error.InvalidArtifact;
529 const lhs_dtype = try valueDType(shape_plan, operands[0]);
530 const rhs_dtype = try valueDType(shape_plan, operands[1]);
531 if (lhs_dtype != rhs_dtype) return error.InvalidArtifact;
532 if (!executableDataDType(lhs_dtype)) return error.CapabilityMismatch;
533 },
534 .select => {
535 if ((try valueDType(shape_plan, operands[0])) != .i1) return error.InvalidArtifact;
536 const true_dtype = try valueDType(shape_plan, operands[1]);
537 const false_dtype = try valueDType(shape_plan, operands[2]);
538 if (true_dtype != false_dtype or result_dtype != true_dtype) return error.InvalidArtifact;
539 if (!executableBufferDType(result_dtype)) return error.CapabilityMismatch;
540 },
541 .pow,
542 .atan2,
543 .sqrt,
544 .exp,
545 .log,
546 .tanh,
547 .sin,
548 .cos,
549 .tan,
550 .floor,
551 .round,
552 .trunc,
553 => {
554 try validateSameDataDType(operands, result_dtype, shape_plan);
555 if (result_dtype != .f32 and result_dtype != .f16) return error.CapabilityMismatch;
556 },
557 .add,
558 .sub,
559 .mul,
560 .div,
561 .min,
562 .max,
563 .neg,
564 .abs,
565 => {
566 try validateSameDataDType(operands, result_dtype, shape_plan);
567 if (!executableDataDType(result_dtype)) return error.CapabilityMismatch;
568 },
569 }
570 }
571
572 fn validateSameDataDType(
573 operands: []const *ir.Value,
574 result_dtype: choir_abi.DType,
575 shape_plan: *const shape_analysis.ShapeLayoutAnalysis,
576 ) common.LoweringError!void {
577 for (operands) |operand| {
578 if ((try valueDType(shape_plan, operand)) != result_dtype) return error.InvalidArtifact;
579 }
580 }
581
582 fn valueDType(
583 shape_plan: *const shape_analysis.ShapeLayoutAnalysis,
584 value: *ir.Value,
585 ) common.LoweringError!choir_abi.DType {
586 return (shape_plan.get(value) orelse return error.InvalidArtifact).dtype;
587 }
588
589 fn executableDataDType(dtype: choir_abi.DType) bool {
590 return dtype == .f32 or dtype == .f16 or dtype == .bf16 or dtype == .f64 or dtype == .i8 or dtype == .i16 or dtype == .i32 or dtype == .u8 or dtype == .u16 or dtype == .u32 or dtype == .i64 or dtype == .u64;
591 }
592
593 fn executableBufferDType(dtype: choir_abi.DType) bool {
594 return executableDataDType(dtype) or dtype == .i1;
595 }
596
597 fn emit_elementwise_index(
598 guarded: anytype,
599 index: kernel_root.Index1D,
600 body_ctx: anytype,
601 ) !void {
602 const flat_index = index.index;
603 const op_results = try body_ctx.allocator.alloc(kernel_root.Value, body_ctx.ops.len);
604 defer body_ctx.allocator.free(op_results);
605 var memo_storage: ?ValueMemo = ValueMemo.init(body_ctx.allocator);
606 defer if (memo_storage) |*memo| memo.deinit();
607 const result = try emitChainAtPosition(
608 guarded,
609 &memo_storage,
610 flat_index,
611 .{},
612 op_results,
613 body_ctx,
614 );
615 try guarded.storeIndex(result, body_ctx.out, flat_index);
616 }
617
618 fn emit_rank2_elementwise_guarded(guarded_inner: anytype, guarded_ctx: anytype) !void {
619 const op_results = try guarded_ctx.payload.allocator.alloc(
620 kernel_root.Value,
621 guarded_ctx.payload.ops.len,
622 );
623 defer guarded_ctx.payload.allocator.free(op_results);
624 var memo_storage: ?ValueMemo = ValueMemo.init(guarded_ctx.payload.allocator);
625 defer if (memo_storage) |*memo| memo.deinit();
626 const facts = IndexFacts{
627 .rank2 = .{
628 .rows = guarded_ctx.payload.rank2.rows,
629 .cols = guarded_ctx.payload.rank2.cols,
630 .row = guarded_ctx.row,
631 .col = guarded_ctx.col,
632 },
633 };
634 const result = try emitChainAtPosition(
635 guarded_inner,
636 &memo_storage,
637 guarded_ctx.flat_index,
638 facts,
639 op_results,
640 guarded_ctx.payload,
641 );
642 try guarded_inner.storeIndex(
643 result,
644 guarded_ctx.payload.out,
645 guarded_ctx.flat_index,
646 );
647 }
648
649 fn emit_rank2_elementwise_index(
650 inner: anytype,
651 active: kernel_root.Index2D,
652 body_ctx: anytype,
653 ) !void {
654 const flat_index = try active.linear(inner);
655 const zero = try inner.constantIndex(0);
656 const count = try inner.loadIndex(body_ctx.count_buffer, zero);
657 const index_i32 = try inner.cast(flat_index, .i32);
658 const in_bounds = try inner.compare(.lt, index_i32, count);
659 try inner.guardDo(in_bounds, .{
660 .flat_index = flat_index,
661 .row = active.y.index,
662 .col = active.x.index,
663 .payload = body_ctx,
664 }, emit_rank2_elementwise_guarded);
665 }
666
667 fn emitElementwiseBody(logical: anytype, ctx: anytype) !void {
668 const logical_index = try logical.index1D("i", ctx.work.element_count);
669 try generated_guard.countIndexDo(logical, logical_index, ctx.abi.count(logical), .{
670 .allocator = ctx.allocator,
671 .work = ctx.work,
672 .ops = ctx.ops,
673 .kinds = ctx.kinds,
674 .result_dtypes = ctx.result_dtypes,
675 .outline = ctx.outline,
676 .buffer_plan = ctx.buffer_plan,
677 .abi = ctx.abi,
678 .out = ctx.abi.output(logical),
679 }, emit_elementwise_index);
680 }
681
682 fn emitRank2ElementwiseBody(logical: anytype, ctx: anytype) !void {
683 const rows: u64 = @intCast(ctx.rank2.rows);
684 const cols: u64 = @intCast(ctx.rank2.cols);
685 const index = try logical.index2D(.{
686 .x = .{ .name = "col", .extent = cols },
687 .y = .{ .name = "row", .extent = rows },
688 });
689 try logical.guardIndex2DDo(index, .{
690 .allocator = ctx.allocator,
691 .work = ctx.work,
692 .ops = ctx.ops,
693 .kinds = ctx.kinds,
694 .result_dtypes = ctx.result_dtypes,
695 .outline = ctx.outline,
696 .buffer_plan = ctx.buffer_plan,
697 .abi = ctx.abi,
698 .out = ctx.abi.output(logical),
699 .rank2 = ctx.rank2,
700 .count_buffer = ctx.abi.count(logical),
701 }, emit_rank2_elementwise_index);
702 }
703
704 fn elementwiseOperandValue(
705 builder: anytype,
706 memo: anytype,
707 index: kernel_root.Value,
708 facts: IndexFacts,
709 operand: *ir.Value,
710 ops: []const *ir.Operation,
711 op_index: usize,
712 op_results: []const kernel_root.Value,
713 outline: kernelization_model.KernelOutline,
714 buffer_plan: *const bufferization.BufferPlanAnalysis,
715 abi: generated_abi.Flat,
716 ) common.LoweringError!kernel_root.Value {
717 if (elementwiseChainValue(operand, ops, op_index, op_results)) |value| return value;
718 if (try broadcastInDimElementwiseOperandValue(builder, index, operand, outline, buffer_plan, abi)) |value| return value;
719 return evaluateValueAtPosition(builder, memo, index, facts, operand, outline, buffer_plan, abi, max_recompute_depth);
720 }
721
722 pub const max_shape_read_depth = 8;
723 pub const max_recompute_depth = 32;
724
725 pub const ValueMemo = std.AutoHashMap(*ir.Value, kernel_root.Value);
726
727 const ValueReadQuery = union(enum) {
728 value: *ir.Value,
729 slice: *ir.Operation,
730 gather: *ir.Operation,
731 pad: *ir.Operation,
732 };
733
734 pub fn shapeReadValue(
735 builder: anytype,
736 index: kernel_root.Value,
737 value: *ir.Value,
738 outline: kernelization_model.KernelOutline,
739 buffer_plan: *const bufferization.BufferPlanAnalysis,
740 abi: anytype,
741 depth: usize,
742 ) common.LoweringError!kernel_root.Value {
743 var memo_storage: ?ValueMemo = null;
744 defer if (memo_storage) |*memo| memo.deinit();
745 return evaluateValueAtPosition(builder, &memo_storage, index, .{}, value, outline, buffer_plan, abi, depth);
746 }
747
748 pub fn evaluateValueAtIndex(
749 builder: anytype,
750 memo: anytype,
751 index: kernel_root.Value,
752 value: *ir.Value,
753 outline: kernelization_model.KernelOutline,
754 buffer_plan: *const bufferization.BufferPlanAnalysis,
755 abi: anytype,
756 depth: usize,
757 ) common.LoweringError!kernel_root.Value {
758 return evaluateValueAtPosition(builder, memo, index, .{}, value, outline, buffer_plan, abi, depth);
759 }
760
761 fn evaluateValueAtPosition(
762 builder: anytype,
763 memo: anytype,
764 index: kernel_root.Value,
765 facts: IndexFacts,
766 value: *ir.Value,
767 outline: kernelization_model.KernelOutline,
768 buffer_plan: *const bufferization.BufferPlanAnalysis,
769 abi: anytype,
770 depth: usize,
771 ) common.LoweringError!kernel_root.Value {
772 return readValueAtPosition(builder, memo, index, facts, .{ .value = value }, outline, buffer_plan, abi, depth);
773 }
774
775 fn readValueAtPosition(
776 builder: anytype,
777 memo: anytype,
778 index: kernel_root.Value,
779 facts: IndexFacts,
780 query: ValueReadQuery,
781 outline: kernelization_model.KernelOutline,
782 buffer_plan: *const bufferization.BufferPlanAnalysis,
783 abi: anytype,
784 depth: usize,
785 ) common.LoweringError!kernel_root.Value {
786 switch (query) {
787 .value => |value| {
788 if (depth == 0) return error.UnsupportedOperation;
789 if (memo.*) |*existing| {
790 if (existing.get(value)) |cached| return cached;
791 }
792
793 const resolved = resolved: {
794 if (value.getDefiningOp()) |def_any| {
795 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
796 if (isName(def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
797 if (try common.optionalSplatConstantValue(
798 builder,
799 def_op,
800 )) |scalar| break :resolved scalar;
801 }
802 if (isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) {
803 const operands = def_op.getOperandValues();
804 if (operands.len != 1) return error.InvalidArtifact;
805 if (operands[0].getDefiningOp()) |source_any| {
806 const source_op: *ir.Operation = @ptrCast(@alignCast(source_any));
807 if (isName(source_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
808 if (try common.optionalSplatConstantValue(
809 builder,
810 source_op,
811 )) |scalar| break :resolved scalar;
812 }
813 }
814 const source_slot = buffer_plan.getSlot(operands[0]) orelse return error.UnsupportedOperation;
815 const source_input_index = externalInputIndex(outline, source_slot.id) orelse return error.UnsupportedOperation;
816 const source_index = try broadcastInDimSourceIndex(builder, index, def_op, source_slot);
817 break :resolved builder.loadIndex(abi.input(builder, source_input_index), source_index) catch |err| return mapKernelBuildError(err);
818 }
819 if (isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastOp.operation_name)) {
820 const operands = def_op.getOperandValues();
821 if (operands.len != 1) return error.InvalidArtifact;
822 if (operands[0].getDefiningOp()) |source_any| {
823 const source_op: *ir.Operation = @ptrCast(@alignCast(source_any));
824 if (isName(source_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
825 if (try common.optionalSplatConstantValue(
826 builder,
827 source_op,
828 )) |scalar| break :resolved scalar;
829 }
830 }
831 }
832 }
833
834 if (buffer_plan.getSlot(value)) |slot| {
835 const input_index = externalInputIndex(outline, slot.id) orelse return error.UnsupportedOperation;
836 break :resolved builder.loadIndex(abi.input(builder, input_index), index) catch |err| return mapKernelBuildError(err);
837 }
838
839 const def_any = value.getDefiningOp() orelse return error.UnsupportedOperation;
840 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
841 if (isName(def_op.name.name, dialect_mod.AccyDialect.IotaOp.operation_name)) {
842 break :resolved try iotaValueAtIndex(builder, index, facts, def_op, value);
843 }
844 if (isName(def_op.name.name, dialect_mod.AccyDialect.SliceOp.operation_name)) {
845 break :resolved try readValueAtPosition(builder, memo, index, facts, .{ .slice = def_op }, outline, buffer_plan, abi, depth);
846 }
847 if (isName(def_op.name.name, dialect_mod.AccyDialect.PadOp.operation_name)) {
848 break :resolved try readValueAtPosition(builder, memo, index, facts, .{ .pad = def_op }, outline, buffer_plan, abi, depth);
849 }
850 if (isName(def_op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name)) {
851 break :resolved try readValueAtPosition(builder, memo, index, facts, .{ .gather = def_op }, outline, buffer_plan, abi, depth);
852 }
853 if (kernelForOperation(def_op)) |kind| {
854 const operands = def_op.getOperandValues();
855 if (operands.len > 3) return error.UnsupportedOperation;
856 var inputs: [3]kernel_root.Value = undefined;
857 for (operands, 0..) |operand, operand_index| {
858 inputs[operand_index] = try readValueAtPosition(builder, memo, index, facts, .{ .value = operand }, outline, buffer_plan, abi, depth - 1);
859 }
860 var dtype_arena_buffer: [160]u8 = undefined;
861 var dtype_arena = alloc_fixed.FixedBuffer.init(dtype_arena_buffer[0..]);
862 const result_type = dialect_mod.decodeTensorType(dtype_arena.allocator(), value.type) catch return error.InvalidArtifact;
863 break :resolved emitElementwiseValue(builder, kind, inputs[0..operands.len], result_type.dtype, def_op) catch |err| return common.mapGeneratedKernelError(err);
864 }
865 return error.UnsupportedOperation;
866 };
867
868 if (memo.*) |*existing| {
869 existing.put(value, resolved) catch return error.OutOfMemory;
870 }
871 return resolved;
872 },
873 .slice => |def_op| {
874 const operands = def_op.getOperandValues();
875 if (operands.len != 1) return error.InvalidArtifact;
876 const result = def_op.getResult(0) orelse return error.InvalidArtifact;
877
878 var result_dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
879 const result_dims = try decodedDims(result, result_dims_buffer[0..]);
880 var source_dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
881 const source_dims = try decodedDims(operands[0], source_dims_buffer[0..]);
882 if (result_dims.len != source_dims.len) return error.InvalidArtifact;
883
884 var starts_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
885 const starts = try readI64ListAttrBounded(def_op, "starts", dialect_mod.AccyDialect.SliceOp.dialectAttrName("starts"), &starts_buffer);
886 var strides_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
887 const strides = try readI64ListAttrBounded(def_op, "strides", dialect_mod.AccyDialect.SliceOp.dialectAttrName("strides"), &strides_buffer);
888 if (starts.len != result_dims.len or strides.len != result_dims.len) return error.InvalidArtifact;
889
890 var source_index = builder.constantIndex(0) catch |err| return mapKernelBuildError(err);
891 for (result_dims, 0..) |dim, dim_index| {
892 if (dim <= 0) return error.CapabilityMismatch;
893 if (strides[dim_index] <= 0) return error.CapabilityMismatch;
894 const coord = try coordinateFromShape(builder, index, facts, result_dims, dim_index);
895 const source_stride = try rowMajorStrideAt(source_dims, dim_index);
896 var source_coord = coord;
897 if (strides[dim_index] != 1) {
898 const stride_value = builder.constantIndex(strides[dim_index]) catch |err| return mapKernelBuildError(err);
899 source_coord = builder.mul(source_coord, stride_value) catch |err| return mapKernelBuildError(err);
900 }
901 if (starts[dim_index] != 0) {
902 const start_value = builder.constantIndex(starts[dim_index]) catch |err| return mapKernelBuildError(err);
903 source_coord = builder.add(source_coord, start_value) catch |err| return mapKernelBuildError(err);
904 }
905 const term = if (source_stride == 1) source_coord else blk: {
906 const stride_value = builder.constantIndex(source_stride) catch |err| return mapKernelBuildError(err);
907 break :blk builder.mul(source_coord, stride_value) catch |err| return mapKernelBuildError(err);
908 };
909 source_index = builder.add(source_index, term) catch |err| return mapKernelBuildError(err);
910 }
911 var shifted_memo: ?ValueMemo = null;
912 defer if (shifted_memo) |*existing| existing.deinit();
913 return readValueAtPosition(builder, &shifted_memo, source_index, .{}, .{ .value = operands[0] }, outline, buffer_plan, abi, depth - 1);
914 },
915 .gather => |def_op| {
916 const operands = def_op.getOperandValues();
917 if (operands.len != 2) return error.InvalidArtifact;
918 const axis = try common.readIntegerAttr(def_op, "axis");
919 if (axis != 0) return error.UnsupportedOperation;
920 var src_dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
921 const src_dims = try decodedDims(operands[0], src_dims_buffer[0..]);
922 if (src_dims.len != 1) return error.UnsupportedOperation;
923 const axis_size = src_dims[0];
924 if (axis_size <= 0) return error.CapabilityMismatch;
925
926 const source_index = if (try gatherIotaSourceIndex(builder, index, facts, operands[1], axis_size)) |proven_index| proven_index else if (try range_analysis.gatherIndexProvenInBounds(operands[1], axis_size, depth - 1)) blk: {
927 const idx_value = try readValueAtPosition(builder, memo, index, facts, .{ .value = operands[1] }, outline, buffer_plan, abi, depth - 1);
928 break :blk builder.castIndex(idx_value) catch |err| return mapKernelBuildError(err);
929 } else blk: {
930 const idx_value = try readValueAtPosition(builder, memo, index, facts, .{ .value = operands[1] }, outline, buffer_plan, abi, depth - 1);
931 const zero_i32 = builder.constantInt(.i32, 0) catch |err| return mapKernelBuildError(err);
932 const limit_value = @min(axis_size - 1, @as(i64, std.math.maxInt(i32)));
933 const limit_i32 = builder.constantInt(.i32, limit_value) catch |err| return mapKernelBuildError(err);
934 const lower_clamped = builder.max(idx_value, zero_i32) catch |err| return mapKernelBuildError(err);
935 const clamped = builder.min(lower_clamped, limit_i32) catch |err| return mapKernelBuildError(err);
936 break :blk builder.castIndex(clamped) catch |err| return mapKernelBuildError(err);
937 };
938
939 if (buffer_plan.getSlot(operands[0])) |slot| {
940 if (externalInputIndex(outline, slot.id)) |input_index| {
941 return builder.loadIndex(abi.input(builder, input_index), source_index) catch |err| return mapKernelBuildError(err);
942 }
943 }
944 var shifted_memo: ?ValueMemo = null;
945 defer if (shifted_memo) |*existing| existing.deinit();
946 return readValueAtPosition(builder, &shifted_memo, source_index, .{}, .{ .value = operands[0] }, outline, buffer_plan, abi, depth - 1);
947 },
948 .pad => |def_op| {
949 const operands = def_op.getOperandValues();
950 if (operands.len != 2) return error.InvalidArtifact;
951 const result = def_op.getResult(0) orelse return error.InvalidArtifact;
952
953 var result_dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
954 const result_dims = try decodedDims(result, result_dims_buffer[0..]);
955 var source_dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
956 const source_dims = try decodedDims(operands[0], source_dims_buffer[0..]);
957 if (result_dims.len != source_dims.len) return error.InvalidArtifact;
958
959 var low_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
960 const edge_low = try readI64ListAttrBounded(def_op, "edge_low", dialect_mod.AccyDialect.PadOp.dialectAttrName("edge_low"), &low_buffer);
961 var interior_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
962 const interior = try readI64ListAttrBounded(def_op, "interior", dialect_mod.AccyDialect.PadOp.dialectAttrName("interior"), &interior_buffer);
963 if (edge_low.len != result_dims.len or interior.len != result_dims.len) return error.InvalidArtifact;
964 for (interior) |dilation| if (dilation != 0) return error.UnsupportedOperation;
965
966 const zero = builder.constantIndex(0) catch |err| return mapKernelBuildError(err);
967 var source_index = zero;
968 var in_bounds: ?kernel_root.Value = null;
969 for (result_dims, 0..) |dim, dim_index| {
970 if (dim <= 0) return error.CapabilityMismatch;
971 const coord = try coordinateFromShape(builder, index, facts, result_dims, dim_index);
972 var source_coord = coord;
973 if (edge_low[dim_index] != 0) {
974 const low_value = builder.constantIndex(edge_low[dim_index]) catch |err| return mapKernelBuildError(err);
975 source_coord = builder.sub(source_coord, low_value) catch |err| return mapKernelBuildError(err);
976 }
977 const dim_extent = builder.constantIndex(source_dims[dim_index]) catch |err| return mapKernelBuildError(err);
978 const dim_edge = builder.constantIndex(source_dims[dim_index] - 1) catch |err| return mapKernelBuildError(err);
979 const low_ok = builder.compare(.ge, source_coord, zero) catch |err| return mapKernelBuildError(err);
980 const high_ok = builder.compare(.lt, source_coord, dim_extent) catch |err| return mapKernelBuildError(err);
981 in_bounds = if (in_bounds) |previous| blk: {
982 const with_low = builder.select(low_ok, previous, low_ok) catch |err| return mapKernelBuildError(err);
983 break :blk builder.select(high_ok, with_low, high_ok) catch |err| return mapKernelBuildError(err);
984 } else blk: {
985 break :blk builder.select(low_ok, high_ok, low_ok) catch |err| return mapKernelBuildError(err);
986 };
987 const clamped_low = builder.max(source_coord, zero) catch |err| return mapKernelBuildError(err);
988 const clamped = builder.min(clamped_low, dim_edge) catch |err| return mapKernelBuildError(err);
989 const source_stride = try rowMajorStrideAt(source_dims, dim_index);
990 const term = if (source_stride == 1) clamped else blk: {
991 const stride_value = builder.constantIndex(source_stride) catch |err| return mapKernelBuildError(err);
992 break :blk builder.mul(clamped, stride_value) catch |err| return mapKernelBuildError(err);
993 };
994 source_index = builder.add(source_index, term) catch |err| return mapKernelBuildError(err);
995 }
996
997 var shifted_memo: ?ValueMemo = null;
998 defer if (shifted_memo) |*existing| existing.deinit();
999 const loaded = try readValueAtPosition(builder, &shifted_memo, source_index, .{}, .{ .value = operands[0] }, outline, buffer_plan, abi, depth - 1);
1000
1001 const pad_slot = buffer_plan.getSlot(operands[1]) orelse return error.UnsupportedOperation;
1002 const pad_input_index = externalInputIndex(outline, pad_slot.id) orelse return error.UnsupportedOperation;
1003 const pad_value = builder.loadIndex(abi.input(builder, pad_input_index), zero) catch |err| return mapKernelBuildError(err);
1004
1005 const bounds = in_bounds orelse return error.InvalidArtifact;
1006 return builder.select(bounds, loaded, pad_value) catch |err| return mapKernelBuildError(err);
1007 },
1008 }
1009 }
1010
1011 fn gatherIotaSourceIndex(
1012 builder: anytype,
1013 index: kernel_root.Value,
1014 facts: IndexFacts,
1015 value: *ir.Value,
1016 axis_size: i64,
1017 ) common.LoweringError!?kernel_root.Value {
1018 const def_any = value.getDefiningOp() orelse return null;
1019 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
1020 if (!isName(def_op.name.name, dialect_mod.AccyDialect.IotaOp.operation_name)) return null;
1021
1022 const axis_attr = def_op.getAttrAs(ir.Attribute.IntegerAttr, "iota_dimension") orelse return error.InvalidArtifact;
1023 const axis = axis_attr.getValue();
1024 if (axis < 0) return error.InvalidArtifact;
1025
1026 var dims_arena_buffer: [256]u8 = undefined;
1027 var dims_arena = alloc_fixed.FixedBuffer.init(dims_arena_buffer[0..]);
1028 const result_type = dialect_mod.decodeTensorType(dims_arena.allocator(), value.type) catch return error.InvalidArtifact;
1029 if (axis >= result_type.dims.len) return error.InvalidArtifact;
1030 const axis_dim = result_type.dims[@intCast(axis)];
1031 if (axis_dim <= 0) return error.CapabilityMismatch;
1032 if (axis_dim > axis_size) return null;
1033 if (!nonNegativeIotaRangeFitsDType(axis_dim, result_type.dtype)) return null;
1034
1035 return try coordinateFromShape(builder, index, facts, result_type.dims, @intCast(axis));
1036 }
1037
1038 fn nonNegativeIotaRangeFitsDType(axis_dim: i64, dtype: choir_abi.DType) bool {
1039 if (axis_dim <= 0) return false;
1040 const last = axis_dim - 1;
1041 return switch (dtype) {
1042 .i8 => last <= std.math.maxInt(i8),
1043 .i16 => last <= std.math.maxInt(i16),
1044 .i32 => last <= std.math.maxInt(i32),
1045 .i64 => true,
1046 .u8 => last <= std.math.maxInt(u8),
1047 .u16 => last <= std.math.maxInt(u16),
1048 .u32 => last <= std.math.maxInt(u32),
1049 .u64 => true,
1050 else => false,
1051 };
1052 }
1053
1054 fn decodedDims(value: *ir.Value, buffer: []i64) common.LoweringError![]const i64 {
1055 var arena_buffer: [256]u8 = undefined;
1056 var arena = alloc_fixed.FixedBuffer.init(arena_buffer[0..]);
1057 const decoded = dialect_mod.decodeTensorType(arena.allocator(), value.type) catch return error.InvalidArtifact;
1058 if (decoded.dims.len > buffer.len) return error.CapabilityMismatch;
1059 @memcpy(buffer[0..decoded.dims.len], decoded.dims);
1060 return buffer[0..decoded.dims.len];
1061 }
1062
1063 fn rowMajorStrideAt(dims: []const i64, dim_index: usize) common.LoweringError!i64 {
1064 var stride: i64 = 1;
1065 for (dims[dim_index + 1 ..]) |dim| {
1066 if (dim <= 0) return error.CapabilityMismatch;
1067 stride = std.math.mul(i64, stride, dim) catch return error.CapabilityMismatch;
1068 }
1069 return stride;
1070 }
1071
1072 fn iota_coordinate(
1073 inner: anytype,
1074 flat: kernel_root.Value,
1075 stride_value: i64,
1076 extent: i64,
1077 ) !kernel_root.Value {
1078 const stride_extent = try inner.constantIndex(stride_value);
1079 const scaled = try inner.div(flat, stride_extent);
1080 const extent_value = try inner.constantIndex(extent);
1081 const wraps = try inner.div(scaled, extent_value);
1082 return inner.sub(scaled, try inner.mul(wraps, extent_value));
1083 }
1084
1085 fn iotaValueAtIndex(
1086 builder: anytype,
1087 index: kernel_root.Value,
1088 facts: IndexFacts,
1089 def_op: *ir.Operation,
1090 value: *ir.Value,
1091 ) common.LoweringError!kernel_root.Value {
1092 const axis_attr = def_op.getAttrAs(ir.Attribute.IntegerAttr, "iota_dimension") orelse return error.InvalidArtifact;
1093 const axis = axis_attr.getValue();
1094 if (axis < 0) return error.InvalidArtifact;
1095
1096 var dims_arena_buffer: [256]u8 = undefined;
1097 var dims_arena = alloc_fixed.FixedBuffer.init(dims_arena_buffer[0..]);
1098 const result_type = dialect_mod.decodeTensorType(dims_arena.allocator(), value.type) catch return error.InvalidArtifact;
1099 if (axis >= result_type.dims.len) return error.InvalidArtifact;
1100
1101 var stride: i64 = 1;
1102 var dim_index = result_type.dims.len;
1103 while (dim_index > @as(usize, @intCast(axis)) + 1) {
1104 dim_index -= 1;
1105 const dim = result_type.dims[dim_index];
1106 if (dim < 1) return error.UnsupportedOperation;
1107 stride *= dim;
1108 }
1109 const axis_dim = result_type.dims[@intCast(axis)];
1110 if (axis_dim < 1) return error.UnsupportedOperation;
1111 if (facts.rank2) |rank2| {
1112 if (rank2.coordinate(result_type.dims, @intCast(axis))) |coord| {
1113 return builder.cast(coord, result_type.dtype) catch |err| return mapKernelBuildError(err);
1114 }
1115 }
1116
1117 const coord = iota_coordinate(builder, index, stride, axis_dim) catch |err| return mapKernelBuildError(err);
1118 const casted = builder.cast(coord, result_type.dtype) catch |err| return mapKernelBuildError(err);
1119 return casted;
1120 }
1121
1122 fn broadcastInDimElementwiseOperandValue(
1123 builder: anytype,
1124 index: kernel_root.Value,
1125 operand: *ir.Value,
1126 outline: kernelization_model.KernelOutline,
1127 buffer_plan: *const bufferization.BufferPlanAnalysis,
1128 abi: generated_abi.Flat,
1129 ) common.LoweringError!?kernel_root.Value {
1130 const def_any = operand.getDefiningOp() orelse return null;
1131 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
1132 if (!isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return null;
1133 const operands = def_op.getOperandValues();
1134 if (operands.len != 1) return error.InvalidArtifact;
1135
1136 const source_slot = buffer_plan.getSlot(operands[0]) orelse return error.UnsupportedOperation;
1137 const source_input_index = externalInputIndex(outline, source_slot.id) orelse return error.UnsupportedOperation;
1138 const source_index = try broadcastInDimSourceIndex(builder, index, def_op, source_slot);
1139 return builder.loadIndex(abi.input(builder, source_input_index), source_index) catch |err| return mapKernelBuildError(err);
1140 }
1141
1142 pub fn broadcastInDimSourceIndex(
1143 builder: anytype,
1144 index: kernel_root.Value,
1145 op: *ir.Operation,
1146 source_slot: *const bufferization.BufferSlot,
1147 ) common.LoweringError!kernel_root.Value {
1148 var broadcast_dims_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
1149 const broadcast_dims = try readI64ListAttrBounded(op, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims"), &broadcast_dims_buffer);
1150 var result_shape_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
1151 const result_shape = try readI64ListAttrBounded(op, "result_shape", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("result_shape"), &result_shape_buffer);
1152 if (broadcast_dims.len != source_slot.dims.len) return error.InvalidArtifact;
1153 try validateBroadcastDims(broadcast_dims, result_shape.len);
1154 const source_strides = source_slot.row_major_strides orelse return error.CapabilityMismatch;
1155 if (source_strides.len != source_slot.dims.len) return error.InvalidArtifact;
1156
1157 var source_index = builder.constantIndex(0) catch |err| return mapKernelBuildError(err);
1158 for (result_shape, 0..) |dim, dim_index| {
1159 if (dim <= 0) return error.CapabilityMismatch;
1160 const coord = try linearCoordinateFromShape(builder, index, result_shape, dim_index);
1161
1162 const source_dim = (try sourceDimForOutputDim(broadcast_dims, dim_index)) orelse continue;
1163 if (source_dim >= source_slot.dims.len) return error.InvalidArtifact;
1164 const source_dim_size = source_slot.dims[source_dim];
1165 if (source_dim_size <= 0) return error.CapabilityMismatch;
1166 if (source_dim_size != 1 and source_dim_size != dim) return error.InvalidArtifact;
1167 if (source_dim_size == 1) continue;
1168
1169 const source_stride = std.math.cast(i64, source_strides[source_dim]) orelse return error.CapabilityMismatch;
1170 const term = if (source_stride == 1) coord else blk: {
1171 const stride_value = builder.constantIndex(source_stride) catch |err| return mapKernelBuildError(err);
1172 break :blk builder.mul(coord, stride_value) catch |err| return mapKernelBuildError(err);
1173 };
1174 source_index = builder.add(source_index, term) catch |err| return mapKernelBuildError(err);
1175 }
1176 return source_index;
1177 }
1178
1179 fn resultStrideI64(result_shape: []const i64, dim_index: usize) common.LoweringError!i64 {
1180 var stride: i64 = 1;
1181 for (result_shape[dim_index + 1 ..]) |dim| {
1182 if (dim <= 0) return error.CapabilityMismatch;
1183 stride = std.math.mul(i64, stride, dim) catch return error.CapabilityMismatch;
1184 }
1185 return stride;
1186 }
1187
1188 fn linearCoordinateFromShape(
1189 builder: anytype,
1190 linear_index: kernel_root.Value,
1191 shape: []const i64,
1192 dim_index: usize,
1193 ) common.LoweringError!kernel_root.Value {
1194 if (dim_index >= shape.len) return error.InvalidArtifact;
1195 const dim = shape[dim_index];
1196 if (dim <= 0) return error.CapabilityMismatch;
1197 const stride = try resultStrideI64(shape, dim_index);
1198 const quotient = if (stride == 1) linear_index else blk: {
1199 const stride_value = builder.constantIndex(stride) catch |err| return mapKernelBuildError(err);
1200 break :blk builder.div(linear_index, stride_value) catch |err| return mapKernelBuildError(err);
1201 };
1202 const dim_value = builder.constantIndex(dim) catch |err| return mapKernelBuildError(err);
1203 const whole_dim_count = builder.div(quotient, dim_value) catch |err| return mapKernelBuildError(err);
1204 const consumed = builder.mul(whole_dim_count, dim_value) catch |err| return mapKernelBuildError(err);
1205 return builder.sub(quotient, consumed) catch |err| return mapKernelBuildError(err);
1206 }
1207
1208 fn coordinateFromShape(
1209 builder: anytype,
1210 linear_index: kernel_root.Value,
1211 facts: IndexFacts,
1212 shape: []const i64,
1213 dim_index: usize,
1214 ) common.LoweringError!kernel_root.Value {
1215 if (facts.rank2) |rank2| {
1216 if (rank2.coordinate(shape, dim_index)) |coord| return coord;
1217 }
1218 return linearCoordinateFromShape(builder, linear_index, shape, dim_index);
1219 }
1220
1221 fn validateBroadcastDims(
1222 broadcast_dims: []const i64,
1223 result_rank: usize,
1224 ) common.LoweringError!void {
1225 for (broadcast_dims, 0..) |dim, index| {
1226 const dim_index = std.math.cast(usize, dim) orelse return error.InvalidArtifact;
1227 if (dim_index >= result_rank) return error.InvalidArtifact;
1228 for (broadcast_dims[0..index]) |previous| {
1229 const previous_index = std.math.cast(usize, previous) orelse return error.InvalidArtifact;
1230 if (previous_index == dim_index) return error.InvalidArtifact;
1231 }
1232 }
1233 }
1234
1235 fn sourceDimForOutputDim(
1236 broadcast_dims: []const i64,
1237 dim_index: usize,
1238 ) common.LoweringError!?usize {
1239 for (broadcast_dims, 0..) |output_dim, source_dim| {
1240 const output_dim_index = std.math.cast(usize, output_dim) orelse return error.InvalidArtifact;
1241 if (output_dim_index == dim_index) return source_dim;
1242 }
1243 return null;
1244 }
1245
1246 fn elementwiseChainValue(
1247 operand: *ir.Value,
1248 ops: []const *ir.Operation,
1249 op_index: usize,
1250 op_results: []const kernel_root.Value,
1251 ) ?kernel_root.Value {
1252 const def_any = operand.getDefiningOp() orelse return null;
1253 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
1254 for (ops[0..op_index], 0..) |op, index| {
1255 if (op == def_op) return op_results[index];
1256 }
1257 return null;
1258 }
1259
1260 fn elementwiseOutputValue(
1261 work: schedule_planning.ScheduleWorkItem,
1262 ops: []const *ir.Operation,
1263 op_results: []const kernel_root.Value,
1264 ) ?kernel_root.Value {
1265 for (ops, 0..) |op, index| {
1266 if (op == work.root) return op_results[index];
1267 }
1268 if (op_results.len == 0) return null;
1269 return op_results[op_results.len - 1];
1270 }
1271
1272 pub fn emitElementwiseValue(
1273 builder: anytype,
1274 kind: ElementwiseKernel,
1275 inputs: []const kernel_root.Value,
1276 work_dtype: choir_abi.DType,
1277 op: *ir.Operation,
1278 ) !kernel_root.Value {
1279 return switch (kind) {
1280 .convert => try convertValue(builder, inputs[0], work_dtype),
1281 .compare => try builder.compare(try comparePredicate(op), inputs[0], inputs[1]),
1282 .select => try builder.select(inputs[0], inputs[1], inputs[2]),
1283 .add => try builder.add(inputs[0], inputs[1]),
1284 .sub => try builder.sub(inputs[0], inputs[1]),
1285 .mul => try builder.mul(inputs[0], inputs[1]),
1286 .div => try builder.div(inputs[0], inputs[1]),
1287 .min => try builder.min(inputs[0], inputs[1]),
1288 .max => try builder.max(inputs[0], inputs[1]),
1289 .neg => try builder.neg(inputs[0]),
1290 .abs => try builder.abs(inputs[0]),
1291 .sqrt => try builder.sqrt(inputs[0]),
1292 .exp => try builder.exp(inputs[0]),
1293 .log => try builder.log(inputs[0]),
1294 .tanh => try builder.tanh(inputs[0]),
1295 .sin => try builder.sin(inputs[0]),
1296 .cos => try builder.cos(inputs[0]),
1297 .tan => try builder.tan(inputs[0]),
1298 .floor => try builder.floor(inputs[0]),
1299 .round => try builder.round(inputs[0]),
1300 .trunc => try builder.trunc(inputs[0]),
1301 .pow => try builder.pow(inputs[0], inputs[1]),
1302 .atan2 => try builder.atan2(inputs[0], inputs[1]),
1303 };
1304 }
1305
1306 fn comparePredicate(op: *ir.Operation) !kernel_root.Compare {
1307 const payload = try common.dialectAttrPayload(
1308 op,
1309 "compare_direction",
1310 dialect_mod.AccyDialect.CompareOp.dialectAttrName("compare_direction"),
1311 );
1312 if (std.mem.eql(u8, payload, "eq")) return .eq;
1313 if (std.mem.eql(u8, payload, "ne")) return .ne;
1314 if (std.mem.eql(u8, payload, "lt")) return .lt;
1315 if (std.mem.eql(u8, payload, "le")) return .le;
1316 if (std.mem.eql(u8, payload, "gt")) return .gt;
1317 if (std.mem.eql(u8, payload, "ge")) return .ge;
1318 return error.UnsupportedOperation;
1319 }
1320
1321 fn convertValue(
1322 builder: anytype,
1323 value: kernel_root.Value,
1324 dtype: choir_abi.DType,
1325 ) !kernel_root.Value {
1326 return switch (dtype) {
1327 .f32 => (try builder.castValue(value, .f32)).raw(),
1328 .i8 => (try builder.castValue(value, .i8)).raw(),
1329 .i16 => (try builder.castValue(value, .i16)).raw(),
1330 .i32 => (try builder.castValue(value, .i32)).raw(),
1331 .u8 => (try builder.castValue(value, .u8)).raw(),
1332 .u16 => (try builder.castValue(value, .u16)).raw(),
1333 .u32 => (try builder.castValue(value, .u32)).raw(),
1334 .i64 => (try builder.castValue(value, .i64)).raw(),
1335 .u64 => (try builder.castValue(value, .u64)).raw(),
1336 .f16 => (try builder.castValue(value, .f16)).raw(),
1337 .bf16 => (try builder.castValue(value, .bf16)).raw(),
1338 .f64 => (try builder.castValue(value, .f64)).raw(),
1339 else => error.CapabilityMismatch,
1340 };
1341 }
1342
1343 pub fn kernelForOperation(op: *ir.Operation) ?ElementwiseKernel {
1344 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ConvertOp.operation_name)) return .convert;
1345 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.CompareOp.operation_name)) return .compare;
1346 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SelectOp.operation_name)) return .select;
1347 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.AddOp.operation_name)) return .add;
1348 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SubOp.operation_name)) return .sub;
1349 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MulOp.operation_name)) return .mul;
1350 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.DivOp.operation_name)) return .div;
1351 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MinOp.operation_name)) return .min;
1352 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MaxOp.operation_name)) return .max;
1353 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.NegOp.operation_name)) return .neg;
1354 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.AbsOp.operation_name)) return .abs;
1355 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SqrtOp.operation_name)) return .sqrt;
1356 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ExpOp.operation_name)) return .exp;
1357 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.LogOp.operation_name)) return .log;
1358 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TanhOp.operation_name)) return .tanh;
1359 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SinOp.operation_name)) return .sin;
1360 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.CosOp.operation_name)) return .cos;
1361 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TanOp.operation_name)) return .tan;
1362 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.FloorOp.operation_name)) return .floor;
1363 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.RoundOp.operation_name)) return .round;
1364 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TruncOp.operation_name)) return .trunc;
1365 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.PowOp.operation_name)) return .pow;
1366 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.Atan2Op.operation_name)) return .atan2;
1367 return null;
1368 }
1369
1370 pub fn expectedInputCount(kind: ElementwiseKernel) usize {
1371 return switch (kind) {
1372 .add, .sub, .mul, .div, .min, .max, .pow, .atan2, .compare => 2,
1373 .neg, .abs, .sqrt, .exp, .log, .tanh, .sin, .cos, .tan, .floor, .round, .trunc, .convert => 1,
1374 .select => 3,
1375 };
1376 }
1377
1378 const testing = std.testing;
1379 const semantic = @import("../../../choir/root.zig").semantic;
1380
1381 const OperandFixture = struct {
1382 module: *semantic.SemanticModule,
1383 root: *ir.Operation,
1384 inputs: [2]*ir.Value,
1385 };
1386
1387 fn operandFixture(
1388 source: *semantic.Builder,
1389 depth: usize,
1390 shared: bool,
1391 gather: bool,
1392 ) !OperandFixture {
1393 const typ = try source.tensor(.f32, &.{4});
1394 const index_type = try source.tensor(.i32, &.{4});
1395 var function = try source.beginFunction("operand_walk", &.{ typ, typ, index_type }, &.{typ});
1396 const inputs = [2]*ir.Value{ function.parameter(0), function.parameter(1) };
1397 var value = inputs[0];
1398 if (gather) value = try function.gather(value, function.parameter(2), typ, 0);
1399 for (0..depth) |_| {
1400 value = if (shared) try function.add(value, value) else try function.neg(value);
1401 }
1402 value = try function.add(inputs[1], value);
1403 try function.return_(&.{value});
1404 try function.finish();
1405 return .{
1406 .module = try source.finish(),
1407 .root = @ptrCast(@alignCast(value.getDefiningOp().?)),
1408 .inputs = inputs,
1409 };
1410 }
1411
1412 test "kernelization operand walk finds gathers beyond former depth and capacity limits" {
1413 for ([_]usize{ 0, 31, 32, 33, 127, 128, 129, 257 }) |depth| {
1414 for ([_]bool{ false, true }) |gather| {
1415 var source = try semantic.Builder.init(testing.allocator, .standard);
1416 defer source.deinit();
1417 const fixture = try operandFixture(&source, depth, false, gather);
1418 defer fixture.module.deinit();
1419 try testing.expectEqual(
1420 gather,
1421 try closureContainsGather(testing.allocator, &.{fixture.root}),
1422 );
1423 }
1424 }
1425 }
1426
1427 test "kernelization operand walk expands shared definitions once" {
1428 for ([_]usize{ 0, 32, 129, 257 }) |depth| {
1429 var source = try semantic.Builder.init(testing.allocator, .standard);
1430 defer source.deinit();
1431 const fixture = try operandFixture(&source, depth, true, false);
1432 defer fixture.module.deinit();
1433 var walk = OperandWalk.init(testing.allocator);
1434 defer walk.deinit();
1435 try walk.expand(fixture.root);
1436 var edges: usize = 0;
1437 while (walk.next()) |value| {
1438 edges += 1;
1439 try testing.expect(edges <= 2 * (depth + 1));
1440 if (value.getDefiningOp()) |def| try walk.expand(@ptrCast(@alignCast(def)));
1441 }
1442 try testing.expectEqual(depth + 1, walk.expanded.count());
1443 try testing.expectEqual(2 * (depth + 1), edges);
1444 try testing.expectEqual(
1445 false,
1446 try closureContainsGather(testing.allocator, &.{fixture.root}),
1447 );
1448 }
1449 }
1450
1451 fn inputFixture(inputs: [2]*ir.Value) !bufferization.BufferPlanAnalysis {
1452 var plan = bufferization.BufferPlanAnalysis.init(testing.allocator);
1453 errdefer plan.deinit();
1454 for (inputs, 0..) |value, index| {
1455 try plan.slots.append(testing.allocator, .{
1456 .id = index,
1457 .value = value,
1458 .producer = null,
1459 .function = null,
1460 .role = .{ .input = true },
1461 .dtype = .f32,
1462 .dims = &.{},
1463 .element_count = 4,
1464 .row_major_strides = null,
1465 .byte_size = 16,
1466 });
1467 try plan.value_to_slot.put(value, index);
1468 }
1469 return plan;
1470 }
1471
1472 fn checkVectorInputWalk(allocator: std.mem.Allocator) !void {
1473 var source = try semantic.Builder.init(testing.allocator, .standard);
1474 defer source.deinit();
1475 const fixture = try operandFixture(&source, 40, true, false);
1476 defer fixture.module.deinit();
1477 var buffers = try inputFixture(fixture.inputs);
1478 defer buffers.deinit();
1479 var input_ids = [_]usize{ 0, 1 };
1480 const outline = kernelization_model.KernelOutline{
1481 .id = 0,
1482 .name = &.{},
1483 .kind = .elementwise,
1484 .work_item_id = 0,
1485 .root = fixture.root,
1486 .input_slot_ids = &input_ids,
1487 .output_slot_id = 2,
1488 .element_count = 4,
1489 .op_count = 1,
1490 };
1491 const inputs = try collectVectorInputs(allocator, &.{fixture.root}, outline, &buffers);
1492 defer allocator.free(inputs);
1493 try testing.expectEqual(@as(usize, 2), inputs.len);
1494 try testing.expect(inputs[0].value == fixture.inputs[1]);
1495 try testing.expectEqual(@as(usize, 1), inputs[0].input_index);
1496 try testing.expect(inputs[1].value == fixture.inputs[0]);
1497 try testing.expectEqual(@as(usize, 0), inputs[1].input_index);
1498 }
1499
1500 test "kernelization operand walk collects ordered unique deep vector inputs with host OOM" {
1501 try testing.checkAllAllocationFailures(testing.allocator, checkVectorInputWalk, .{});
1502 }
1503
1504 fn checkGatherWalkFailure(allocator: std.mem.Allocator, gather: bool) !void {
1505 var source = try semantic.Builder.init(testing.allocator, .standard);
1506 defer source.deinit();
1507 const fixture = try operandFixture(&source, 40, false, gather);
1508 defer fixture.module.deinit();
1509 try testing.expectEqual(gather, try closureContainsGather(allocator, &.{fixture.root}));
1510 }
1511
1512 test "kernelization operand walk preserves gather discovery allocation failure" {
1513 for ([_]bool{ false, true }) |gather| {
1514 try testing.checkAllAllocationFailures(
1515 testing.allocator,
1516 checkGatherWalkFailure,
1517 .{gather},
1518 );
1519 }
1520 }
1521
1522 test "kernelization operand walk fits graph derived reserved storage" {
1523 for ([_]usize{ 0, 32, 129, 257 }) |depth| {
1524 try checkOperandWalkStorage(depth);
1525 }
1526 }
1527
1528 fn checkOperandWalkStorage(depth: usize) !void {
1529 const accounting = @import("choir").passes.pass.work;
1530 var source = try semantic.Builder.init(testing.allocator, .standard);
1531 defer source.deinit();
1532 const fixture = try operandFixture(&source, depth, true, false);
1533 defer fixture.module.deinit();
1534 const census = try accounting.Census.inspect(fixture.module.choir_module);
1535 const bound = try accounting.add(
1536 try accounting.hashMapGrowth(*ir.Operation, void, census.operations),
1537 try accounting.arrayListGrowth(*ir.Value, census.operands),
1538 );
1539 const bytes = try testing.allocator.alignedAlloc(u8, .@"64", @intCast(bound));
1540 defer testing.allocator.free(bytes);
1541 var backing = alloc_fixed.Tracked.init(bytes);
1542 var retained = alloc_fixed.Monotonic.init(backing.allocator(), bytes.len);
1543 var walk = OperandWalk.init(retained.allocator());
1544 defer walk.deinit();
1545 try walk.expand(fixture.root);
1546 var visits: usize = 0;
1547 while (walk.next()) |value| {
1548 visits += 1;
1549 try testing.expect(visits <= census.operands);
1550 if (value.getDefiningOp()) |def| try walk.expand(@ptrCast(@alignCast(def)));
1551 }
1552 try testing.expectEqual(2 * (depth + 1), visits);
1553 try testing.expect(walk.expanded.count() <= census.operations);
1554 const used = if (retained.current) |*current| alloc_fixed.used(current) else 0;
1555 try testing.expect(used <= bound);
1556 try testing.expect(!backing.exhausted);
1557 }