lib/accy/src/tensor/trace/builder.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const tensor = @import("../root.zig");
3 const trace = @import("root.zig");
4
5 const Binary = trace.Binary;
6 const CompareDirection = trace.CompareDirection;
7 const DType = trace.DType;
8 const Dim = trace.Dim;
9 const Id = trace.Id;
10 const Program = trace.Program;
11 const Reducer = trace.Reducer;
12 const Spec = trace.Spec;
13 const Type = trace.Type;
14 const Unary = trace.Unary;
15 const Value = trace.Value;
16 const program_mod = tensor.program;
17 const type_mod = tensor.types;
18
19 pub const ScanScope = struct {
20 parent: *Builder,
21 child: Builder,
22 length: i64,
23 init_ids: []const Id,
24 phase: Phase = .open,
25
26 const Phase = enum { open, done };
27
28 pub fn body(self: *ScanScope) *Builder {
29 return &self.child;
30 }
31
32 pub fn carry(self: *ScanScope, index: usize) Value {
33 return self.child.makeValue(self.child.parameters_list.items[index]);
34 }
35
36 pub fn finish(self: *ScanScope, next: []const Value) !Value {
37 if (self.phase != .open) return error.ScanScopeMisused;
38 if (next.len != self.init_ids.len) return error.ScanArityMismatch;
39 const output_ids = try self.child.arena.allocator().alloc(Id, next.len);
40 for (next, output_ids) |value, *slot| {
41 try self.child.ensureValue(value);
42 slot.* = value.id;
43 }
44 const body_view = program_mod.Subgraph{
45 .values = self.child.values.items,
46 .operations = self.child.operations.items,
47 .parameters = self.child.parameters_list.items,
48 .outputs = output_ids,
49 };
50 const result = try self.parent.emitScanIds(self.length, self.init_ids, &body_view);
51 self.abort();
52 return result;
53 }
54
55 pub fn abort(self: *ScanScope) void {
56 if (self.phase == .done) return;
57 self.child.deinit();
58 self.phase = .done;
59 }
60 };
61
62 pub const Builder = struct {
63 allocator: std.mem.Allocator,
64 arena: std.heap.ArenaAllocator,
65 name: []const u8,
66 values: std.ArrayListUnmanaged(Type) = .empty,
67 operations: std.ArrayListUnmanaged(program_mod.Operation) = .empty,
68 parameters_list: std.ArrayListUnmanaged(Id) = .empty,
69 finished: bool = false,
70
71 pub fn init(allocator: std.mem.Allocator, name: []const u8) !Builder {
72 var arena = std.heap.ArenaAllocator.init(allocator);
73 errdefer arena.deinit();
74 var builder = Builder{
75 .allocator = allocator,
76 .arena = arena,
77 .name = &.{},
78 };
79 builder.name = try builder.arena.allocator().dupe(u8, name);
80 return builder;
81 }
82
83 pub fn deinit(self: *Builder) void {
84 if (!self.finished) {
85 self.arena.deinit();
86 self.finished = true;
87 }
88 }
89
90 pub fn input(self: *Builder, dtype: DType, dims_struct: anytype) !Value {
91 var buffer: [type_mod.dimCount(@TypeOf(dims_struct))]Dim = undefined;
92 type_mod.fillDims(dims_struct, &buffer);
93 return self.inputDims(dtype, &buffer);
94 }
95
96 pub fn inputDims(self: *Builder, dtype: DType, dims: []const Dim) !Value {
97 try type_mod.validateAuthoredDims(dims);
98 return self.inputTyped(.{ .dtype = dtype, .dims = dims });
99 }
100
101 pub fn inputSpec(self: *Builder, input_spec: Spec) !Value {
102 return self.inputDims(input_spec.dtype, input_spec.dims);
103 }
104
105 pub fn inputs(self: *Builder, specs: []const Spec) ![]Value {
106 const result = try self.arena.allocator().alloc(Value, specs.len);
107 for (specs, 0..) |input_spec, index| {
108 result[index] = try self.inputSpec(input_spec);
109 }
110 return result;
111 }
112
113 pub fn inputTyped(self: *Builder, ty: Type) !Value {
114 const owned = try self.copyType(ty);
115 const id = try self.append(
116 owned,
117 .{ .parameter = .{ .index = self.parameters_list.items.len } },
118 );
119 try self.parameters_list.append(self.arena.allocator(), id);
120 return self.makeValue(id);
121 }
122
123 pub fn operation(self: *Builder, op: *const program_mod.Operation, args: []const Value) !Value {
124 return switch (op.kind) {
125 .parameter => self.inputTyped(op.result),
126 .constant => |constant| self.constantBytes(op.result, constant.payload),
127 .iota => |iota_op| self.iotaTyped(op.result, iota_op.axis),
128 .unary => |unary_op| self.unary(unary_op.op, args[0]),
129 .binary => |binary_op| self.binary(binary_op.op, args[0], args[1]),
130 .broadcast => self.broadcastOp(args[0], op.result.dims),
131 .broadcast_in_dim => |broadcast_op| self.broadcastTo(args[0], op.result.dims, broadcast_op.broadcast_dims),
132 .reshape => self.reshapeTo(args[0], op.result.dims),
133 .transpose => |transpose_op| self.transposeBy(args[0], transpose_op.permutation),
134 .reduce => |reduce_op| self.reduce(args[0], args[1], reduce_op.reducer, reduce_op.dimensions),
135 .gather => |gather_op| self.gatherOp(args[0], args[1], gather_op.axis),
136 .scatter_add => |scatter_add| self.scatterAddOp(args[0], args[1], args[2], scatter_add.axis),
137 .sparse_cross_entropy => |sparse_cross_entropy_op| self.sparseCrossEntropyOp(args[0], args[1], sparse_cross_entropy_op.axis),
138 .compare => |compare_op| self.compare(compare_op.direction, args[0], args[1]),
139 .select => self.select(args[0], args[1], args[2]),
140 .custom_call => |custom_call| self.customCall(custom_call.target, custom_call.version, args, op.result),
141 .dot_general => |dot| self.dotGeneralOp(
142 args[0],
143 args[1],
144 dot.lhs_contract,
145 dot.rhs_contract,
146 dot.lhs_batch,
147 dot.rhs_batch,
148 ),
149 .scan => |scan_op| self.emitScan(scan_op.length, args[0..scan_op.inits.len], scan_op.body),
150 .projection => |projection_op| self.projection(args[0], projection_op.index),
151 };
152 }
153
154 pub fn scalar(self: *Builder, comptime dtype: DType, scalar_value: dtype.ZigType()) !Value {
155 const ty = Type.scalar(dtype);
156 return self.constantBytes(ty, std.mem.asBytes(&scalar_value));
157 }
158
159 pub fn full(self: *Builder, comptime dtype: DType, dims_struct: anytype, fill_value: dtype.ZigType()) !Value {
160 var buffer: [type_mod.dimCount(@TypeOf(dims_struct))]Dim = undefined;
161 type_mod.fillDims(dims_struct, &buffer);
162 return self.fullDims(dtype, &buffer, fill_value);
163 }
164
165 pub fn fullDims(self: *Builder, comptime dtype: DType, dims: []const Dim, fill_value: dtype.ZigType()) !Value {
166 try type_mod.validateAuthoredDims(dims);
167 const ty = try Type.init(self.arena.allocator(), dtype, dims);
168 const count = try ty.elementCount();
169 const payload = try self.arena.allocator().alloc(dtype.ZigType(), count);
170 for (payload) |*slot| {
171 slot.* = fill_value;
172 }
173 return self.constantOwned(ty, std.mem.sliceAsBytes(payload));
174 }
175
176 pub fn fullFloat(self: *Builder, ty: Type, fill_value: f64) !Value {
177 const owned = try self.copyType(ty);
178 const count = try owned.elementCount();
179 return switch (owned.dtype) {
180 .f16 => blk: {
181 const payload = try self.arena.allocator().alloc(f16, count);
182 const value: f16 = @floatCast(fill_value);
183 for (payload) |*slot| slot.* = value;
184 break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));
185 },
186 .bf16 => blk: {
187 const Bf16 = DType.bf16.ZigType();
188 const payload = try self.arena.allocator().alloc(Bf16, count);
189 const value = Bf16.fromF32(@floatCast(fill_value));
190 for (payload) |*slot| slot.* = value;
191 break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));
192 },
193 .f32 => blk: {
194 const payload = try self.arena.allocator().alloc(f32, count);
195 const value: f32 = @floatCast(fill_value);
196 for (payload) |*slot| slot.* = value;
197 break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));
198 },
199 .f64 => blk: {
200 const payload = try self.arena.allocator().alloc(f64, count);
201 for (payload) |*slot| slot.* = fill_value;
202 break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));
203 },
204 else => error.NonFloatDType,
205 };
206 }
207
208 pub fn constantBytes(self: *Builder, ty: Type, payload: []const u8) !Value {
209 if (payload.len != try ty.byteCount()) return error.PayloadLengthMismatch;
210 return self.constantOwned(try self.copyType(ty), try self.arena.allocator().dupe(u8, payload));
211 }
212
213 pub fn zeros(self: *Builder, ty: Type) !Value {
214 const owned = try self.copyType(ty);
215 const payload = try self.arena.allocator().alloc(u8, try owned.byteCount());
216 @memset(payload, 0);
217 return self.constantOwned(owned, payload);
218 }
219
220 pub fn isStructuralZero(self: *Builder, id: Id) bool {
221 const index: usize = @intCast(id.index);
222 if (index >= self.operations.items.len) return false;
223 return switch (self.operations.items[index].kind) {
224 .constant => |constant| program_mod.isZeroPayload(constant.payload),
225 .unary => |unary_op| switch (unary_op.op) {
226 .neg => self.isStructuralZero(unary_op.input),
227 else => false,
228 },
229 .binary => |binary_op| switch (binary_op.op) {
230 .add, .sub => self.isStructuralZero(binary_op.lhs) and self.isStructuralZero(binary_op.rhs),
231 .mul => self.isStructuralZero(binary_op.lhs) or self.isStructuralZero(binary_op.rhs),
232 else => false,
233 },
234 .broadcast => |broadcast_op| self.isStructuralZero(broadcast_op.input),
235 .broadcast_in_dim => |broadcast_op| self.isStructuralZero(broadcast_op.input),
236 .reshape => |reshape_op| self.isStructuralZero(reshape_op.input),
237 .transpose => |transpose_op| self.isStructuralZero(transpose_op.input),
238 .gather => |gather_op| self.isStructuralZero(gather_op.input),
239 .scatter_add => |scatter_add| self.isStructuralZero(scatter_add.input) and self.isStructuralZero(scatter_add.updates),
240 .select => |select_op| self.isStructuralZero(select_op.on_true) and self.isStructuralZero(select_op.on_false),
241 else => false,
242 };
243 }
244
245 pub fn iota(self: *Builder, dtype: DType, dims_struct: anytype, comptime axis: anytype) !Value {
246 var buffer: [type_mod.dimCount(@TypeOf(dims_struct))]Dim = undefined;
247 type_mod.fillDims(dims_struct, &buffer);
248 return self.iotaDims(dtype, &buffer, comptime nameOf(axis));
249 }
250
251 pub fn iotaDims(self: *Builder, dtype: DType, dims: []const Dim, axis_name: []const u8) !Value {
252 try type_mod.validateAuthoredDims(dims);
253 const index = type_mod.findDim(dims, axis_name) orelse return error.AxisNotFound;
254 return self.iotaTyped(.{ .dtype = dtype, .dims = dims }, @intCast(index));
255 }
256
257 fn iotaTyped(self: *Builder, ty: Type, axis: i64) !Value {
258 try type_mod.validateAxes(ty.dims.len, &.{axis});
259 return self.emit(try self.copyType(ty), .{ .iota = .{ .axis = axis } });
260 }
261
262 pub fn unary(self: *Builder, op: Unary, operand_value: Value) !Value {
263 try self.ensureValue(operand_value);
264 return self.emit(try self.copyType(operand_value.ty), .{
265 .unary = .{ .op = op, .input = operand_value.id },
266 });
267 }
268
269 pub fn binary(self: *Builder, op: Binary, lhs: Value, rhs: Value) !Value {
270 try self.ensureValue(lhs);
271 try self.ensureValue(rhs);
272 try type_mod.sameType(lhs.ty, rhs.ty);
273 return self.emit(try self.copyType(lhs.ty), .{
274 .binary = .{ .op = op, .lhs = lhs.id, .rhs = rhs.id },
275 });
276 }
277
278 pub fn alignedBinary(self: *Builder, op: Binary, lhs: Value, rhs: Value) !Value {
279 try self.ensureValue(lhs);
280 try self.ensureValue(rhs);
281 if (lhs.ty.dtype != rhs.ty.dtype) return error.DTypeMismatch;
282 const target = try type_mod.unionDims(self.arena.allocator(), lhs.ty.dims, rhs.ty.dims);
283 const left = try self.alignTo(lhs, target);
284 const right = try self.alignTo(rhs, target);
285 return self.binary(op, left, right);
286 }
287
288 pub fn customCall(
289 self: *Builder,
290 target: []const u8,
291 version: u32,
292 operands: []const Value,
293 result_ty: Type,
294 ) !Value {
295 if (target.len == 0) return error.InvalidCustomCallContract;
296 if (version == 0) return error.InvalidCustomCallContract;
297 if (operands.len > program_mod.max_custom_call_operands) return error.UnsupportedCustomCallArity;
298 const arena = self.arena.allocator();
299 const ids = try arena.alloc(program_mod.Id, operands.len);
300 for (operands, ids) |operand, *id| {
301 try self.ensureValue(operand);
302 id.* = operand.id;
303 }
304 return self.emit(try self.copyType(result_ty), .{
305 .custom_call = .{
306 .target = try arena.dupe(u8, target),
307 .version = version,
308 .operands = ids,
309 },
310 });
311 }
312
313 pub fn compare(self: *Builder, direction: CompareDirection, lhs: Value, rhs: Value) !Value {
314 try self.ensureValue(lhs);
315 try self.ensureValue(rhs);
316 try type_mod.sameType(lhs.ty, rhs.ty);
317 const pred_ty = try self.copyType(.{ .dtype = .i1, .dims = lhs.ty.dims });
318 return self.emit(pred_ty, .{
319 .compare = .{ .lhs = lhs.id, .rhs = rhs.id, .direction = direction },
320 });
321 }
322
323 pub fn alignedCompare(self: *Builder, direction: CompareDirection, lhs: Value, rhs: Value) !Value {
324 try self.ensureValue(lhs);
325 try self.ensureValue(rhs);
326 if (lhs.ty.dtype != rhs.ty.dtype) return error.DTypeMismatch;
327 const target = try type_mod.unionDims(self.arena.allocator(), lhs.ty.dims, rhs.ty.dims);
328 const left = try self.alignTo(lhs, target);
329 const right = try self.alignTo(rhs, target);
330 return self.compare(direction, left, right);
331 }
332
333 pub fn select(self: *Builder, pred: Value, on_true: Value, on_false: Value) !Value {
334 try self.ensureValue(pred);
335 try self.ensureValue(on_true);
336 try self.ensureValue(on_false);
337 const ty = try type_mod.select(self.arena.allocator(), pred.ty, on_true.ty, on_false.ty);
338 return self.emit(ty, .{
339 .select = .{ .pred = pred.id, .on_true = on_true.id, .on_false = on_false.id },
340 });
341 }
342
343 pub fn alignedSelect(self: *Builder, pred: Value, on_true: Value, on_false: Value) !Value {
344 try self.ensureValue(pred);
345 try self.ensureValue(on_true);
346 try self.ensureValue(on_false);
347 if (on_true.ty.dtype != on_false.ty.dtype) return error.DTypeMismatch;
348 const arena = self.arena.allocator();
349 var target = try type_mod.unionDims(arena, on_true.ty.dims, on_false.ty.dims);
350 if (pred.ty.rank() != 0) {
351 target = try type_mod.unionDims(arena, target, pred.ty.dims);
352 }
353 const chosen_true = try self.alignTo(on_true, target);
354 const chosen_false = try self.alignTo(on_false, target);
355 const chosen_pred = if (pred.ty.rank() == 0) pred else try self.alignTo(pred, target);
356 return self.select(chosen_pred, chosen_true, chosen_false);
357 }
358
359 pub fn broadcastAxes(self: *Builder, operand_value: Value, added: []const Dim) !Value {
360 try self.ensureValue(operand_value);
361 try type_mod.validateAuthoredDims(added);
362 const target = try type_mod.appendDims(self.arena.allocator(), operand_value.ty.dims, added);
363 return self.alignTo(operand_value, target);
364 }
365
366 pub fn renameAxis(self: *Builder, operand_value: Value, old_name: []const u8, new_name: []const u8) !Value {
367 try self.ensureValue(operand_value);
368 try type_mod.validateAuthoredName(new_name);
369 const new_dims = try type_mod.renamed(self.arena.allocator(), operand_value.ty.dims, old_name, new_name);
370 return self.reshapeTo(operand_value, new_dims);
371 }
372
373 pub fn splitAxis(self: *Builder, operand_value: Value, axis_name: []const u8, parts: []const Dim) !Value {
374 try self.ensureValue(operand_value);
375 try type_mod.validateAuthoredDims(parts);
376 const new_dims = try type_mod.splitDims(self.arena.allocator(), operand_value.ty.dims, axis_name, parts);
377 return self.reshapeTo(operand_value, new_dims);
378 }
379
380 pub fn mergeAxes(self: *Builder, operand_value: Value, names: []const []const u8, merged_name: []const u8) !Value {
381 try self.ensureValue(operand_value);
382 try type_mod.validateAuthoredName(merged_name);
383 const plan = try type_mod.mergeDims(self.arena.allocator(), operand_value.ty.dims, names, merged_name);
384 var current = operand_value;
385 if (plan.permutation) |permutation| {
386 current = try self.transposeBy(current, permutation);
387 }
388 return self.reshapeTo(current, plan.result);
389 }
390
391 pub fn alignTo(self: *Builder, operand_value: Value, target: []const Dim) !Value {
392 if (type_mod.sameDims(operand_value.ty.dims, target)) return operand_value;
393 const plan = try type_mod.alignment(self.arena.allocator(), operand_value.ty.dims, target);
394 var current = operand_value;
395 if (plan.permutation) |permutation| {
396 current = try self.transposeBy(current, permutation);
397 }
398 if (type_mod.sameDims(current.ty.dims, target)) return current;
399 return self.broadcastTo(current, target, plan.mapping.?);
400 }
401
402 pub fn broadcastOp(self: *Builder, operand_value: Value, result_dims: []const Dim) !Value {
403 try self.ensureValue(operand_value);
404 const ty = try Type.init(self.arena.allocator(), operand_value.ty.dtype, result_dims);
405 return self.emit(ty, .{
406 .broadcast = .{
407 .input = operand_value.id,
408 .sizes = try type_mod.extents(self.arena.allocator(), ty.dims),
409 },
410 });
411 }
412
413 pub fn broadcastTo(self: *Builder, operand_value: Value, result_dims: []const Dim, broadcast_dims: []const i64) !Value {
414 try self.ensureValue(operand_value);
415 const ty = try type_mod.broadcastInDim(self.arena.allocator(), operand_value.ty, result_dims, broadcast_dims);
416 return self.emit(ty, .{
417 .broadcast_in_dim = .{
418 .input = operand_value.id,
419 .broadcast_dims = try self.arena.allocator().dupe(i64, broadcast_dims),
420 },
421 });
422 }
423
424 pub fn reshapeTo(self: *Builder, operand_value: Value, new_dims: []const Dim) !Value {
425 try self.ensureValue(operand_value);
426 const ty = try type_mod.reshaped(self.arena.allocator(), operand_value.ty, new_dims);
427 return self.emit(ty, .{
428 .reshape = .{
429 .input = operand_value.id,
430 .new_shape = try type_mod.extents(self.arena.allocator(), ty.dims),
431 },
432 });
433 }
434
435 pub fn transposeBy(self: *Builder, operand_value: Value, permutation: []const i64) !Value {
436 try self.ensureValue(operand_value);
437 const result_dims = try type_mod.permuted(self.arena.allocator(), operand_value.ty.dims, permutation);
438 const ty = try Type.init(self.arena.allocator(), operand_value.ty.dtype, result_dims);
439 return self.emit(ty, .{
440 .transpose = .{
441 .input = operand_value.id,
442 .permutation = try self.arena.allocator().dupe(i64, permutation),
443 },
444 });
445 }
446
447 pub fn reduce(self: *Builder, operand_value: Value, init_value: Value, reducer: Reducer, dimensions: []const i64) !Value {
448 try self.ensureValue(operand_value);
449 try self.ensureValue(init_value);
450 if (operand_value.ty.dtype != init_value.ty.dtype) return error.DTypeMismatch;
451 if (init_value.ty.rank() != 0) return error.ReduceInitNotScalar;
452 const result_dims = try type_mod.removeAxes(self.arena.allocator(), operand_value.ty.dims, dimensions);
453 const ty = try Type.init(self.arena.allocator(), operand_value.ty.dtype, result_dims);
454 return self.emit(ty, .{
455 .reduce = .{
456 .input = operand_value.id,
457 .init = init_value.id,
458 .reducer = reducer,
459 .dimensions = try self.arena.allocator().dupe(i64, dimensions),
460 },
461 });
462 }
463
464 pub fn reduceNamed(self: *Builder, operand_value: Value, reducer: Reducer, names: []const []const u8) !Value {
465 try self.ensureValue(operand_value);
466 const indices = try type_mod.axisIndices(self.arena.allocator(), operand_value.ty.dims, names);
467 const init_value = try self.reducerInit(operand_value.ty.dtype, reducer);
468 return self.reduce(operand_value, init_value, reducer, indices);
469 }
470
471 pub fn reduceWith(self: *Builder, operand_value: Value, init_value: Value, reducer: Reducer, names: []const []const u8) !Value {
472 try self.ensureValue(operand_value);
473 const indices = try type_mod.axisIndices(self.arena.allocator(), operand_value.ty.dims, names);
474 return self.reduce(operand_value, init_value, reducer, indices);
475 }
476
477 pub fn gather(self: *Builder, operand_value: Value, indices_value: Value, comptime axis: anytype) !Value {
478 return self.gatherNamed(operand_value, indices_value, comptime nameOf(axis));
479 }
480
481 pub fn gatherNamed(self: *Builder, operand_value: Value, indices_value: Value, axis_name: []const u8) !Value {
482 try self.ensureValue(operand_value);
483 const axis = type_mod.findDim(operand_value.ty.dims, axis_name) orelse return error.AxisNotFound;
484 return self.gatherOp(operand_value, indices_value, @intCast(axis));
485 }
486
487 pub fn gatherOp(self: *Builder, operand_value: Value, indices_value: Value, axis: i64) !Value {
488 try self.ensureValue(operand_value);
489 try self.ensureValue(indices_value);
490 const ty = try type_mod.gather(self.arena.allocator(), operand_value.ty, indices_value.ty, axis);
491 return self.emit(ty, .{
492 .gather = .{
493 .input = operand_value.id,
494 .indices = indices_value.id,
495 .axis = axis,
496 },
497 });
498 }
499
500 pub fn scatterAdd(self: *Builder, input_value: Value, indices_value: Value, updates_value: Value, comptime axis: anytype) !Value {
501 return self.scatterAddNamed(input_value, indices_value, updates_value, comptime nameOf(axis));
502 }
503
504 pub fn scatterAddNamed(self: *Builder, input_value: Value, indices_value: Value, updates_value: Value, axis_name: []const u8) !Value {
505 try self.ensureValue(input_value);
506 const axis = type_mod.findDim(input_value.ty.dims, axis_name) orelse return error.AxisNotFound;
507 return self.scatterAddOp(input_value, indices_value, updates_value, @intCast(axis));
508 }
509
510 pub fn scatterAddOp(self: *Builder, input_value: Value, indices_value: Value, updates_value: Value, axis: i64) !Value {
511 try self.ensureValue(input_value);
512 try self.ensureValue(indices_value);
513 try self.ensureValue(updates_value);
514 const ty = try type_mod.scatterAdd(self.arena.allocator(), input_value.ty, indices_value.ty, updates_value.ty, axis);
515 return self.emit(ty, .{
516 .scatter_add = .{
517 .input = input_value.id,
518 .indices = indices_value.id,
519 .updates = updates_value.id,
520 .axis = axis,
521 },
522 });
523 }
524
525 pub fn sparseCrossEntropyLoss(self: *Builder, logits_value: Value, targets_value: Value, comptime axis: anytype) !Value {
526 return self.sparseCrossEntropyLossNamed(logits_value, targets_value, comptime nameOf(axis));
527 }
528
529 pub fn sparseCrossEntropyLossNamed(self: *Builder, logits_value: Value, targets_value: Value, axis_name: []const u8) !Value {
530 try self.ensureValue(logits_value);
531 const axis = type_mod.findDim(logits_value.ty.dims, axis_name) orelse return error.AxisNotFound;
532 return self.sparseCrossEntropyOp(logits_value, targets_value, @intCast(axis));
533 }
534
535 pub fn sparseCrossEntropyOp(self: *Builder, logits_value: Value, targets_value: Value, axis: i64) !Value {
536 try self.ensureValue(logits_value);
537 try self.ensureValue(targets_value);
538 const ty = try type_mod.sparseCrossEntropy(self.arena.allocator(), logits_value.ty, targets_value.ty, axis);
539 return self.emit(ty, .{
540 .sparse_cross_entropy = .{
541 .logits = logits_value.id,
542 .targets = targets_value.id,
543 .axis = axis,
544 },
545 });
546 }
547
548 pub fn meanNamed(self: *Builder, operand_value: Value, names: []const []const u8) !Value {
549 try self.ensureValue(operand_value);
550 const indices = try type_mod.axisIndices(self.arena.allocator(), operand_value.ty.dims, names);
551 const count = try type_mod.reducedExtentProduct(operand_value.ty.dims, indices);
552 const init_value = try self.reducerInit(operand_value.ty.dtype, .sum);
553 const summed = try self.reduce(operand_value, init_value, .sum, indices);
554 const divisor = try self.fullFloat(Type.scalar(operand_value.ty.dtype), @floatFromInt(count));
555 return self.alignedBinary(.div, summed, divisor);
556 }
557
558 fn reducerInit(self: *Builder, dtype: DType, reducer: Reducer) !Value {
559 const ty = Type.scalar(dtype);
560 return switch (reducer) {
561 .sum => self.zeros(ty),
562 .max => self.extremum(ty, .lowest),
563 .min => self.extremum(ty, .highest),
564 };
565 }
566
567 fn extremum(self: *Builder, ty: Type, comptime bound: enum { lowest, highest }) !Value {
568 return switch (ty.dtype) {
569 inline .f16, .f32, .f64 => |tag| blk: {
570 const Element = tag.ZigType();
571 const value: Element = if (bound == .lowest) -std.math.inf(Element) else std.math.inf(Element);
572 break :blk self.constantBytes(ty, std.mem.asBytes(&value));
573 },
574 .bf16 => blk: {
575 const Bf16 = DType.bf16.ZigType();
576 const value = Bf16.fromF32(if (bound == .lowest) -std.math.inf(f32) else std.math.inf(f32));
577 break :blk self.constantBytes(ty, std.mem.asBytes(&value));
578 },
579 inline .i32, .i64 => |tag| blk: {
580 const Element = tag.ZigType();
581 const value: Element = if (bound == .lowest) std.math.minInt(Element) else std.math.maxInt(Element);
582 break :blk self.constantBytes(ty, std.mem.asBytes(&value));
583 },
584 else => error.DTypeMismatch,
585 };
586 }
587
588 pub fn contract(self: *Builder, lhs: Value, rhs: Value, names: []const []const u8) !Value {
589 try self.ensureValue(lhs);
590 try self.ensureValue(rhs);
591 if (lhs.ty.dtype != rhs.ty.dtype) return error.ContractDTypeMismatch;
592 const arena = self.arena.allocator();
593 const plan = try type_mod.contraction(arena, lhs.ty.dims, rhs.ty.dims, names);
594
595 const batch_count = plan.lhs_batch.len;
596 const contract_count = plan.lhs_contract.len;
597 const lhs_free = lhs.ty.dims.len - batch_count - contract_count;
598
599 const lhs_order = try arena.alloc(Dim, lhs.ty.dims.len);
600 for (plan.lhs_batch, 0..) |axis, index| {
601 lhs_order[index] = lhs.ty.dims[@intCast(axis)];
602 }
603 fillFreeDims(lhs.ty.dims, plan.lhs_batch, plan.lhs_contract, lhs_order[batch_count .. batch_count + lhs_free]);
604 for (plan.lhs_contract, 0..) |axis, index| {
605 lhs_order[batch_count + lhs_free + index] = lhs.ty.dims[@intCast(axis)];
606 }
607
608 const rhs_order = try arena.alloc(Dim, rhs.ty.dims.len);
609 for (plan.rhs_batch, 0..) |axis, index| {
610 rhs_order[index] = rhs.ty.dims[@intCast(axis)];
611 }
612 for (plan.rhs_contract, 0..) |axis, index| {
613 rhs_order[batch_count + index] = rhs.ty.dims[@intCast(axis)];
614 }
615 fillFreeDims(rhs.ty.dims, plan.rhs_batch, plan.rhs_contract, rhs_order[batch_count + contract_count ..]);
616
617 const lhs_canonical = try self.alignTo(lhs, lhs_order);
618 const rhs_canonical = try self.alignTo(rhs, rhs_order);
619
620 const lhs_batch = try arena.alloc(i64, batch_count);
621 const rhs_batch = try arena.alloc(i64, batch_count);
622 for (0..batch_count) |index| {
623 lhs_batch[index] = @intCast(index);
624 rhs_batch[index] = @intCast(index);
625 }
626 const lhs_contract = try arena.alloc(i64, contract_count);
627 const rhs_contract = try arena.alloc(i64, contract_count);
628 for (0..contract_count) |index| {
629 lhs_contract[index] = @intCast(batch_count + lhs_free + index);
630 rhs_contract[index] = @intCast(batch_count + index);
631 }
632
633 const ty = try Type.init(arena, lhs.ty.dtype, plan.result);
634 return self.emit(ty, .{
635 .dot_general = .{
636 .lhs = lhs_canonical.id,
637 .rhs = rhs_canonical.id,
638 .lhs_contract = lhs_contract,
639 .rhs_contract = rhs_contract,
640 .lhs_batch = lhs_batch,
641 .rhs_batch = rhs_batch,
642 },
643 });
644 }
645
646 pub fn dotGeneralOp(
647 self: *Builder,
648 lhs: Value,
649 rhs: Value,
650 lhs_contract: []const i64,
651 rhs_contract: []const i64,
652 lhs_batch: []const i64,
653 rhs_batch: []const i64,
654 ) !Value {
655 try self.ensureValue(lhs);
656 try self.ensureValue(rhs);
657 if (lhs.ty.dtype != rhs.ty.dtype) return error.ContractDTypeMismatch;
658 const arena = self.arena.allocator();
659 const result_dims = try type_mod.dotGeneralDims(
660 arena,
661 lhs.ty.dims,
662 rhs.ty.dims,
663 lhs_contract,
664 rhs_contract,
665 lhs_batch,
666 rhs_batch,
667 );
668 const ty = try Type.init(arena, lhs.ty.dtype, result_dims);
669 return self.emit(ty, .{
670 .dot_general = .{
671 .lhs = lhs.id,
672 .rhs = rhs.id,
673 .lhs_contract = try arena.dupe(i64, lhs_contract),
674 .rhs_contract = try arena.dupe(i64, rhs_contract),
675 .lhs_batch = try arena.dupe(i64, lhs_batch),
676 .rhs_batch = try arena.dupe(i64, rhs_batch),
677 },
678 });
679 }
680
681 pub fn scan(self: *Builder, spec: anytype) !ScanResult(@TypeOf(spec.init)) {
682 const init_count = comptime scanValueCount(@TypeOf(spec.init));
683 var inits: [init_count]Value = undefined;
684 scanValues(spec.init, &inits);
685
686 var scope = try self.scanScope(spec.length, &inits);
687 errdefer scope.abort();
688
689 const carry = scanCarry(@TypeOf(spec.init), &scope);
690 const next = try callScanBody(spec.body, scope.body(), carry);
691 var next_values: [init_count]Value = undefined;
692 scanValues(next, &next_values);
693
694 const walked = try scope.finish(&next_values);
695 return scanResult(self, @TypeOf(spec.init), walked);
696 }
697
698 pub fn scanScope(self: *Builder, length: i64, inits: []const Value) !ScanScope {
699 if (length < 0) return error.ScanLengthNegative;
700 if (inits.len == 0) return error.ScanWithoutCarries;
701 if (inits.len > program_mod.max_scan_carries) return error.ScanArityUnsupported;
702 const init_ids = try self.arena.allocator().alloc(Id, inits.len);
703 for (inits, init_ids) |value, *slot| {
704 try self.ensureValue(value);
705 slot.* = value.id;
706 }
707 var child = try Builder.init(self.allocator, "scan_body");
708 errdefer child.deinit();
709 for (inits) |value| {
710 _ = try child.inputTyped(value.ty);
711 }
712 return .{
713 .parent = self,
714 .child = child,
715 .length = length,
716 .init_ids = init_ids,
717 };
718 }
719
720 pub fn projection(self: *Builder, source: Value, index: usize) !Value {
721 try self.ensureValue(source);
722 const source_op = self.operations.items[source.id.index];
723 const scan_op = switch (source_op.kind) {
724 .scan => |scan_op| scan_op,
725 else => return error.ProjectionSourceNotScan,
726 };
727 if (index == 0 or index >= scan_op.inits.len) return error.ProjectionIndexInvalid;
728 const ty = try self.copyType(scan_op.body.typeOf(scan_op.body.outputs[index]));
729 return self.emit(ty, .{ .projection = .{ .source = source.id, .index = index } });
730 }
731
732 pub fn emitScan(self: *Builder, length: i64, inits: []const Value, body: *const program_mod.Subgraph) !Value {
733 const init_ids = try self.arena.allocator().alloc(Id, inits.len);
734 for (inits, init_ids) |value, *slot| {
735 try self.ensureValue(value);
736 slot.* = value.id;
737 }
738 return self.emitScanIds(length, init_ids, body);
739 }
740
741 fn emitScanIds(self: *Builder, length: i64, inits: []const Id, body: *const program_mod.Subgraph) !Value {
742 if (length < 0) return error.ScanLengthNegative;
743 if (inits.len == 0) return error.ScanWithoutCarries;
744 if (inits.len > program_mod.max_scan_carries) return error.ScanArityUnsupported;
745 if (body.parameters.len != inits.len) return error.ScanArityMismatch;
746 if (body.outputs.len != inits.len) return error.ScanArityMismatch;
747 for (inits, 0..) |init_id, index| {
748 const init_ty = self.values.items[init_id.index];
749 try type_mod.sameType(init_ty, body.typeOf(body.parameters[index]));
750 try type_mod.sameType(init_ty, body.typeOf(body.outputs[index]));
751 }
752 const owned_body = try program_mod.cloneSubgraph(self.arena.allocator(), body);
753 const result_ty = try self.copyType(self.values.items[inits[0].index]);
754 return self.emit(result_ty, .{ .scan = .{
755 .length = length,
756 .inits = inits,
757 .body = owned_body,
758 } });
759 }
760
761 pub fn finish(self: *Builder, outputs: []const Value) !Program {
762 if (self.finished) return error.BuilderFinished;
763 const owned_outputs = try self.arena.allocator().alloc(Id, outputs.len);
764 for (outputs, 0..) |output, index| {
765 try self.ensureValue(output);
766 owned_outputs[index] = output.id;
767 }
768 self.finished = true;
769 return .{
770 .arena = self.arena,
771 .name = self.name,
772 .values = self.values.items,
773 .operations = self.operations.items,
774 .parameters = self.parameters_list.items,
775 .outputs = owned_outputs,
776 };
777 }
778
779 fn constantOwned(self: *Builder, ty: Type, payload: []const u8) !Value {
780 return self.emit(ty, .{ .constant = .{ .payload = payload } });
781 }
782
783 fn emit(self: *Builder, ty: Type, kind: program_mod.Kind) !Value {
784 const id = try self.append(ty, kind);
785 return self.makeValue(id);
786 }
787
788 fn append(self: *Builder, result: Type, kind: program_mod.Kind) !Id {
789 if (self.finished) return error.BuilderFinished;
790 const id = Id{ .index = @intCast(self.values.items.len) };
791 try self.values.append(self.arena.allocator(), result);
792 try self.operations.append(self.arena.allocator(), .{
793 .id = id,
794 .result = result,
795 .kind = kind,
796 });
797 return id;
798 }
799
800 fn makeValue(self: *Builder, id: Id) Value {
801 return .{
802 .builder = self,
803 .id = id,
804 .ty = self.values.items[id.index],
805 };
806 }
807
808 fn ensureValue(self: *Builder, candidate: Value) !void {
809 if (self.finished) return error.BuilderFinished;
810 if (candidate.builder != self) return error.CrossBuilderValue;
811 }
812
813 fn copyType(self: *Builder, ty: Type) !Type {
814 return Type.init(self.arena.allocator(), ty.dtype, ty.dims);
815 }
816 };
817
818 fn fillFreeDims(dims: []const Dim, batch: []const i64, contracted: []const i64, out: []Dim) void {
819 var index: usize = 0;
820 for (dims, 0..) |dim, position| {
821 if (containsIndex(batch, position) or containsIndex(contracted, position)) continue;
822 out[index] = dim;
823 index += 1;
824 }
825 }
826
827 fn containsIndex(axes: []const i64, position: usize) bool {
828 for (axes) |axis| {
829 if (axis == @as(i64, @intCast(position))) return true;
830 }
831 return false;
832 }
833
834 pub fn nameOf(comptime name_value: anytype) []const u8 {
835 return switch (@typeInfo(@TypeOf(name_value))) {
836 .enum_literal => @tagName(name_value),
837 .pointer, .array => name_value[0..],
838 else => @compileError("tensor axis names must be enum literals or string literals"),
839 };
840 }
841
842 fn ScanResult(comptime Init: type) type {
843 if (comptime Init == Value) return Value;
844 return switch (@typeInfo(Init)) {
845 .array => |array| blk: {
846 if (array.child != Value) @compileError("tensor scan array init values must be tensor Values");
847 break :blk [array.len]Value;
848 },
849 .@"struct" => |info| blk: {
850 if (info.field_names.len == 0) @compileError("tensor scan init struct must contain at least one carry");
851 inline for (info.field_types) |field_type| {
852 if (field_type != Value) @compileError("tensor scan struct init fields must be tensor Values");
853 }
854 break :blk Init;
855 },
856 else => @compileError("tensor scan init must be a Value, array of Values, tuple of Values, or struct of Values"),
857 };
858 }
859
860 fn scanValueCount(comptime Init: type) usize {
861 if (comptime Init == Value) return 1;
862 return switch (@typeInfo(Init)) {
863 .array => |array| blk: {
864 if (array.child != Value) @compileError("tensor scan array init values must be tensor Values");
865 break :blk array.len;
866 },
867 .@"struct" => |info| blk: {
868 if (info.field_names.len == 0) @compileError("tensor scan init struct must contain at least one carry");
869 inline for (info.field_types) |field_type| {
870 if (field_type != Value) @compileError("tensor scan struct init fields must be tensor Values");
871 }
872 break :blk info.field_names.len;
873 },
874 else => @compileError("tensor scan init must be a Value, array of Values, tuple of Values, or struct of Values"),
875 };
876 }
877
878 fn scanValues(source: anytype, values: *[scanValueCount(@TypeOf(source))]Value) void {
879 const Source = @TypeOf(source);
880 if (comptime Source == Value) {
881 values[0] = source;
882 return;
883 }
884 switch (@typeInfo(Source)) {
885 .array => for (source, 0..) |value, index| {
886 values[index] = value;
887 },
888 .@"struct" => |info| inline for (info.field_names, 0..) |field_name, index| {
889 values[index] = @field(source, field_name);
890 },
891 else => unreachable,
892 }
893 }
894
895 fn scanCarry(comptime Init: type, scope: *ScanScope) ScanResult(Init) {
896 if (comptime Init == Value) return scope.carry(0);
897 return switch (@typeInfo(Init)) {
898 .array => |array| blk: {
899 var result: [array.len]Value = undefined;
900 inline for (0..array.len) |index| {
901 result[index] = scope.carry(index);
902 }
903 break :blk result;
904 },
905 .@"struct" => |info| blk: {
906 var result: Init = undefined;
907 inline for (info.field_names, 0..) |field_name, index| {
908 @field(result, field_name) = scope.carry(index);
909 }
910 break :blk result;
911 },
912 else => unreachable,
913 };
914 }
915
916 fn scanResult(builder: *Builder, comptime Init: type, walked: Value) !ScanResult(Init) {
917 if (comptime Init == Value) return walked;
918 return switch (@typeInfo(Init)) {
919 .array => |array| blk: {
920 var result: [array.len]Value = undefined;
921 inline for (0..array.len) |index| {
922 result[index] = if (index == 0) walked else try builder.projection(walked, index);
923 }
924 break :blk result;
925 },
926 .@"struct" => |info| blk: {
927 var result: Init = undefined;
928 inline for (info.field_names, 0..) |field_name, index| {
929 @field(result, field_name) = if (index == 0) walked else try builder.projection(walked, index);
930 }
931 break :blk result;
932 },
933 else => unreachable,
934 };
935 }
936
937 fn callScanBody(comptime body: anytype, builder: *Builder, carry: anytype) !ScanResult(@TypeOf(carry)) {
938 if (comptime @typeInfo(@TypeOf(body)) != .@"fn") @compileError("tensor scan body must be a function value");
939 return body(builder, carry);
940 }
941
942 test "indexing operations trace with named axis shapes" {
943 var builder = try Builder.init(std.testing.allocator, "trace_indexing");
944 defer builder.deinit();
945
946 const table = try builder.input(.f32, .{ .vocab = 32, .channel = 8 });
947 const ids = try builder.input(.i32, .{ .token = 5 });
948 const gathered = try table.gather(ids, .vocab);
949 try std.testing.expectEqual(@as(usize, 2), gathered.ty.rank());
950 try std.testing.expectEqualStrings("token", gathered.ty.dims[0].name);
951 try std.testing.expectEqualStrings("channel", gathered.ty.dims[1].name);
952
953 const zero_table = try builder.full(.f32, .{ .vocab = 32, .channel = 8 }, 0.0);
954 const accumulated = try zero_table.scatterAdd(ids, gathered, .vocab);
955 try std.testing.expect(type_mod.sameDims(table.ty.dims, accumulated.ty.dims));
956
957 var program = try builder.finish(&.{accumulated});
958 defer program.deinit();
959
960 try std.testing.expect(program.operation(gathered.id).kind == .gather);
961 try std.testing.expect(program.operation(accumulated.id).kind == .scatter_add);
962 }
963
964 test "scan scope traces a carry loop into one scan operation" {
965 var builder = try Builder.init(std.testing.allocator, "scan_trace");
966 defer builder.deinit();
967
968 const x0 = try builder.input(.f32, .{ .lane = 4 });
969 const acc0 = try builder.full(.f32, .{ .lane = 4 }, 0.0);
970
971 var scope = try builder.scanScope(5, &.{ x0, acc0 });
972 errdefer scope.abort();
973 const half = try scope.body().scalar(.f32, 0.5);
974 const next_x = try (try scope.carry(0).mul(scope.carry(0))).mul(half);
975 const next_acc = try scope.carry(1).add(next_x);
976 const walked = try scope.finish(&.{ next_x, next_acc });
977 const final_acc = try builder.projection(walked, 1);
978
979 var program = try builder.finish(&.{ walked, final_acc });
980 defer program.deinit();
981
982 const scan_op = program.operation(walked.id);
983 try std.testing.expectEqual(@as(i64, 5), scan_op.kind.scan.length);
984 try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.inits.len);
985 try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.body.parameters.len);
986 try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.body.outputs.len);
987 try std.testing.expectEqualStrings("lane", program.typeOf(walked.id).dims[0].name);
988 try std.testing.expectEqual(@as(i64, 4), program.typeOf(final_acc.id).dims[0].extent);
989
990 const projection_op = program.operation(final_acc.id);
991 try std.testing.expectEqual(walked.id, projection_op.kind.projection.source);
992 try std.testing.expectEqual(@as(usize, 1), projection_op.kind.projection.index);
993 }
994
995 fn namedScanStep(scan_builder: *Builder, carry: anytype) !@TypeOf(carry) {
996 const half = try scan_builder.scalar(.f32, 0.5);
997 const next_x = try (try carry.x.mul(carry.x)).mul(half);
998 return .{
999 .x = next_x,
1000 .acc = try carry.acc.add(next_x),
1001 };
1002 }
1003
1004 test "scan definition returns named carries" {
1005 var builder = try Builder.init(std.testing.allocator, "scan_named");
1006 defer builder.deinit();
1007
1008 const x0 = try builder.input(.f32, .{ .lane = 4 });
1009 const acc0 = try builder.full(.f32, .{ .lane = 4 }, 0.0);
1010
1011 const walked = try builder.scan(.{
1012 .length = 5,
1013 .init = .{ .x = x0, .acc = acc0 },
1014 .body = namedScanStep,
1015 });
1016
1017 var program = try builder.finish(&.{ walked.x, walked.acc });
1018 defer program.deinit();
1019
1020 const scan_op = program.operation(walked.x.id);
1021 try std.testing.expectEqual(@as(i64, 5), scan_op.kind.scan.length);
1022 try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.inits.len);
1023 try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.body.outputs.len);
1024
1025 const projection_op = program.operation(walked.acc.id);
1026 try std.testing.expectEqual(walked.x.id, projection_op.kind.projection.source);
1027 try std.testing.expectEqual(@as(usize, 1), projection_op.kind.projection.index);
1028 }
1029
1030 fn singleScanStep(_: *Builder, carry: Value) !Value {
1031 return carry.add(carry);
1032 }
1033
1034 test "scan definition returns a single carry value" {
1035 var builder = try Builder.init(std.testing.allocator, "scan_single");
1036 defer builder.deinit();
1037
1038 const x0 = try builder.input(.f32, .{ .lane = 4 });
1039
1040 const walked = try builder.scan(.{
1041 .length = 3,
1042 .init = x0,
1043 .body = singleScanStep,
1044 });
1045
1046 var program = try builder.finish(&.{walked});
1047 defer program.deinit();
1048
1049 const scan_op = program.operation(walked.id);
1050 try std.testing.expectEqual(@as(i64, 3), scan_op.kind.scan.length);
1051 try std.testing.expectEqual(@as(usize, 1), scan_op.kind.scan.inits.len);
1052 try std.testing.expectEqual(@as(usize, 1), scan_op.kind.scan.body.outputs.len);
1053 }
1054
1055 fn arrayScanStep(_: *Builder, carry: [2]Value) ![2]Value {
1056 const next = try carry[0].add(carry[0]);
1057 return .{ next, try carry[1].add(next) };
1058 }
1059
1060 fn tupleScanStep(_: *Builder, carry: anytype) !@TypeOf(carry) {
1061 const next = try carry[0].add(carry[0]);
1062 return .{ next, try carry[1].add(next) };
1063 }
1064
1065 test "scan definition returns array and tuple carries" {
1066 var array_builder = try Builder.init(std.testing.allocator, "scan_array");
1067 defer array_builder.deinit();
1068
1069 const array_x0 = try array_builder.input(.f32, .{ .lane = 2 });
1070 const array_acc0 = try array_builder.full(.f32, .{ .lane = 2 }, 0.0);
1071 const array_walked = try array_builder.scan(.{
1072 .length = 2,
1073 .init = [2]Value{ array_x0, array_acc0 },
1074 .body = arrayScanStep,
1075 });
1076 var array_program = try array_builder.finish(&.{ array_walked[0], array_walked[1] });
1077 defer array_program.deinit();
1078
1079 const array_projection = array_program.operation(array_walked[1].id);
1080 try std.testing.expectEqual(array_walked[0].id, array_projection.kind.projection.source);
1081 try std.testing.expectEqual(@as(usize, 1), array_projection.kind.projection.index);
1082
1083 var tuple_builder = try Builder.init(std.testing.allocator, "scan_tuple");
1084 defer tuple_builder.deinit();
1085
1086 const tuple_x0 = try tuple_builder.input(.f32, .{ .lane = 2 });
1087 const tuple_acc0 = try tuple_builder.full(.f32, .{ .lane = 2 }, 0.0);
1088 const tuple_walked = try tuple_builder.scan(.{
1089 .length = 2,
1090 .init = .{ tuple_x0, tuple_acc0 },
1091 .body = tupleScanStep,
1092 });
1093 var tuple_program = try tuple_builder.finish(&.{ tuple_walked[0], tuple_walked[1] });
1094 defer tuple_program.deinit();
1095
1096 const tuple_projection = tuple_program.operation(tuple_walked[1].id);
1097 try std.testing.expectEqual(tuple_walked[0].id, tuple_projection.kind.projection.source);
1098 try std.testing.expectEqual(@as(usize, 1), tuple_projection.kind.projection.index);
1099 }
1100
1101 test "scan scope rejects misuse" {
1102 var builder = try Builder.init(std.testing.allocator, "scan_misuse");
1103 defer builder.deinit();
1104
1105 const x0 = try builder.input(.f32, .{ .lane = 2 });
1106
1107 try std.testing.expectError(error.ScanLengthNegative, builder.scanScope(-1, &.{x0}));
1108 try std.testing.expectError(error.ScanWithoutCarries, builder.scanScope(3, &.{}));
1109
1110 var scope = try builder.scanScope(3, &.{x0});
1111 defer scope.abort();
1112 const stranger = try builder.full(.f32, .{ .lane = 2 }, 1.0);
1113 try std.testing.expectError(error.CrossBuilderValue, scope.finish(&.{stranger}));
1114
1115 var second = try builder.scanScope(3, &.{x0});
1116 defer second.abort();
1117 const wide = try second.body().full(.f32, .{ .lane = 3 }, 0.0);
1118 try std.testing.expectError(error.ShapeMismatch, second.finish(&.{wide}));
1119
1120 var third = try builder.scanScope(3, &.{x0});
1121 defer third.abort();
1122 try std.testing.expectError(error.ScanArityMismatch, third.finish(&.{ third.carry(0), third.carry(0) }));
1123
1124 const not_scan = try builder.full(.f32, .{ .lane = 2 }, 2.0);
1125 try std.testing.expectError(error.ProjectionSourceNotScan, builder.projection(not_scan, 1));
1126
1127 var fourth = try builder.scanScope(2, &.{x0});
1128 const bumped = try fourth.carry(0).add(try fourth.body().full(.f32, .{ .lane = 2 }, 1.0));
1129 const walked = try fourth.finish(&.{bumped});
1130 try std.testing.expectError(error.ProjectionIndexInvalid, builder.projection(walked, 0));
1131 try std.testing.expectError(error.ProjectionIndexInvalid, builder.projection(walked, 1));
1132 }
1133
1134 test "scan operations replay through the structural graph copy" {
1135 var builder = try Builder.init(std.testing.allocator, "scan_replay_source");
1136 defer builder.deinit();
1137
1138 const x0 = try builder.input(.f32, .{ .lane = 3 });
1139 const gain = try builder.input(.f32, .{ .lane = 3 });
1140
1141 var scope = try builder.scanScope(4, &.{ x0, gain });
1142 const next = try scope.carry(0).mul(scope.carry(1));
1143 const walked = try scope.finish(&.{ next, scope.carry(1) });
1144 const kept_gain = try builder.projection(walked, 1);
1145 var source = try builder.finish(&.{ walked, kept_gain });
1146 defer source.deinit();
1147
1148 var replay_builder = try Builder.init(std.testing.allocator, "scan_replay_source");
1149 errdefer replay_builder.deinit();
1150 var values = std.ArrayListUnmanaged(Value).empty;
1151 defer values.deinit(std.testing.allocator);
1152 for (source.operations) |*op| {
1153 var buffer: [program_mod.max_operation_operands]Value = undefined;
1154 var count: usize = 0;
1155 switch (op.kind) {
1156 .parameter => {},
1157 .scan => |scan| {
1158 for (scan.inits) |init_id| {
1159 buffer[count] = values.items[init_id.index];
1160 count += 1;
1161 }
1162 },
1163 .projection => |projection_op| {
1164 buffer[0] = values.items[projection_op.source.index];
1165 count = 1;
1166 },
1167 else => unreachable,
1168 }
1169 try values.append(std.testing.allocator, try replay_builder.operation(op, buffer[0..count]));
1170 }
1171 var outputs: [2]Value = .{ values.items[source.outputs[0].index], values.items[source.outputs[1].index] };
1172 var replayed = try replay_builder.finish(&outputs);
1173 defer replayed.deinit();
1174
1175 try std.testing.expectEqual(source.fingerprint(), replayed.fingerprint());
1176 }
1177
1178 test "scan fingerprints see body changes" {
1179 var first_builder = try Builder.init(std.testing.allocator, "scan_fp");
1180 defer first_builder.deinit();
1181 const first_x = try first_builder.input(.f32, .{ .lane = 2 });
1182 var first_scope = try first_builder.scanScope(3, &.{first_x});
1183 const first_next = try first_scope.carry(0).add(try first_scope.body().full(.f32, .{ .lane = 2 }, 1.0));
1184 const first_out = try first_scope.finish(&.{first_next});
1185 var first = try first_builder.finish(&.{first_out});
1186 defer first.deinit();
1187
1188 var second_builder = try Builder.init(std.testing.allocator, "scan_fp");
1189 defer second_builder.deinit();
1190 const second_x = try second_builder.input(.f32, .{ .lane = 2 });
1191 var second_scope = try second_builder.scanScope(3, &.{second_x});
1192 const second_next = try second_scope.carry(0).add(try second_scope.body().full(.f32, .{ .lane = 2 }, 2.0));
1193 const second_out = try second_scope.finish(&.{second_next});
1194 var second = try second_builder.finish(&.{second_out});
1195 defer second.deinit();
1196
1197 try std.testing.expect(first.fingerprint() != second.fingerprint());
1198 }
1199
1200 test "aligned binary broadcasts missing axes by name" {
1201 var builder = try Builder.init(std.testing.allocator, "align_binary");
1202 defer builder.deinit();
1203
1204 const xs = try builder.input(.f32, .{ .point = 4 });
1205 const noise = try builder.input(.f32, .{ .point = 4, .sample = 8 });
1206 const sum = try xs.add(noise);
1207
1208 try std.testing.expectEqual(@as(usize, 2), sum.ty.rank());
1209 try std.testing.expectEqualStrings("point", sum.ty.dims[0].name);
1210 try std.testing.expectEqualStrings("sample", sum.ty.dims[1].name);
1211
1212 const flipped = try noise.add(xs);
1213 try std.testing.expectEqualStrings("point", flipped.ty.dims[0].name);
1214
1215 const scalar_gain = try builder.scalar(.f32, 2.0);
1216 const scaled = try noise.mul(scalar_gain);
1217 try std.testing.expect(type_mod.sameDims(noise.ty.dims, scaled.ty.dims));
1218
1219 const wrong = try builder.input(.f32, .{ .point = 5 });
1220 try std.testing.expectError(error.AxisExtentMismatch, xs.add(wrong));
1221 }
1222
1223 test "aligned binary transposes shared axes into agreement" {
1224 var builder = try Builder.init(std.testing.allocator, "align_transpose");
1225 defer builder.deinit();
1226
1227 const row_major = try builder.input(.f32, .{ .row = 2, .col = 3 });
1228 const col_major = try builder.input(.f32, .{ .col = 3, .row = 2 });
1229 const sum = try row_major.add(col_major);
1230
1231 try std.testing.expectEqualStrings("row", sum.ty.dims[0].name);
1232 try std.testing.expectEqualStrings("col", sum.ty.dims[1].name);
1233
1234 const transpose_op = builder.operations.items[sum.id.index - 1];
1235 try std.testing.expect(transpose_op.kind == .transpose);
1236 }
1237
1238 test "named reductions use canonical inits" {
1239 var builder = try Builder.init(std.testing.allocator, "named_reduce");
1240 defer builder.deinit();
1241
1242 const x = try builder.input(.f32, .{ .point = 4, .sample = 8 });
1243
1244 const summed = try x.sum(.sample);
1245 try std.testing.expectEqual(@as(usize, 1), summed.ty.rank());
1246 try std.testing.expectEqualStrings("point", summed.ty.dims[0].name);
1247
1248 const biggest = try x.max(.{ .point, .sample });
1249 try std.testing.expectEqual(@as(usize, 0), biggest.ty.rank());
1250
1251 const smallest = try x.min(.point);
1252 try std.testing.expectEqualStrings("sample", smallest.ty.dims[0].name);
1253
1254 const averaged = try x.mean(.sample);
1255 try std.testing.expectEqualStrings("point", averaged.ty.dims[0].name);
1256
1257 try std.testing.expectError(error.AxisNotFound, x.sum(.missing));
1258 }
1259
1260 test "contract derives matmul attention and batched shapes" {
1261 var builder = try Builder.init(std.testing.allocator, "named_contract");
1262 defer builder.deinit();
1263
1264 const q = try builder.input(.f32, .{ .pos = 8, .head = 16 });
1265 const k = try builder.input(.f32, .{ .ctx = 6, .head = 16 });
1266 const scores = try q.contract(k, .head);
1267 try std.testing.expectEqualStrings("pos", scores.ty.dims[0].name);
1268 try std.testing.expectEqualStrings("ctx", scores.ty.dims[1].name);
1269
1270 const v = try builder.input(.f32, .{ .ctx = 6, .val = 32 });
1271 const out = try scores.contract(v, .ctx);
1272 try std.testing.expectEqualStrings("pos", out.ty.dims[0].name);
1273 try std.testing.expectEqualStrings("val", out.ty.dims[1].name);
1274
1275 const lhs = try builder.input(.f32, .{ .walk = 5, .m = 2, .k = 3 });
1276 const rhs = try builder.input(.f32, .{ .k = 3, .walk = 5, .n = 4 });
1277 const batched = try lhs.contract(rhs, .k);
1278 try std.testing.expectEqualStrings("walk", batched.ty.dims[0].name);
1279 try std.testing.expectEqualStrings("m", batched.ty.dims[1].name);
1280 try std.testing.expectEqualStrings("n", batched.ty.dims[2].name);
1281 }
1282
1283 test "structure ops rename split merge and broadcast by name" {
1284 var builder = try Builder.init(std.testing.allocator, "named_structure");
1285 defer builder.deinit();
1286
1287 const x = try builder.input(.f32, .{ .pixels = 12 });
1288
1289 const image = try x.split(.pixels, .{ .row = 3, .col = 4 });
1290 try std.testing.expectEqualStrings("row", image.ty.dims[0].name);
1291 try std.testing.expectEqualStrings("col", image.ty.dims[1].name);
1292
1293 const relabeled = try image.rename(.row, .line);
1294 try std.testing.expectEqualStrings("line", relabeled.ty.dims[0].name);
1295
1296 const flat = try relabeled.merge(.{ .line, .col }, .pixels);
1297 try std.testing.expectEqual(@as(usize, 1), flat.ty.rank());
1298 try std.testing.expectEqual(@as(i64, 12), flat.ty.dims[0].extent);
1299
1300 const widened = try flat.broadcast(.{ .sample = 2 });
1301 try std.testing.expectEqualStrings("pixels", widened.ty.dims[0].name);
1302 try std.testing.expectEqualStrings("sample", widened.ty.dims[1].name);
1303
1304 try std.testing.expectError(error.AxisNameReserved, flat.broadcast(.{ .@"#batch" = 2 }));
1305 }
1306
1307 test "aligned select broadcasts predicate and branches" {
1308 var builder = try Builder.init(std.testing.allocator, "named_select");
1309 defer builder.deinit();
1310
1311 const pred = try builder.input(.i1, .{ .point = 4 });
1312 const wide = try builder.input(.f32, .{ .point = 4, .sample = 8 });
1313 const fallback = try builder.scalar(.f32, 0.0);
1314
1315 const chosen = try pred.select(wide, fallback);
1316 try std.testing.expectEqual(@as(usize, 2), chosen.ty.rank());
1317 try std.testing.expectEqualStrings("point", chosen.ty.dims[0].name);
1318
1319 const scalar_pred = try builder.scalar(.i1, true);
1320 const scalar_chosen = try scalar_pred.select(wide, fallback);
1321 try std.testing.expectEqual(@as(usize, 2), scalar_chosen.ty.rank());
1322 }