lib/accy/src/tensor/transform.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 const Operand = struct {
8 id: program_mod.Id,
9 value: trace.Value,
10 ty: trace.Type,
11 };
12
13 pub const Context = struct {
14 source: *const program_mod.Program,
15 builder: *trace.Builder,
16 op: *const program_mod.Operation,
17 args: []const trace.Value,
18
19 pub fn arg(self: Context, index: usize) trace.Value {
20 return self.args[index];
21 }
22
23 pub fn operand(self: Context, index: usize) Operand {
24 const id = self.operandId(index).?;
25 const value = self.arg(index);
26 return .{
27 .id = id,
28 .value = value,
29 .ty = value.ty,
30 };
31 }
32
33 pub fn isZero(self: Context, index: usize) bool {
34 const id = self.operandId(index) orelse return false;
35 if (id.index >= self.source.valueCount()) return false;
36 return self.source.isZeroConstant(id);
37 }
38
39 pub fn constantPayload(self: Context, index: usize) ?[]const u8 {
40 const id = self.operandId(index) orelse return null;
41 if (id.index >= self.source.valueCount()) return null;
42 return self.source.constantPayload(id);
43 }
44
45 fn operandId(self: Context, index: usize) ?program_mod.Id {
46 return switch (self.op.kind) {
47 .parameter, .constant, .iota => null,
48 .unary => |unary| if (index == 0) unary.input else null,
49 .binary => |binary| switch (index) {
50 0 => binary.lhs,
51 1 => binary.rhs,
52 else => null,
53 },
54 .broadcast => |broadcast| if (index == 0) broadcast.input else null,
55 .broadcast_in_dim => |broadcast| if (index == 0) broadcast.input else null,
56 .reshape => |reshape| if (index == 0) reshape.input else null,
57 .transpose => |transpose| if (index == 0) transpose.input else null,
58 .reduce => |reduce| switch (index) {
59 0 => reduce.input,
60 1 => reduce.init,
61 else => null,
62 },
63 .gather => |gather| switch (index) {
64 0 => gather.input,
65 1 => gather.indices,
66 else => null,
67 },
68 .scatter_add => |scatter_add| switch (index) {
69 0 => scatter_add.input,
70 1 => scatter_add.indices,
71 2 => scatter_add.updates,
72 else => null,
73 },
74 .sparse_cross_entropy => |sparse_cross_entropy| switch (index) {
75 0 => sparse_cross_entropy.logits,
76 1 => sparse_cross_entropy.targets,
77 else => null,
78 },
79 .dot_general => |dot| switch (index) {
80 0 => dot.lhs,
81 1 => dot.rhs,
82 else => null,
83 },
84 .compare => |compare| switch (index) {
85 0 => compare.lhs,
86 1 => compare.rhs,
87 else => null,
88 },
89 .select => |select| switch (index) {
90 0 => select.pred,
91 1 => select.on_true,
92 2 => select.on_false,
93 else => null,
94 },
95 .custom_call => |custom| if (index < custom.operands.len) custom.operands[index] else null,
96 .scan => |scan| if (index < scan.inits.len) scan.inits[index] else null,
97 .projection => |projection| if (index == 0) projection.source else null,
98 };
99 }
100 };
101
102 pub fn apply(allocator: std.mem.Allocator, source: *const program_mod.Program, pass: anytype) !program_mod.Program {
103 var builder = try trace.Builder.init(allocator, source.name);
104 errdefer builder.deinit();
105
106 const graph = interpret.Graph{ .builder = &builder };
107 return try interpret.run(allocator, source, semantics(source, graph, pass));
108 }
109
110 pub fn semantics(source: *const program_mod.Program, next: anytype, pass: anytype) interpret.Layer(trace.Value, @TypeOf(next), Rewrite(@TypeOf(pass))) {
111 return interpret.layer(trace.Value, next, Rewrite(@TypeOf(pass)){
112 .source = source,
113 .pass = pass,
114 });
115 }
116
117 pub fn Rewrite(comptime Pass: type) type {
118 return struct {
119 source: *const program_mod.Program,
120 pass: Pass,
121
122 pub fn bind(self: *@This(), layer: anytype) !trace.Value {
123 var ctx = Context{
124 .source = self.source,
125 .builder = layer.builderHandle(),
126 .op = layer.op,
127 .args = layer.args,
128 };
129 return (try dispatch(&self.pass, &ctx, layer.op)) orelse layer.default();
130 }
131 };
132 }
133
134 fn dispatch(pass: anytype, ctx: *Context, op: *const program_mod.Operation) !?trace.Value {
135 const Pass = @TypeOf(pass.*);
136 return switch (op.kind) {
137 .binary => |binary| switch (binary.op) {
138 .add => if (comptime @hasDecl(Pass, "add"))
139 pass.add(ctx)
140 else
141 null,
142 .sub => if (comptime @hasDecl(Pass, "sub"))
143 pass.sub(ctx)
144 else
145 null,
146 .mul => if (comptime @hasDecl(Pass, "mul"))
147 pass.mul(ctx)
148 else
149 null,
150 .div => if (comptime @hasDecl(Pass, "div"))
151 pass.div(ctx)
152 else
153 null,
154 .max => if (comptime @hasDecl(Pass, "max"))
155 pass.max(ctx)
156 else
157 null,
158 .min => if (comptime @hasDecl(Pass, "min"))
159 pass.min(ctx)
160 else
161 null,
162 .pow => if (comptime @hasDecl(Pass, "pow"))
163 pass.pow(ctx)
164 else
165 null,
166 },
167 .unary => |unary| switch (unary.op) {
168 .neg => if (comptime @hasDecl(Pass, "neg"))
169 pass.neg(ctx)
170 else
171 null,
172 .abs => if (comptime @hasDecl(Pass, "abs"))
173 pass.abs(ctx)
174 else
175 null,
176 .exp => if (comptime @hasDecl(Pass, "exp"))
177 pass.exp(ctx)
178 else
179 null,
180 .log => if (comptime @hasDecl(Pass, "log"))
181 pass.log(ctx)
182 else
183 null,
184 .sqrt => if (comptime @hasDecl(Pass, "sqrt"))
185 pass.sqrt(ctx)
186 else
187 null,
188 .tanh => if (comptime @hasDecl(Pass, "tanh"))
189 pass.tanh(ctx)
190 else
191 null,
192 .sin => if (comptime @hasDecl(Pass, "sin"))
193 pass.sin(ctx)
194 else
195 null,
196 .cos => if (comptime @hasDecl(Pass, "cos"))
197 pass.cos(ctx)
198 else
199 null,
200 .tan => if (comptime @hasDecl(Pass, "tan"))
201 pass.tan(ctx)
202 else
203 null,
204 },
205 .reduce => if (comptime @hasDecl(Pass, "reduce"))
206 pass.reduce(ctx)
207 else
208 null,
209 .gather => if (comptime @hasDecl(Pass, "gather"))
210 pass.gather(ctx)
211 else
212 null,
213 .scatter_add => if (comptime @hasDecl(Pass, "scatterAdd"))
214 pass.scatterAdd(ctx)
215 else
216 null,
217 .sparse_cross_entropy => if (comptime @hasDecl(Pass, "sparseCrossEntropy"))
218 pass.sparseCrossEntropy(ctx)
219 else
220 null,
221 .dot_general => if (comptime @hasDecl(Pass, "dotGeneral"))
222 pass.dotGeneral(ctx)
223 else
224 null,
225 .compare => if (comptime @hasDecl(Pass, "compare"))
226 pass.compare(ctx)
227 else
228 null,
229 .select => if (comptime @hasDecl(Pass, "select"))
230 pass.select(ctx)
231 else
232 null,
233 .custom_call => if (comptime @hasDecl(Pass, "customCall"))
234 pass.customCall(ctx)
235 else
236 null,
237 .broadcast => if (comptime @hasDecl(Pass, "broadcast"))
238 pass.broadcast(ctx)
239 else
240 null,
241 .broadcast_in_dim => if (comptime @hasDecl(Pass, "broadcastInDim"))
242 pass.broadcastInDim(ctx)
243 else
244 null,
245 .reshape => if (comptime @hasDecl(Pass, "reshape"))
246 pass.reshape(ctx)
247 else
248 null,
249 .transpose => if (comptime @hasDecl(Pass, "transpose"))
250 pass.transpose(ctx)
251 else
252 null,
253 .iota => if (comptime @hasDecl(Pass, "iota"))
254 pass.iota(ctx)
255 else
256 null,
257 .parameter => if (comptime @hasDecl(Pass, "parameter"))
258 pass.parameter(ctx)
259 else
260 null,
261 .scan => if (comptime @hasDecl(Pass, "scan"))
262 pass.scan(ctx)
263 else
264 null,
265 .projection => if (comptime @hasDecl(Pass, "projection"))
266 pass.projection(ctx)
267 else
268 null,
269 .constant => null,
270 };
271 }
272
273 test "tensor transform defaults to structural copy" {
274 var builder = try trace.Builder.init(std.testing.allocator, "copy");
275 defer builder.deinit();
276
277 const x = try builder.input(.f32, .{ .lane = 4 });
278 const y = try builder.input(.f32, .{ .lane = 4 });
279 const out = try (try x.add(y)).tanh();
280 var source = try builder.finish(&.{out});
281 defer source.deinit();
282
283 var rewritten = try apply(std.testing.allocator, &source, struct {}{});
284 defer rewritten.deinit();
285
286 try std.testing.expectEqual(source.valueCount(), rewritten.valueCount());
287 try std.testing.expectEqual(source.operationCount(), rewritten.operationCount());
288 try std.testing.expect(types.sameDims(source.typeOf(source.outputs[0]).dims, rewritten.typeOf(rewritten.outputs[0]).dims));
289 }
290
291 const DropAddZero = struct {
292 pub fn add(_: *@This(), ctx: *Context) !?trace.Value {
293 const lhs = ctx.operand(0);
294 const rhs = ctx.operand(1);
295 if (ctx.isZero(1)) return lhs.value;
296 if (ctx.isZero(0)) return rhs.value;
297 return null;
298 }
299 };
300
301 test "tensor transform lets users rewrite one primitive" {
302 var builder = try trace.Builder.init(std.testing.allocator, "drop_add_zero");
303 defer builder.deinit();
304
305 const x = try builder.input(.f32, .{ .lane = 4 });
306 const zero = try builder.full(.f32, .{ .lane = 4 }, 0.0);
307 const out = try x.add(zero);
308 var source = try builder.finish(&.{out});
309 defer source.deinit();
310
311 var rewritten = try apply(std.testing.allocator, &source, DropAddZero{});
312 defer rewritten.deinit();
313
314 try std.testing.expectEqual(@as(usize, 2), rewritten.valueCount());
315 try std.testing.expectEqual(@as(u32, 0), rewritten.outputs[0].index);
316 }
317
318 const BindSecondParameter = struct {
319 pub fn parameter(_: *@This(), ctx: *Context) !?trace.Value {
320 const parameter_info = ctx.op.kind.parameter;
321 if (parameter_info.index == 1) return try ctx.builder.fullFloat(ctx.op.result, 1.0);
322 return null;
323 }
324 };
325
326 test "tensor transform lets users bind parameters" {
327 var builder = try trace.Builder.init(std.testing.allocator, "bind_parameter");
328 defer builder.deinit();
329
330 const x = try builder.input(.f32, .{ .lane = 4 });
331 const scale = try builder.input(.f32, .{ .lane = 4 });
332 const out = try x.mul(scale);
333 var source = try builder.finish(&.{out});
334 defer source.deinit();
335
336 var rewritten = try apply(std.testing.allocator, &source, BindSecondParameter{});
337 defer rewritten.deinit();
338
339 try std.testing.expectEqual(@as(usize, 1), rewritten.parameters.len);
340 try types.expectExtents(&.{4}, rewritten.typeOf(rewritten.outputs[0]));
341 }