lib/accy/src/tensor/dsl/body.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const tensor = @import("../root.zig");
 3 const trace = tensor.trace;
 4 const parameter = @import("parameter.zig");
 5 
 6 pub fn buildInputs(builder: *trace.Builder, comptime parameters: anytype) ![]trace.Value {
 7     const input_count = comptime parameter.arity(parameters);
 8     const args = try builder.arena.allocator().alloc(trace.Value, input_count);
 9     if (comptime parameter.named(parameters)) {
10         inline for (@typeInfo(@TypeOf(parameters)).@"struct".field_names, 0..) |field_name, index_value| {
11             args[index_value] = try builder.inputSpec(@field(parameters, field_name));
12         }
13     } else {
14         inline for (0..input_count) |index_value| {
15             args[index_value] = try builder.inputSpec(parameters[index_value]);
16         }
17     }
18     return args;
19 }
20 
21 pub fn call(comptime definition: anytype, builder: *trace.Builder, args: []const trace.Value) ![]const trace.Value {
22     const result = if (comptime parameter.named(definition.parameters))
23         try definition.body(builder, parameter.Args(definition.parameters){ .values = args })
24     else
25         try definition.body(builder, args);
26     return collectOutputs(builder, result);
27 }
28 
29 fn collectOutputs(builder: *trace.Builder, result: anytype) ![]const trace.Value {
30     const Result = @TypeOf(result);
31     if (comptime Result == trace.Value) {
32         const outputs = try builder.arena.allocator().alloc(trace.Value, 1);
33         outputs[0] = result;
34         return outputs;
35     }
36 
37     return switch (@typeInfo(Result)) {
38         .array => |array| blk: {
39             if (array.child != trace.Value) @compileError("tensor Program array outputs must contain tensor Values");
40             const outputs = try builder.arena.allocator().alloc(trace.Value, array.len);
41             for (result, 0..) |output, index_value| {
42                 outputs[index_value] = output;
43             }
44             break :blk outputs;
45         },
46         .pointer => |pointer| collectPointerOutputs(builder, result, pointer),
47         .@"struct" => |info| blk: {
48             if (info.field_names.len == 0) @compileError("tensor Program bodies must return at least one output");
49             const outputs = try builder.arena.allocator().alloc(trace.Value, info.field_names.len);
50             inline for (info.field_names, info.field_types, 0..) |field_name, field_type, index_value| {
51                 if (field_type != trace.Value) @compileError("tensor Program struct outputs must contain tensor Values");
52                 outputs[index_value] = @field(result, field_name);
53             }
54             break :blk outputs;
55         },
56         else => @compileError("tensor Program bodies must return a Value, array, slice, or struct of Values"),
57     };
58 }
59 
60 fn collectPointerOutputs(builder: *trace.Builder, result: anytype, comptime pointer: std.builtin.Type.Pointer) ![]const trace.Value {
61     switch (pointer.size) {
62         .slice => {
63             if (pointer.child != trace.Value) @compileError("tensor Program slice outputs must contain tensor Values");
64             const outputs = try builder.arena.allocator().alloc(trace.Value, result.len);
65             for (result, 0..) |output, index_value| {
66                 outputs[index_value] = output;
67             }
68             return outputs;
69         },
70         .one => {
71             return switch (@typeInfo(pointer.child)) {
72                 .array => |array| blk: {
73                     if (array.child != trace.Value) @compileError("tensor Program array outputs must contain tensor Values");
74                     const outputs = try builder.arena.allocator().alloc(trace.Value, array.len);
75                     for (result.*, 0..) |output, index_value| {
76                         outputs[index_value] = output;
77                     }
78                     break :blk outputs;
79                 },
80                 else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),
81             };
82         },
83         else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),
84     }
85 }