lib/accy/src/tensor/autodiff.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const emit = @import("emit.zig");
3 const hook = @import("hook.zig");
4 const interpret = @import("interpret/root.zig");
5 const program_mod = @import("program.zig");
6 const trace = @import("trace/root.zig");
7 const transform = @import("transform.zig");
8 const unroll = @import("unroll.zig");
9 const types = @import("type/root.zig");
10
11 pub const LinearizeOptions = struct {
12 wrt: []const usize,
13 };
14
15 pub const JvpOptions = LinearizeOptions;
16
17 pub const Dual = struct {
18 primal: trace.Value,
19 tangent: trace.Value,
20 };
21
22 const ResultMode = enum {
23 linearization,
24 downstream,
25 };
26
27 pub const Linearization = struct {
28 program: program_mod.Program,
29 differentiated_parameters: []const usize,
30 primal_parameter_count: usize,
31 tangent_parameter_count: usize,
32 primal_output_count: usize,
33 tangent_output_count: usize,
34
35 pub fn deinit(self: *Linearization) void {
36 self.program.deinit();
37 }
38
39 pub fn primalOutputs(self: *const Linearization) []const program_mod.Id {
40 return self.program.outputs[0..self.primal_output_count];
41 }
42
43 pub fn tangentOutputs(self: *const Linearization) []const program_mod.Id {
44 return self.program.outputs[self.primal_output_count .. self.primal_output_count + self.tangent_output_count];
45 }
46 };
47
48 pub fn semantics(
49 source: *const program_mod.Program,
50 next: anytype,
51 options: LinearizeOptions,
52 ) Semantics(@TypeOf(next), .linearization) {
53 return .{
54 .source = source,
55 .next = next,
56 .options = options,
57 };
58 }
59
60 pub fn jvpSemantics(
61 source: *const program_mod.Program,
62 next: anytype,
63 options: JvpOptions,
64 ) Semantics(@TypeOf(next), .downstream) {
65 return .{
66 .source = source,
67 .next = next,
68 .options = options,
69 };
70 }
71
72 pub fn jvpSemanticsWith(
73 source: *const program_mod.Program,
74 next: anytype,
75 options: JvpOptions,
76 hooks: anytype,
77 ) @TypeOf(hook.attach("jvp", hooks, jvpSemantics(source, hook.attach("linearize", hooks, next), options))) {
78 const generated = hook.attach("linearize", hooks, next);
79 const linear = jvpSemantics(source, generated, options);
80 return hook.attach("jvp", hooks, linear);
81 }
82
83 pub fn Semantics(comptime Next: type, comptime result_mode: ResultMode) type {
84 return struct {
85 source: *const program_mod.Program,
86 next: Next,
87 options: LinearizeOptions,
88 parameter_values: ?[]Value = null,
89
90 pub const Value: type = Dual;
91 pub const Result: type = if (result_mode == .linearization) Linearization else Next.Result;
92
93 pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {
94 var buffer: [program_mod.max_operation_operands]Value = undefined;
95 return self.bind(step.op, interpret.arguments(Value, step.op, step.values, &buffer));
96 }
97
98 pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const Value) !Value {
99 return switch (op.kind) {
100 .parameter => |parameter_info| self.parameterValue(parameter_info),
101 .constant => self.constantValue(op),
102 .iota => self.iotaValue(op),
103 .unary => |unary_info| self.unaryValue(op, unary_info, args[0]),
104 .binary => |binary_info| self.binaryValue(op, binary_info, args[0], args[1]),
105 .broadcast => self.broadcastValue(op, args[0]),
106 .broadcast_in_dim => |broadcast_info| self.broadcastInDimValue(op, broadcast_info, args[0]),
107 .reshape => self.reshapeValue(op, args[0]),
108 .transpose => |transpose_info| self.transposeValue(op, transpose_info, args[0]),
109 .reduce => |reduce_info| self.reduceValue(op, reduce_info, args[0], args[1]),
110 .gather => |gather_info| self.gatherValue(op, gather_info, args[0], args[1]),
111 .scatter_add => |scatter_add| self.scatterAddValue(op, scatter_add, args[0], args[1], args[2]),
112 .sparse_cross_entropy => |sparse_cross_entropy| self.sparseCrossEntropyValue(op, sparse_cross_entropy, args[0], args[1]),
113 .dot_general => |dot_info| self.dotGeneralValue(op, dot_info, args[0], args[1]),
114 .compare => self.compareValue(op, args[0], args[1]),
115 .select => self.selectValue(op, args[0], args[1], args[2]),
116 .custom_call => error.CustomCallRequiresJvpContract,
117 .scan, .projection => error.ScanLinearizationUnsupported,
118 };
119 }
120
121 pub fn finish(self: *@This(), outputs: []const Value) !Result {
122 const values = try self.builderHandle().arena.allocator().alloc(trace.Value, outputs.len * 2);
123 for (outputs, 0..) |output, index| {
124 values[index] = output.primal;
125 values[index + outputs.len] = output.tangent;
126 }
127
128 if (comptime result_mode == .downstream) return self.next.finish(values);
129
130 var program = try self.next.finish(values);
131 errdefer program.deinit();
132
133 const differentiated_parameters = try differentiatedParameters(&program, self.source, self.options);
134 const tangent_parameter_count = differentiated_parameters.len;
135
136 return .{
137 .program = program,
138 .differentiated_parameters = differentiated_parameters,
139 .primal_parameter_count = self.source.parameters.len,
140 .tangent_parameter_count = tangent_parameter_count,
141 .primal_output_count = outputs.len,
142 .tangent_output_count = outputs.len,
143 };
144 }
145
146 pub fn builderHandle(self: *@This()) *trace.Builder {
147 return self.next.builderHandle();
148 }
149
150 fn parameterValue(self: *@This(), parameter_info: program_mod.Parameter) !Value {
151 if (self.parameter_values == null) try self.prepareParameters();
152 return self.parameter_values.?[parameter_info.index];
153 }
154
155 fn prepareParameters(self: *@This()) !void {
156 const values = try self.builderHandle().arena.allocator().alloc(Value, self.source.parameters.len);
157 for (self.source.parameters) |id| {
158 const op = &self.source.operations[id.index];
159 const parameter_info = op.kind.parameter;
160 values[parameter_info.index].primal = try self.next.bind(op, &.{});
161 }
162 for (self.source.parameters) |id| {
163 const op = &self.source.operations[id.index];
164 const parameter_info = op.kind.parameter;
165 values[parameter_info.index].tangent = if (isWrt(self.options.wrt, parameter_info.index)) tangent: {
166 if (!op.result.dtype.isFloat()) return error.NonDifferentiableParameter;
167 break :tangent try self.builderHandle().inputTyped(op.result);
168 } else try self.emitZeros(op.result);
169 }
170 self.parameter_values = values;
171 }
172
173 fn compareValue(self: *@This(), op: *const program_mod.Operation, lhs: Value, rhs: Value) !Value {
174 return .{
175 .primal = try self.next.bind(op, &.{ lhs.primal, rhs.primal }),
176 .tangent = try self.emitZeros(op.result),
177 };
178 }
179
180 fn selectValue(self: *@This(), op: *const program_mod.Operation, pred: Value, on_true: Value, on_false: Value) !Value {
181 return .{
182 .primal = try self.next.bind(op, &.{ pred.primal, on_true.primal, on_false.primal }),
183 .tangent = try emit.select(self, pred.primal, on_true.tangent, on_false.tangent),
184 };
185 }
186
187 fn constantValue(self: *@This(), op: *const program_mod.Operation) !Value {
188 return .{
189 .primal = try self.next.bind(op, &.{}),
190 .tangent = try self.emitZeros(op.result),
191 };
192 }
193
194 fn iotaValue(self: *@This(), op: *const program_mod.Operation) !Value {
195 return .{
196 .primal = try self.next.bind(op, &.{}),
197 .tangent = try self.emitZeros(op.result),
198 };
199 }
200
201 fn unaryValue(self: *@This(), op: *const program_mod.Operation, unary_info: program_mod.UnaryOp, input: Value) !Value {
202 const primal = try self.next.bind(op, &.{input.primal});
203 return .{
204 .primal = primal,
205 .tangent = try unaryJvp(self, unary_info.op, input.primal, primal, input.tangent),
206 };
207 }
208
209 fn binaryValue(self: *@This(), op: *const program_mod.Operation, binary_info: program_mod.BinaryOp, lhs: Value, rhs: Value) !Value {
210 const primal = try self.next.bind(op, &.{ lhs.primal, rhs.primal });
211 return .{
212 .primal = primal,
213 .tangent = try binaryJvp(self, binary_info.op, lhs.primal, rhs.primal, lhs.tangent, rhs.tangent),
214 };
215 }
216
217 fn broadcastValue(self: *@This(), op: *const program_mod.Operation, input: Value) !Value {
218 return .{
219 .primal = try self.next.bind(op, &.{input.primal}),
220 .tangent = try self.emitBroadcast(input.tangent, op.result.dims),
221 };
222 }
223
224 fn broadcastInDimValue(self: *@This(), op: *const program_mod.Operation, broadcast_info: program_mod.BroadcastInDim, input: Value) !Value {
225 return .{
226 .primal = try self.next.bind(op, &.{input.primal}),
227 .tangent = try self.emitBroadcastInDim(input.tangent, op.result.dims, broadcast_info.broadcast_dims),
228 };
229 }
230
231 fn reshapeValue(self: *@This(), op: *const program_mod.Operation, input: Value) !Value {
232 return .{
233 .primal = try self.next.bind(op, &.{input.primal}),
234 .tangent = try self.emitReshape(input.tangent, op.result.dims),
235 };
236 }
237
238 fn transposeValue(self: *@This(), op: *const program_mod.Operation, transpose_info: program_mod.Transpose, input: Value) !Value {
239 return .{
240 .primal = try self.next.bind(op, &.{input.primal}),
241 .tangent = try self.emitTranspose(input.tangent, transpose_info.permutation),
242 };
243 }
244
245 fn reduceValue(self: *@This(), op: *const program_mod.Operation, reduce_info: program_mod.Reduce, input: Value, init: Value) !Value {
246 const primal = try self.next.bind(op, &.{ input.primal, init.primal });
247 const tangent = switch (reduce_info.reducer) {
248 .sum => try self.emitReduce(input.tangent, init.tangent, .sum, reduce_info.dimensions),
249 .max, .min => blk: {
250 const input_active = !input.tangent.isStructuralZero();
251 const init_active = !init.tangent.isStructuralZero();
252 if (!input_active and !init_active) break :blk try self.emitZeros(primal.ty);
253
254 var input_term: ?trace.Value = null;
255 if (input_active) {
256 const allocator = self.builderHandle().arena.allocator();
257 const input_dims = input.primal.ty.dims;
258 const kept = try allocator.alloc(i64, input_dims.len - reduce_info.dimensions.len);
259 var kept_index: usize = 0;
260 for (0..input_dims.len) |axis| {
261 var reduced = false;
262 for (reduce_info.dimensions) |dimension| {
263 if (dimension == @as(i64, @intCast(axis))) reduced = true;
264 }
265 if (reduced) continue;
266 kept[kept_index] = @intCast(axis);
267 kept_index += 1;
268 }
269 const expanded = try self.emitBroadcastInDim(primal, input_dims, kept);
270 const mask = try self.emitCompare(.eq, input.primal, expanded);
271 const zeros = try self.emitZeros(input.primal.ty);
272 const masked = try self.emitSelect(mask, input.tangent, zeros);
273 const zero_init = try self.emitZeros(init.tangent.ty);
274 input_term = try self.emitReduce(masked, zero_init, .sum, reduce_info.dimensions);
275 }
276
277 var init_term: ?trace.Value = null;
278 if (init_active) {
279 const init_primal = try self.emitToShape(init.primal, primal.ty);
280 const init_tangent = try self.emitToShape(init.tangent, primal.ty);
281 const init_mask = try self.emitCompare(.eq, primal, init_primal);
282 const zeros = try self.emitZeros(primal.ty);
283 init_term = try self.emitSelect(init_mask, init_tangent, zeros);
284 }
285
286 break :blk if (input_term) |input_value|
287 if (init_term) |init_value|
288 try self.emitBinary(.add, input_value, init_value)
289 else
290 input_value
291 else
292 init_term.?;
293 },
294 };
295 return .{
296 .primal = primal,
297 .tangent = tangent,
298 };
299 }
300
301 fn gatherValue(self: *@This(), op: *const program_mod.Operation, gather_info: program_mod.Gather, input: Value, indices: Value) !Value {
302 return .{
303 .primal = try self.next.bind(op, &.{ input.primal, indices.primal }),
304 .tangent = try self.emitGather(input.tangent, indices.primal, gather_info.axis),
305 };
306 }
307
308 fn scatterAddValue(
309 self: *@This(),
310 op: *const program_mod.Operation,
311 scatter_add: program_mod.ScatterAdd,
312 input: Value,
313 indices: Value,
314 updates: Value,
315 ) !Value {
316 return .{
317 .primal = try self.next.bind(op, &.{ input.primal, indices.primal, updates.primal }),
318 .tangent = try self.emitScatterAdd(input.tangent, indices.primal, updates.tangent, scatter_add.axis),
319 };
320 }
321
322 fn sparseCrossEntropyValue(
323 self: *@This(),
324 op: *const program_mod.Operation,
325 sparse_cross_entropy: program_mod.SparseCrossEntropy,
326 logits: Value,
327 targets: Value,
328 ) !Value {
329 const primal = try self.next.bind(op, &.{ logits.primal, targets.primal });
330
331 const logits_ty = logits.primal.ty;
332 const axis = sparse_cross_entropy.axis;
333 const class_axes = [_]i64{axis};
334 const scalar_ty = trace.Type.scalar(logits_ty.dtype);
335
336 const max_init = try self.emitFloatLowest(logits_ty.dtype);
337 const row_max = try emit.reduce(self, logits.primal, max_init, .max, class_axes[0..]);
338 const row_max_full = try self.emitExpandAxis(row_max, logits_ty.dims, axis);
339 const shifted = try emit.binary(self, .sub, logits.primal, row_max_full);
340 const exponentials = try emit.unary(self, .exp, shifted);
341 const sum_init = try self.emitZeros(scalar_ty);
342 const denominator = try emit.reduce(self, exponentials, sum_init, .sum, class_axes[0..]);
343 const denominator_full = try self.emitExpandAxis(denominator, logits_ty.dims, axis);
344 const softmax = try emit.binary(self, .div, exponentials, denominator_full);
345
346 const weighted = try emit.binary(self, .mul, softmax, logits.tangent);
347 const weighted_sum = try emit.reduce(self, weighted, sum_init, .sum, class_axes[0..]);
348
349 const positions = try emit.iota(self, .i32, logits_ty.dims, axis);
350 const target_positions = try self.emitExpandAxis(targets.primal, logits_ty.dims, axis);
351 const mask = try emit.compare(self, .eq, positions, target_positions);
352 const tangent_zeros = try self.emitZeros(logits_ty);
353 const masked_tangent = try emit.select(self, mask, logits.tangent, tangent_zeros);
354 const target_tangent = try emit.reduce(self, masked_tangent, sum_init, .sum, class_axes[0..]);
355
356 return .{
357 .primal = primal,
358 .tangent = try emit.binary(self, .sub, weighted_sum, target_tangent),
359 };
360 }
361
362 fn emitExpandAxis(self: *@This(), value: trace.Value, result_dims: []const trace.Dim, axis: i64) !trace.Value {
363 const allocator = self.builderHandle().arena.allocator();
364 const mapping = try allocator.alloc(i64, value.ty.rank());
365 var out: usize = 0;
366 for (0..result_dims.len) |position| {
367 if (position == @as(usize, @intCast(axis))) continue;
368 mapping[out] = @intCast(position);
369 out += 1;
370 }
371 const expanded_dims = try allocator.alloc(trace.Dim, result_dims.len);
372 for (result_dims, expanded_dims) |dim, *slot| slot.* = dim;
373 return emit.broadcastInDim(self, value, expanded_dims, mapping);
374 }
375
376 fn emitFloatLowest(self: *@This(), dtype: trace.DType) !trace.Value {
377 const scalar_ty = trace.Type.scalar(dtype);
378 return switch (dtype) {
379 .f16 => emit.fullFloat(self, scalar_ty, -std.math.floatMax(f16)),
380 .bf16, .f32 => emit.fullFloat(self, scalar_ty, -std.math.floatMax(f32)),
381 .f64 => emit.fullFloat(self, scalar_ty, -std.math.floatMax(f64)),
382 else => error.NonFloatDType,
383 };
384 }
385
386 fn dotGeneralValue(self: *@This(), op: *const program_mod.Operation, dot: program_mod.DotGeneral, lhs: Value, rhs: Value) !Value {
387 const primal = try self.next.bind(op, &.{ lhs.primal, rhs.primal });
388 const left = try self.emitDotGeneral(
389 lhs.tangent,
390 rhs.primal,
391 dot.lhs_contract,
392 dot.rhs_contract,
393 dot.lhs_batch,
394 dot.rhs_batch,
395 );
396 const right = try self.emitDotGeneral(
397 lhs.primal,
398 rhs.tangent,
399 dot.lhs_contract,
400 dot.rhs_contract,
401 dot.lhs_batch,
402 dot.rhs_batch,
403 );
404 return .{
405 .primal = primal,
406 .tangent = try self.emitBinary(.add, left, right),
407 };
408 }
409
410 fn emitZeros(self: *@This(), ty: trace.Type) !trace.Value {
411 return emit.zeros(self, ty);
412 }
413
414 fn emitToShape(self: *@This(), value: trace.Value, ty: trace.Type) !trace.Value {
415 if (value.ty.dtype != ty.dtype) return error.DTypeMismatch;
416 if (types.sameDims(value.ty.dims, ty.dims)) return value;
417 if (value.ty.rank() == 0) return self.emitBroadcastInDim(value, ty.dims, &.{});
418 return error.ShapeMismatch;
419 }
420
421 fn emitCompare(self: *@This(), direction: program_mod.CompareDirection, lhs: trace.Value, rhs: trace.Value) !trace.Value {
422 return emit.compare(self, direction, lhs, rhs);
423 }
424
425 fn emitSelect(self: *@This(), pred: trace.Value, on_true: trace.Value, on_false: trace.Value) !trace.Value {
426 return emit.select(self, pred, on_true, on_false);
427 }
428
429 fn emitFullFloat(self: *@This(), ty: trace.Type, fill_value: f64) !trace.Value {
430 return emit.fullFloat(self, ty, fill_value);
431 }
432
433 fn emitConstantBytes(self: *@This(), ty: trace.Type, payload: []const u8) !trace.Value {
434 return emit.constantBytes(self, ty, payload);
435 }
436
437 fn emitUnary(self: *@This(), op_kind: program_mod.Unary, input: trace.Value) !trace.Value {
438 return emit.unary(self, op_kind, input);
439 }
440
441 fn emitBinary(self: *@This(), op_kind: program_mod.Binary, lhs: trace.Value, rhs: trace.Value) !trace.Value {
442 return emit.binary(self, op_kind, lhs, rhs);
443 }
444
445 fn emitBroadcast(self: *@This(), input: trace.Value, result_dims: []const trace.Dim) !trace.Value {
446 return emit.broadcast(self, input, result_dims);
447 }
448
449 fn emitBroadcastInDim(self: *@This(), input: trace.Value, result_dims: []const trace.Dim, broadcast_dims: []const i64) !trace.Value {
450 return emit.broadcastInDim(self, input, result_dims, broadcast_dims);
451 }
452
453 fn emitReshape(self: *@This(), input: trace.Value, new_dims: []const trace.Dim) !trace.Value {
454 return emit.reshape(self, input, new_dims);
455 }
456
457 fn emitTranspose(self: *@This(), input: trace.Value, permutation: []const i64) !trace.Value {
458 return emit.transpose(self, input, permutation);
459 }
460
461 fn emitReduce(self: *@This(), input: trace.Value, init: trace.Value, reducer: program_mod.Reducer, dimensions: []const i64) !trace.Value {
462 return emit.reduce(self, input, init, reducer, dimensions);
463 }
464
465 fn emitGather(self: *@This(), input: trace.Value, indices: trace.Value, axis: i64) !trace.Value {
466 return emit.gather(self, input, indices, axis);
467 }
468
469 fn emitScatterAdd(self: *@This(), input: trace.Value, indices: trace.Value, updates: trace.Value, axis: i64) !trace.Value {
470 return emit.scatterAdd(self, input, indices, updates, axis);
471 }
472
473 fn emitDotGeneral(
474 self: *@This(),
475 lhs: trace.Value,
476 rhs: trace.Value,
477 lhs_contract: []const i64,
478 rhs_contract: []const i64,
479 lhs_batch: []const i64,
480 rhs_batch: []const i64,
481 ) !trace.Value {
482 return emit.dotGeneral(self, lhs, rhs, lhs_contract, rhs_contract, lhs_batch, rhs_batch);
483 }
484 };
485 }
486
487 pub fn linearize(
488 allocator: std.mem.Allocator,
489 source: *const program_mod.Program,
490 options: LinearizeOptions,
491 ) !Linearization {
492 if (source.containsScan()) {
493 var expanded = try unroll.apply(allocator, source);
494 defer expanded.deinit();
495 return linearize(allocator, &expanded, options);
496 }
497
498 try validate(source, options);
499
500 var builder = try trace.Builder.init(allocator, source.name);
501 errdefer builder.deinit();
502
503 const graph = interpret.Graph{ .builder = &builder };
504 return linearizeInto(allocator, source, options, graph);
505 }
506
507 pub fn linearizeWith(
508 allocator: std.mem.Allocator,
509 source: *const program_mod.Program,
510 options: LinearizeOptions,
511 hooks: anytype,
512 ) !Linearization {
513 if (source.containsScan()) {
514 var expanded = try unroll.apply(allocator, source);
515 defer expanded.deinit();
516 return linearizeWith(allocator, &expanded, options, hooks);
517 }
518
519 try validate(source, options);
520
521 var builder = try trace.Builder.init(allocator, source.name);
522 errdefer builder.deinit();
523
524 const graph = interpret.Graph{ .builder = &builder };
525 const generated = hook.attach("linearize", hooks, graph);
526 const linear = semantics(source, generated, options);
527 return interpret.run(allocator, source, hook.attach("jvp", hooks, linear));
528 }
529
530 fn linearizeInto(
531 allocator: std.mem.Allocator,
532 source: *const program_mod.Program,
533 options: LinearizeOptions,
534 initial: anytype,
535 ) !Linearization {
536 return interpret.run(allocator, source, semantics(source, initial, options));
537 }
538
539 pub fn jvp(
540 allocator: std.mem.Allocator,
541 source: *const program_mod.Program,
542 options: JvpOptions,
543 ) !program_mod.Program {
544 const linearized = try linearize(allocator, source, options);
545 return linearized.program;
546 }
547
548 pub fn jvpWith(
549 allocator: std.mem.Allocator,
550 source: *const program_mod.Program,
551 options: JvpOptions,
552 hooks: anytype,
553 ) !program_mod.Program {
554 const linearized = try linearizeWith(allocator, source, options, hooks);
555 return linearized.program;
556 }
557
558 fn unaryJvp(layer: anytype, op: program_mod.Unary, input: trace.Value, primal: trace.Value, tangent: trace.Value) !trace.Value {
559 if (!input.ty.dtype.isFloat()) return error.UnsupportedDerivative;
560 return switch (op) {
561 .neg => layer.emitUnary(.neg, tangent),
562 .exp => layer.emitBinary(.mul, primal, tangent),
563 .log => layer.emitBinary(.div, tangent, input),
564 .sqrt => blk: {
565 const two = try layer.emitFullFloat(input.ty, 2.0);
566 break :blk layer.emitBinary(.div, tangent, try layer.emitBinary(.mul, two, primal));
567 },
568 .tanh => blk: {
569 const one = try layer.emitFullFloat(input.ty, 1.0);
570 const square = try layer.emitBinary(.mul, primal, primal);
571 break :blk layer.emitBinary(.mul, tangent, try layer.emitBinary(.sub, one, square));
572 },
573 .sin => layer.emitBinary(.mul, tangent, try layer.emitUnary(.cos, input)),
574 .cos => layer.emitUnary(.neg, try layer.emitBinary(.mul, tangent, try layer.emitUnary(.sin, input))),
575 .abs => blk: {
576 const zeros = try layer.emitZeros(input.ty);
577 const non_negative = try layer.emitCompare(.ge, input, zeros);
578 break :blk layer.emitSelect(non_negative, tangent, try layer.emitUnary(.neg, tangent));
579 },
580 .tan => blk: {
581 const cosine = try layer.emitUnary(.cos, input);
582 break :blk layer.emitBinary(.div, tangent, try layer.emitBinary(.mul, cosine, cosine));
583 },
584 };
585 }
586
587 fn binaryJvp(
588 layer: anytype,
589 op: program_mod.Binary,
590 lhs: trace.Value,
591 rhs: trace.Value,
592 lhs_tangent: trace.Value,
593 rhs_tangent: trace.Value,
594 ) !trace.Value {
595 if (!lhs.ty.dtype.isFloat()) return error.UnsupportedDerivative;
596 return switch (op) {
597 .add => layer.emitBinary(.add, lhs_tangent, rhs_tangent),
598 .sub => layer.emitBinary(.sub, lhs_tangent, rhs_tangent),
599 .mul => layer.emitBinary(.add, try layer.emitBinary(.mul, lhs_tangent, rhs), try layer.emitBinary(.mul, lhs, rhs_tangent)),
600 .div => blk: {
601 const numerator = try layer.emitBinary(.sub, try layer.emitBinary(.mul, lhs_tangent, rhs), try layer.emitBinary(.mul, lhs, rhs_tangent));
602 const denominator = try layer.emitBinary(.mul, rhs, rhs);
603 break :blk layer.emitBinary(.div, numerator, denominator);
604 },
605 .max => blk: {
606 const mask = try layer.emitCompare(.ge, lhs, rhs);
607 break :blk layer.emitSelect(mask, lhs_tangent, rhs_tangent);
608 },
609 .min => blk: {
610 const mask = try layer.emitCompare(.le, lhs, rhs);
611 break :blk layer.emitSelect(mask, lhs_tangent, rhs_tangent);
612 },
613 .pow => blk: {
614 const lhs_active = !lhs_tangent.isStructuralZero();
615 const rhs_active = !rhs_tangent.isStructuralZero();
616 if (!lhs_active and !rhs_active) break :blk layer.emitZeros(lhs.ty);
617
618 var base_term: ?trace.Value = null;
619 if (lhs_active) {
620 const one = try layer.emitFullFloat(lhs.ty, 1.0);
621 const exponent_minus_one = try layer.emitBinary(.sub, rhs, one);
622 const power_step_down = try layer.emitBinary(.pow, lhs, exponent_minus_one);
623 base_term = try layer.emitBinary(
624 .mul,
625 try layer.emitBinary(.mul, rhs, power_step_down),
626 lhs_tangent,
627 );
628 }
629
630 var exponent_term: ?trace.Value = null;
631 if (rhs_active) {
632 exponent_term = try layer.emitBinary(
633 .mul,
634 try layer.emitBinary(.mul, try layer.emitBinary(.pow, lhs, rhs), try layer.emitUnary(.log, lhs)),
635 rhs_tangent,
636 );
637 }
638
639 break :blk if (base_term) |base|
640 if (exponent_term) |exponent|
641 try layer.emitBinary(.add, base, exponent)
642 else
643 base
644 else
645 exponent_term.?;
646 },
647 };
648 }
649
650 fn differentiatedParameters(
651 program: *program_mod.Program,
652 source: *const program_mod.Program,
653 options: LinearizeOptions,
654 ) ![]const usize {
655 const result = try program.arena.allocator().alloc(usize, differentiatedParameterCount(source, options.wrt));
656 var out: usize = 0;
657 for (source.parameters) |id| {
658 const parameter = source.operation(id).kind.parameter;
659 if (isWrt(options.wrt, parameter.index)) {
660 result[out] = parameter.index;
661 out += 1;
662 }
663 }
664 return result;
665 }
666
667 fn differentiatedParameterCount(source: *const program_mod.Program, wrt: []const usize) usize {
668 var count: usize = 0;
669 for (source.parameters) |id| {
670 const parameter = source.operation(id).kind.parameter;
671 if (isWrt(wrt, parameter.index)) count += 1;
672 }
673 return count;
674 }
675
676 pub fn validate(source: *const program_mod.Program, options: LinearizeOptions) !void {
677 for (options.wrt, 0..) |index, offset| {
678 if (index >= source.parameters.len) return error.ParameterOutOfRange;
679 for (options.wrt[0..offset]) |seen| {
680 if (seen == index) return error.DuplicateParameter;
681 }
682 }
683 }
684
685 fn isWrt(wrt: []const usize, index: usize) bool {
686 for (wrt) |item| {
687 if (item == index) return true;
688 }
689 return false;
690 }
691
692 fn jvpBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
693 return try (try args[0].mul(args[1])).tanh();
694 }
695
696 fn addZeroBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
697 const zero = try builder.full(.f32, .{ .lane = 4 }, 0.0);
698 return try args[0].add(zero);
699 }
700
701 fn powInactiveExponentBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
702 const three = try builder.scalar(.f32, 3.0);
703 const exponent = three;
704 return try args[0].pow(exponent);
705 }
706
707 fn reduceMaxWithInitBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
708 return try args[0].reduce(args[1], .max, .lane);
709 }
710
711 fn reduceMinWithInitBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
712 return try args[0].reduce(args[1], .min, .lane);
713 }
714
715 fn expectExtremaInitTangentMasked(source: *const program_mod.Program) !void {
716 var differentiated = try linearize(std.testing.allocator, source, .{ .wrt = &.{ 0, 1 } });
717 defer differentiated.deinit();
718
719 try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);
720 const init_tangent = differentiated.program.parameters[3];
721 var reduce_seeded_by_init_tangent = false;
722 var masked_init_tangent = false;
723 for (differentiated.program.operations) |op| {
724 switch (op.kind) {
725 .reduce => |reduce| {
726 if (reduce.reducer == .sum and reduce.init.index == init_tangent.index) {
727 reduce_seeded_by_init_tangent = true;
728 }
729 },
730 .select => |select| {
731 if (select.on_true.index == init_tangent.index) masked_init_tangent = true;
732 },
733 else => {},
734 }
735 }
736
737 try std.testing.expect(!reduce_seeded_by_init_tangent);
738 try std.testing.expect(masked_init_tangent);
739 }
740
741 const CountDualBinary = struct {
742 count: *usize,
743
744 pub fn bind(self: *@This(), ctx: anytype) !Dual {
745 switch (ctx.op.kind) {
746 .binary => self.count.* += 1,
747 else => {},
748 }
749 return ctx.default();
750 }
751 };
752
753 const TraceBinaryCounts = struct {
754 add: usize = 0,
755 mul: usize = 0,
756 };
757
758 const TraceBinaryCounter = struct {
759 counts: *TraceBinaryCounts,
760
761 pub fn add(self: *@This(), ctx: anytype) !trace.Value {
762 self.counts.add += 1;
763 return ctx.default();
764 }
765
766 pub fn mul(self: *@This(), ctx: anytype) !trace.Value {
767 self.counts.mul += 1;
768 return ctx.default();
769 }
770 };
771
772 test "tensor linearize semantics accepts user layers over dual values" {
773 var source = try trace.define(std.testing.allocator, "jvp_layer", &.{
774 types.spec(.f32, .{ .lane = 4 }),
775 }, addZeroBody);
776 defer source.deinit();
777
778 var builder = try trace.Builder.init(std.testing.allocator, source.name);
779 errdefer builder.deinit();
780
781 var count: usize = 0;
782 const graph = interpret.Graph{ .builder = &builder };
783 const linear = semantics(&source, graph, .{ .wrt = &.{0} });
784 var differentiated = try interpret.run(std.testing.allocator, &source, interpret.layer(Dual, linear, CountDualBinary{ .count = &count }));
785 defer differentiated.deinit();
786
787 try std.testing.expectEqual(@as(usize, 1), count);
788 try std.testing.expectEqual(@as(usize, 2), differentiated.program.parameters.len);
789 try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);
790 }
791
792 test "tensor linearize binds generated tangent arithmetic through downstream semantics" {
793 var source = try trace.define(std.testing.allocator, "jvp_generated_arithmetic", &.{
794 types.spec(.f32, .{ .lane = 4 }),
795 types.spec(.f32, .{ .lane = 4 }),
796 }, struct {
797 fn body(_: *trace.Builder, args: []const trace.Value) !trace.Value {
798 return try args[0].mul(args[1]);
799 }
800 }.body);
801 defer source.deinit();
802
803 var builder = try trace.Builder.init(std.testing.allocator, source.name);
804 errdefer builder.deinit();
805
806 var counts: TraceBinaryCounts = .{};
807 const graph = interpret.Graph{ .builder = &builder };
808 const counted = interpret.bind(TraceBinaryCounter{ .counts = &counts }).attach(graph);
809 var differentiated = try interpret.run(std.testing.allocator, &source, semantics(&source, counted, .{ .wrt = &.{ 0, 1 } }));
810 defer differentiated.deinit();
811
812 try std.testing.expectEqual(@as(usize, 1), counts.add);
813 try std.testing.expectEqual(@as(usize, 3), counts.mul);
814 try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);
815 try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);
816 }
817
818 const GeneratedAddMetadata = struct {
819 seen: *usize,
820
821 pub fn add(self: *@This(), ctx: *transform.Context) !?trace.Value {
822 self.seen.* += 1;
823 try std.testing.expect(!ctx.isZero(0));
824 try std.testing.expect(!ctx.isZero(1));
825 try std.testing.expect(ctx.constantPayload(0) == null);
826 try std.testing.expect(ctx.constantPayload(1) == null);
827 return null;
828 }
829 };
830
831 test "tensor linearize generated arithmetic is safe for downstream rewrite metadata queries" {
832 var source = try trace.define(std.testing.allocator, "jvp_generated_metadata", &.{
833 types.spec(.f32, .{ .lane = 4 }),
834 types.spec(.f32, .{ .lane = 4 }),
835 }, struct {
836 fn body(_: *trace.Builder, args: []const trace.Value) !trace.Value {
837 return try args[0].mul(args[1]);
838 }
839 }.body);
840 defer source.deinit();
841
842 var builder = try trace.Builder.init(std.testing.allocator, source.name);
843 errdefer builder.deinit();
844
845 var seen: usize = 0;
846 const graph = interpret.Graph{ .builder = &builder };
847 const rewrite = transform.semantics(&source, graph, GeneratedAddMetadata{ .seen = &seen });
848 var differentiated = try interpret.run(std.testing.allocator, &source, semantics(&source, rewrite, .{ .wrt = &.{ 0, 1 } }));
849 defer differentiated.deinit();
850
851 try std.testing.expectEqual(@as(usize, 1), seen);
852 try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);
853 try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);
854 }
855
856 const DropAddZero = struct {
857 pub fn add(_: *@This(), ctx: *transform.Context) !?trace.Value {
858 if (ctx.isZero(1)) return ctx.arg(0);
859 if (ctx.isZero(0)) return ctx.arg(1);
860 return null;
861 }
862 };
863
864 test "tensor linearize semantics delegates primal binds through rewrite layers" {
865 var source = try trace.define(std.testing.allocator, "jvp_rewrite_stack", &.{
866 types.spec(.f32, .{ .lane = 4 }),
867 }, addZeroBody);
868 defer source.deinit();
869
870 var plain = try linearize(std.testing.allocator, &source, .{ .wrt = &.{0} });
871 defer plain.deinit();
872
873 var builder = try trace.Builder.init(std.testing.allocator, source.name);
874 errdefer builder.deinit();
875
876 const graph = interpret.Graph{ .builder = &builder };
877 const rewrite = transform.semantics(&source, graph, DropAddZero{});
878 var rewritten = try interpret.run(std.testing.allocator, &source, semantics(&source, rewrite, .{ .wrt = &.{0} }));
879 defer rewritten.deinit();
880
881 try std.testing.expect(rewritten.program.operationCount() < plain.program.operationCount());
882 try std.testing.expectEqual(@as(u32, 0), rewritten.primalOutputs()[0].index);
883 try types.expectExtents(&.{4}, rewritten.program.typeOf(rewritten.tangentOutputs()[0]));
884 }
885
886 test "tensor linearize exposes primal and tangent structure" {
887 var source = try trace.define(std.testing.allocator, "jvp", &.{
888 types.spec(.f32, .{ .lane = 4 }),
889 types.spec(.f32, .{ .lane = 4 }),
890 }, jvpBody);
891 defer source.deinit();
892
893 var differentiated = try linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
894 defer differentiated.deinit();
895
896 try std.testing.expectEqual(@as(usize, 2), differentiated.primal_parameter_count);
897 try std.testing.expectEqual(@as(usize, 2), differentiated.tangent_parameter_count);
898 try std.testing.expectEqual(@as(usize, 1), differentiated.primal_output_count);
899 try std.testing.expectEqual(@as(usize, 1), differentiated.tangent_output_count);
900 try std.testing.expectEqualSlices(usize, &.{ 0, 1 }, differentiated.differentiated_parameters);
901 try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);
902 try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);
903 try types.expectExtents(&.{4}, differentiated.program.typeOf(differentiated.primalOutputs()[0]));
904 try types.expectExtents(&.{4}, differentiated.program.typeOf(differentiated.tangentOutputs()[0]));
905 }
906
907 test "tensor linearize omits inactive pow exponent branch" {
908 var source = try trace.define(std.testing.allocator, "jvp_pow_inactive_exponent", &.{
909 types.spec(.f32, .{ .lane = 4 }),
910 }, powInactiveExponentBody);
911 defer source.deinit();
912
913 var differentiated = try linearize(std.testing.allocator, &source, .{ .wrt = &.{0} });
914 defer differentiated.deinit();
915
916 var logs: usize = 0;
917 for (differentiated.program.operations) |op| {
918 switch (op.kind) {
919 .unary => |unary_op| {
920 if (unary_op.op == .log) logs += 1;
921 },
922 else => {},
923 }
924 }
925
926 try std.testing.expectEqual(@as(usize, 0), logs);
927 try types.expectExtents(&.{4}, differentiated.program.typeOf(differentiated.tangentOutputs()[0]));
928 }
929
930 test "tensor linearize masks reduce max init tangents" {
931 var source = try trace.define(std.testing.allocator, "jvp_reduce_max_init", &.{
932 types.spec(.f32, .{ .lane = 4 }),
933 types.spec(.f32, .{}),
934 }, reduceMaxWithInitBody);
935 defer source.deinit();
936
937 try expectExtremaInitTangentMasked(&source);
938 }
939
940 test "tensor linearize masks reduce min init tangents" {
941 var source = try trace.define(std.testing.allocator, "jvp_reduce_min_init", &.{
942 types.spec(.f32, .{ .lane = 4 }),
943 types.spec(.f32, .{}),
944 }, reduceMinWithInitBody);
945 defer source.deinit();
946
947 try expectExtremaInitTangentMasked(&source);
948 }
949
950 test "tensor jvp is a linearize convenience wrapper" {
951 var source = try trace.define(std.testing.allocator, "jvp_wrapper", &.{
952 types.spec(.f32, .{ .lane = 4 }),
953 types.spec(.f32, .{ .lane = 4 }),
954 }, jvpBody);
955 defer source.deinit();
956
957 var differentiated = try jvp(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
958 defer differentiated.deinit();
959
960 try std.testing.expectEqual(@as(usize, 4), differentiated.parameters.len);
961 try std.testing.expectEqual(@as(usize, 2), differentiated.outputs.len);
962 try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[0]));
963 try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[1]));
964 }
965
966 test "tensor jvpWith routes linearize generated ops through user semantics" {
967 var source = try trace.define(std.testing.allocator, "jvp_with_generated_arithmetic", &.{
968 types.spec(.f32, .{ .lane = 4 }),
969 types.spec(.f32, .{ .lane = 4 }),
970 }, struct {
971 fn body(_: *trace.Builder, args: []const trace.Value) !trace.Value {
972 return try args[0].mul(args[1]);
973 }
974 }.body);
975 defer source.deinit();
976
977 var counts: TraceBinaryCounts = .{};
978 var differentiated = try jvpWith(
979 std.testing.allocator,
980 &source,
981 .{ .wrt = &.{ 0, 1 } },
982 .{ .linearize = interpret.bind(TraceBinaryCounter{ .counts = &counts }) },
983 );
984 defer differentiated.deinit();
985
986 try std.testing.expectEqual(@as(usize, 1), counts.add);
987 try std.testing.expectEqual(@as(usize, 3), counts.mul);
988 try std.testing.expectEqual(@as(usize, 4), differentiated.parameters.len);
989 try std.testing.expectEqual(@as(usize, 2), differentiated.outputs.len);
990 try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[0]));
991 try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[1]));
992 }
993
994 fn doublingCallBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
995 _ = builder;
996 return args[0].builder.customCall("accy.custom.double", 1, &.{args[0]}, args[0].ty);
997 }
998
999 test "tensor linearize rejects opaque custom calls naming the contract" {
1000 var source = try trace.define(std.testing.allocator, "jvp_opaque_custom", &.{
1001 types.spec(.f32, .{ .lane = 4 }),
1002 }, doublingCallBody);
1003 defer source.deinit();
1004
1005 try std.testing.expectError(
1006 error.CustomCallRequiresJvpContract,
1007 linearize(std.testing.allocator, &source, .{ .wrt = &.{0} }),
1008 );
1009 }
1010
1011 const DoubleJvpRule = struct {
1012 applied: *usize,
1013
1014 pub fn bind(self: *@This(), ctx: anytype) !Dual {
1015 switch (ctx.op.kind) {
1016 .custom_call => |custom| {
1017 if (std.mem.eql(u8, custom.target, "accy.custom.double")) {
1018 self.applied.* += 1;
1019 const builder = ctx.builderHandle();
1020 return .{
1021 .primal = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].primal}, ctx.op.result),
1022 .tangent = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].tangent}, ctx.op.result),
1023 };
1024 }
1025 },
1026 else => {},
1027 }
1028 return ctx.default();
1029 }
1030 };
1031
1032 test "tensor linearize accepts custom call jvp contracts through user layers" {
1033 var source = try trace.define(std.testing.allocator, "jvp_custom_contract", &.{
1034 types.spec(.f32, .{ .lane = 4 }),
1035 }, doublingCallBody);
1036 defer source.deinit();
1037
1038 var builder = try trace.Builder.init(std.testing.allocator, source.name);
1039 errdefer builder.deinit();
1040
1041 var applied: usize = 0;
1042 const graph = interpret.Graph{ .builder = &builder };
1043 const linear = semantics(&source, graph, .{ .wrt = &.{0} });
1044 var differentiated = try interpret.run(
1045 std.testing.allocator,
1046 &source,
1047 interpret.layer(Dual, linear, DoubleJvpRule{ .applied = &applied }),
1048 );
1049 defer differentiated.deinit();
1050
1051 try std.testing.expectEqual(@as(usize, 1), applied);
1052 try std.testing.expectEqual(@as(usize, 2), differentiated.program.parameters.len);
1053 try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);
1054
1055 var custom_calls: usize = 0;
1056 for (differentiated.program.operations) |op| {
1057 switch (op.kind) {
1058 .custom_call => custom_calls += 1,
1059 else => {},
1060 }
1061 }
1062 try std.testing.expectEqual(@as(usize, 2), custom_calls);
1063 }
1064
1065 test "tensor jvpWith accepts custom call jvp contracts through public hooks" {
1066 var source = try trace.define(std.testing.allocator, "jvp_public_custom_contract", &.{
1067 types.spec(.f32, .{ .lane = 4 }),
1068 }, doublingCallBody);
1069 defer source.deinit();
1070
1071 var applied: usize = 0;
1072 var differentiated = try jvpWith(
1073 std.testing.allocator,
1074 &source,
1075 .{ .wrt = &.{0} },
1076 .{ .jvp = interpret.bind(DoubleJvpRule{ .applied = &applied }) },
1077 );
1078 defer differentiated.deinit();
1079
1080 try std.testing.expectEqual(@as(usize, 1), applied);
1081 try std.testing.expectEqual(@as(usize, 2), differentiated.parameters.len);
1082 try std.testing.expectEqual(@as(usize, 2), differentiated.outputs.len);
1083
1084 var custom_calls: usize = 0;
1085 for (differentiated.operations) |op| {
1086 switch (op.kind) {
1087 .custom_call => custom_calls += 1,
1088 else => {},
1089 }
1090 }
1091 try std.testing.expectEqual(@as(usize, 2), custom_calls);
1092 }