lib/accy/src/tensor/unroll.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const interpret = @import("interpret/root.zig");
3 const program_mod = @import("program.zig");
4 const trace = @import("trace/root.zig");
5 const types = @import("type/root.zig");
6
7 pub fn apply(allocator: std.mem.Allocator, source: *const program_mod.Program) !program_mod.Program {
8 var builder = try trace.Builder.init(allocator, source.name);
9 errdefer builder.deinit();
10
11 var scratch = std.heap.ArenaAllocator.init(allocator);
12 defer scratch.deinit();
13
14 const graph = interpret.Graph{ .builder = &builder };
15 return try interpret.run(allocator, source, Semantics(@TypeOf(graph)){
16 .next = graph,
17 .scratch = scratch.allocator(),
18 });
19 }
20
21 fn Semantics(comptime Next: type) type {
22 return struct {
23 next: Next,
24 scratch: std.mem.Allocator,
25 finals: std.AutoHashMapUnmanaged(u32, []trace.Value) = .empty,
26
27 pub const Value: type = trace.Value;
28 pub const Result: type = Next.Result;
29
30 pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {
31 switch (step.op.kind) {
32 .scan => |scan| {
33 var buffer: [program_mod.max_scan_carries]trace.Value = undefined;
34 for (scan.inits, 0..) |init_id, index| buffer[index] = step.read(init_id);
35 const expanded = try self.expandScan(scan, buffer[0..scan.inits.len]);
36 try self.finals.put(self.scratch, step.op.id.index, expanded);
37 return expanded[0];
38 },
39 .projection => |projection| {
40 return self.finals.get(projection.source.index).?[projection.index];
41 },
42 else => {
43 var buffer: [program_mod.max_operation_operands]trace.Value = undefined;
44 return self.next.bind(step.op, interpret.arguments(Value, step.op, step.values, &buffer));
45 },
46 }
47 }
48
49 fn expandScan(self: *@This(), scan: program_mod.Scan, inits: []const trace.Value) anyerror![]trace.Value {
50 const carries = try self.scratch.dupe(trace.Value, inits);
51 var iteration: i64 = 0;
52 while (iteration < scan.length) : (iteration += 1) {
53 try self.expandStep(scan.body, carries);
54 }
55 return carries;
56 }
57
58 fn expandStep(self: *@This(), body: *const program_mod.Subgraph, carries: []trace.Value) !void {
59 const values = try self.scratch.alloc(trace.Value, body.values.len);
60 var nested_finals: std.AutoHashMapUnmanaged(u32, []trace.Value) = .empty;
61 for (body.operations) |*op| {
62 values[op.id.index] = switch (op.kind) {
63 .parameter => |parameter| carries[parameter.index],
64 .scan => |nested| blk: {
65 var buffer: [program_mod.max_scan_carries]trace.Value = undefined;
66 for (nested.inits, 0..) |init_id, index| buffer[index] = values[init_id.index];
67 const expanded = try self.expandScan(nested, buffer[0..nested.inits.len]);
68 try nested_finals.put(self.scratch, op.id.index, expanded);
69 break :blk expanded[0];
70 },
71 .projection => |projection| nested_finals.get(projection.source.index).?[projection.index],
72 else => blk: {
73 var buffer: [program_mod.max_operation_operands]trace.Value = undefined;
74 const args = interpret.arguments(Value, op, values, &buffer);
75 break :blk try self.next.bind(op, args);
76 },
77 };
78 }
79 var buffer: [program_mod.max_scan_carries]trace.Value = undefined;
80 for (body.outputs, 0..) |id, index| buffer[index] = values[id.index];
81 @memcpy(carries, buffer[0..body.outputs.len]);
82 }
83
84 pub fn finish(self: *@This(), outputs: []const Value) !Result {
85 return self.next.finish(outputs);
86 }
87
88 pub fn builderHandle(self: *@This()) *trace.Builder {
89 return self.next.builderHandle();
90 }
91 };
92 }
93
94 fn incrementScanStep(scan_builder: *trace.Builder, carry: trace.Value) !trace.Value {
95 const bump = try scan_builder.full(.f32, .{ .lane = 2 }, 1.0);
96 return carry.add(bump);
97 }
98
99 test "unroll expands scan operations into scan free graphs" {
100 var builder = try trace.Builder.init(std.testing.allocator, "unroll_expand");
101 defer builder.deinit();
102
103 const x0 = try builder.input(.f32, .{ .lane = 2 });
104 const walked = try builder.scan(.{
105 .length = 3,
106 .init = x0,
107 .body = incrementScanStep,
108 });
109 var source = try builder.finish(&.{walked});
110 defer source.deinit();
111
112 var expanded = try apply(std.testing.allocator, &source);
113 defer expanded.deinit();
114
115 for (expanded.operations) |op| {
116 try std.testing.expect(op.kind != .scan);
117 try std.testing.expect(op.kind != .projection);
118 }
119 try std.testing.expectEqual(@as(usize, 7), expanded.operationCount());
120 try types.expectExtents(&.{2}, expanded.typeOf(expanded.outputs[0]));
121 }
122
123 test "unroll expands across carry arities and lengths including zero" {
124 inline for (1..4) |arity| {
125 inline for (0..4) |length| {
126 var builder = try trace.Builder.init(std.testing.allocator, "unroll_grid");
127 defer builder.deinit();
128
129 var inits: [arity]trace.Value = undefined;
130 for (0..arity) |index| {
131 inits[index] = try builder.full(.f32, .{ .lane = 2 }, @floatFromInt(index + 1));
132 }
133 var scope = try builder.scanScope(@intCast(length), inits[0..]);
134 var nexts: [arity]trace.Value = undefined;
135 for (0..arity) |index| {
136 const partner = scope.carry((index + 1) % arity);
137 nexts[index] = try scope.carry(index).add(partner);
138 }
139 const walked = try scope.finish(nexts[0..]);
140 var outputs_buffer: [arity]trace.Value = undefined;
141 outputs_buffer[0] = walked;
142 for (1..arity) |index| {
143 outputs_buffer[index] = try builder.projection(walked, index);
144 }
145 var source = try builder.finish(outputs_buffer[0..]);
146 defer source.deinit();
147
148 var expanded = try apply(std.testing.allocator, &source);
149 defer expanded.deinit();
150
151 for (expanded.operations) |op| {
152 try std.testing.expect(op.kind != .scan);
153 try std.testing.expect(op.kind != .projection);
154 }
155 try std.testing.expectEqual(arity + length * arity, expanded.operationCount());
156 for (expanded.outputs) |id| {
157 try types.expectExtents(&.{2}, expanded.typeOf(id));
158 }
159 }
160 }
161 }
162
163 fn projectionScanStep(_: *trace.Builder, carry: anytype) !@TypeOf(carry) {
164 const doubled = try carry.x.add(carry.x);
165 return .{
166 .x = doubled,
167 .acc = try carry.acc.add(doubled),
168 };
169 }
170
171 test "unroll resolves projections to expanded carries" {
172 var builder = try trace.Builder.init(std.testing.allocator, "unroll_projection");
173 defer builder.deinit();
174
175 const x0 = try builder.input(.f32, .{ .lane = 2 });
176 const acc0 = try builder.full(.f32, .{ .lane = 2 }, 0.0);
177 const walked = try builder.scan(.{
178 .length = 2,
179 .init = .{ .x = x0, .acc = acc0 },
180 .body = projectionScanStep,
181 });
182 var source = try builder.finish(&.{walked.acc});
183 defer source.deinit();
184
185 var expanded = try apply(std.testing.allocator, &source);
186 defer expanded.deinit();
187
188 for (expanded.operations) |op| {
189 try std.testing.expect(op.kind != .scan);
190 try std.testing.expect(op.kind != .projection);
191 }
192 try types.expectExtents(&.{2}, expanded.typeOf(expanded.outputs[0]));
193 }