lib/accy/src/tensor/emit.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const program_mod = @import("program.zig");
3 const trace = @import("trace/root.zig");
4 const types = @import("type/root.zig");
5
6 pub fn zeros(layer: anytype, ty: trace.Type) !trace.Value {
7 const payload = try layer.builderHandle().arena.allocator().alloc(u8, try ty.byteCount());
8 @memset(payload, 0);
9 return constantBytes(layer, ty, payload);
10 }
11
12 pub fn fullFloat(layer: anytype, ty: trace.Type, fill_value: f64) !trace.Value {
13 const owned = try trace.Type.init(layer.builderHandle().arena.allocator(), ty.dtype, ty.dims);
14 const count = try owned.elementCount();
15 return switch (owned.dtype) {
16 .f16 => blk: {
17 const payload = try layer.builderHandle().arena.allocator().alloc(f16, count);
18 const value: f16 = @floatCast(fill_value);
19 for (payload) |*slot| slot.* = value;
20 break :blk constantBytes(layer, owned, std.mem.sliceAsBytes(payload));
21 },
22 .bf16 => blk: {
23 const Bf16 = trace.DType.bf16.ZigType();
24 const payload = try layer.builderHandle().arena.allocator().alloc(Bf16, count);
25 const value = Bf16.fromF32(@floatCast(fill_value));
26 for (payload) |*slot| slot.* = value;
27 break :blk constantBytes(layer, owned, std.mem.sliceAsBytes(payload));
28 },
29 .f32 => blk: {
30 const payload = try layer.builderHandle().arena.allocator().alloc(f32, count);
31 const value: f32 = @floatCast(fill_value);
32 for (payload) |*slot| slot.* = value;
33 break :blk constantBytes(layer, owned, std.mem.sliceAsBytes(payload));
34 },
35 .f64 => blk: {
36 const payload = try layer.builderHandle().arena.allocator().alloc(f64, count);
37 for (payload) |*slot| slot.* = fill_value;
38 break :blk constantBytes(layer, owned, std.mem.sliceAsBytes(payload));
39 },
40 else => error.NonFloatDType,
41 };
42 }
43
44 pub fn constantBytes(layer: anytype, ty: trace.Type, payload: []const u8) !trace.Value {
45 var op = program_mod.Operation{
46 .id = program_mod.synthetic_id,
47 .result = ty,
48 .kind = .{ .constant = .{ .payload = payload } },
49 };
50 return layer.next.bind(&op, &.{});
51 }
52
53 pub fn unary(layer: anytype, op_kind: program_mod.Unary, input: trace.Value) !trace.Value {
54 var op = program_mod.Operation{
55 .id = program_mod.synthetic_id,
56 .result = input.ty,
57 .kind = .{ .unary = .{ .op = op_kind, .input = program_mod.synthetic_id } },
58 };
59 return layer.next.bind(&op, &.{input});
60 }
61
62 pub fn binary(layer: anytype, op_kind: program_mod.Binary, lhs: trace.Value, rhs: trace.Value) !trace.Value {
63 var op = program_mod.Operation{
64 .id = program_mod.synthetic_id,
65 .result = lhs.ty,
66 .kind = .{ .binary = .{ .op = op_kind, .lhs = program_mod.synthetic_id, .rhs = program_mod.synthetic_id } },
67 };
68 return layer.next.bind(&op, &.{ lhs, rhs });
69 }
70
71 pub fn compare(layer: anytype, direction: program_mod.CompareDirection, lhs: trace.Value, rhs: trace.Value) !trace.Value {
72 var op = program_mod.Operation{
73 .id = program_mod.synthetic_id,
74 .result = .{ .dtype = .i1, .dims = lhs.ty.dims },
75 .kind = .{ .compare = .{
76 .lhs = program_mod.synthetic_id,
77 .rhs = program_mod.synthetic_id,
78 .direction = direction,
79 } },
80 };
81 return layer.next.bind(&op, &.{ lhs, rhs });
82 }
83
84 pub fn select(layer: anytype, pred: trace.Value, on_true: trace.Value, on_false: trace.Value) !trace.Value {
85 const ty = try types.select(layer.builderHandle().arena.allocator(), pred.ty, on_true.ty, on_false.ty);
86 var op = program_mod.Operation{
87 .id = program_mod.synthetic_id,
88 .result = ty,
89 .kind = .{ .select = .{
90 .pred = program_mod.synthetic_id,
91 .on_true = program_mod.synthetic_id,
92 .on_false = program_mod.synthetic_id,
93 } },
94 };
95 return layer.next.bind(&op, &.{ pred, on_true, on_false });
96 }
97
98 test "tensor emit select rejects non-boolean predicates before binding" {
99 var builder = try trace.Builder.init(std.testing.allocator, "emit_select_predicate");
100 defer builder.deinit();
101
102 const pred = try builder.input(.f32, .{ .lane = 4 });
103 const on_true = try builder.input(.f32, .{ .lane = 4 });
104 const on_false = try builder.input(.f32, .{ .lane = 4 });
105 var layer = struct {
106 builder: *trace.Builder,
107 next: struct {
108 pub fn bind(_: *@This(), _: *const program_mod.Operation, _: []const trace.Value) !trace.Value {
109 return error.UnexpectedBind;
110 }
111 } = .{},
112
113 pub fn builderHandle(self: *@This()) *trace.Builder {
114 return self.builder;
115 }
116 }{ .builder = &builder };
117
118 try std.testing.expectError(error.DTypeMismatch, select(&layer, pred, on_true, on_false));
119 }
120
121 pub fn broadcast(layer: anytype, input: trace.Value, result_dims: []const trace.Dim) !trace.Value {
122 const ty = trace.Type{ .dtype = input.ty.dtype, .dims = result_dims };
123 var op = program_mod.Operation{
124 .id = program_mod.synthetic_id,
125 .result = ty,
126 .kind = .{ .broadcast = .{
127 .input = program_mod.synthetic_id,
128 .sizes = try types.extents(layer.builderHandle().arena.allocator(), result_dims),
129 } },
130 };
131 return layer.next.bind(&op, &.{input});
132 }
133
134 pub fn broadcastInDim(layer: anytype, input: trace.Value, result_dims: []const trace.Dim, broadcast_dims: []const i64) !trace.Value {
135 const ty = trace.Type{ .dtype = input.ty.dtype, .dims = result_dims };
136 var op = program_mod.Operation{
137 .id = program_mod.synthetic_id,
138 .result = ty,
139 .kind = .{ .broadcast_in_dim = .{ .input = program_mod.synthetic_id, .broadcast_dims = broadcast_dims } },
140 };
141 return layer.next.bind(&op, &.{input});
142 }
143
144 pub fn reshape(layer: anytype, input: trace.Value, new_dims: []const trace.Dim) !trace.Value {
145 const allocator = layer.builderHandle().arena.allocator();
146 var op = program_mod.Operation{
147 .id = program_mod.synthetic_id,
148 .result = try types.reshaped(allocator, input.ty, new_dims),
149 .kind = .{ .reshape = .{
150 .input = program_mod.synthetic_id,
151 .new_shape = try types.extents(allocator, new_dims),
152 } },
153 };
154 return layer.next.bind(&op, &.{input});
155 }
156
157 pub fn transpose(layer: anytype, input: trace.Value, permutation: []const i64) !trace.Value {
158 const result_dims = try types.permuted(layer.builderHandle().arena.allocator(), input.ty.dims, permutation);
159 var op = program_mod.Operation{
160 .id = program_mod.synthetic_id,
161 .result = .{ .dtype = input.ty.dtype, .dims = result_dims },
162 .kind = .{ .transpose = .{ .input = program_mod.synthetic_id, .permutation = permutation } },
163 };
164 return layer.next.bind(&op, &.{input});
165 }
166
167 pub fn reduce(layer: anytype, input: trace.Value, init: trace.Value, reducer: program_mod.Reducer, dimensions: []const i64) !trace.Value {
168 const result_dims = try types.removeAxes(layer.builderHandle().arena.allocator(), input.ty.dims, dimensions);
169 var op = program_mod.Operation{
170 .id = program_mod.synthetic_id,
171 .result = .{ .dtype = input.ty.dtype, .dims = result_dims },
172 .kind = .{ .reduce = .{ .input = program_mod.synthetic_id, .init = program_mod.synthetic_id, .reducer = reducer, .dimensions = dimensions } },
173 };
174 return layer.next.bind(&op, &.{ input, init });
175 }
176
177 pub fn iota(layer: anytype, dtype: trace.DType, result_dims: []const trace.Dim, axis: i64) !trace.Value {
178 var op = program_mod.Operation{
179 .id = program_mod.synthetic_id,
180 .result = .{ .dtype = dtype, .dims = result_dims },
181 .kind = .{ .iota = .{ .axis = axis } },
182 };
183 return layer.next.bind(&op, &.{});
184 }
185
186 pub fn gather(layer: anytype, input: trace.Value, indices: trace.Value, axis: i64) !trace.Value {
187 const ty = try types.gather(layer.builderHandle().arena.allocator(), input.ty, indices.ty, axis);
188 var op = program_mod.Operation{
189 .id = program_mod.synthetic_id,
190 .result = ty,
191 .kind = .{ .gather = .{
192 .input = program_mod.synthetic_id,
193 .indices = program_mod.synthetic_id,
194 .axis = axis,
195 } },
196 };
197 return layer.next.bind(&op, &.{ input, indices });
198 }
199
200 pub fn scatterAdd(layer: anytype, input: trace.Value, indices: trace.Value, updates: trace.Value, axis: i64) !trace.Value {
201 const ty = try types.scatterAdd(layer.builderHandle().arena.allocator(), input.ty, indices.ty, updates.ty, axis);
202 var op = program_mod.Operation{
203 .id = program_mod.synthetic_id,
204 .result = ty,
205 .kind = .{ .scatter_add = .{
206 .input = program_mod.synthetic_id,
207 .indices = program_mod.synthetic_id,
208 .updates = program_mod.synthetic_id,
209 .axis = axis,
210 } },
211 };
212 return layer.next.bind(&op, &.{ input, indices, updates });
213 }
214
215 pub fn sparseCrossEntropy(layer: anytype, logits: trace.Value, targets: trace.Value, axis: i64) !trace.Value {
216 const ty = try types.sparseCrossEntropy(layer.builderHandle().arena.allocator(), logits.ty, targets.ty, axis);
217 var op = program_mod.Operation{
218 .id = program_mod.synthetic_id,
219 .result = ty,
220 .kind = .{ .sparse_cross_entropy = .{
221 .logits = program_mod.synthetic_id,
222 .targets = program_mod.synthetic_id,
223 .axis = axis,
224 } },
225 };
226 return layer.next.bind(&op, &.{ logits, targets });
227 }
228
229 pub fn dotGeneral(
230 layer: anytype,
231 lhs: trace.Value,
232 rhs: trace.Value,
233 lhs_contract: []const i64,
234 rhs_contract: []const i64,
235 lhs_batch: []const i64,
236 rhs_batch: []const i64,
237 ) !trace.Value {
238 const result_dims = try types.dotGeneralDims(
239 layer.builderHandle().arena.allocator(),
240 lhs.ty.dims,
241 rhs.ty.dims,
242 lhs_contract,
243 rhs_contract,
244 lhs_batch,
245 rhs_batch,
246 );
247 var op = program_mod.Operation{
248 .id = program_mod.synthetic_id,
249 .result = .{ .dtype = lhs.ty.dtype, .dims = result_dims },
250 .kind = .{
251 .dot_general = .{
252 .lhs = program_mod.synthetic_id,
253 .rhs = program_mod.synthetic_id,
254 .lhs_contract = lhs_contract,
255 .rhs_contract = rhs_contract,
256 .lhs_batch = lhs_batch,
257 .rhs_batch = rhs_batch,
258 },
259 },
260 };
261 return layer.next.bind(&op, &.{ lhs, rhs });
262 }