lib/accy/src/choir/publication.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const operation = choir.product.operation;
4 const revision = choir.product.revision;
5
6 pub const Stage = enum(u8) {
7 semantic,
8 contract,
9 tensor,
10 dispatch,
11 memory,
12 kernel,
13 target,
14
15 pub fn name(self: Stage) []const u8 {
16 return switch (self) {
17 .semantic => "accy.semantic",
18 .contract => "accy.contract",
19 .tensor => "accy.tensor_opt",
20 .dispatch => "accy.dispatch",
21 .memory => "accy.memory",
22 .kernel => "accy.kernel/gpu",
23 .target => "accy.target",
24 };
25 }
26
27 pub fn schema(comptime self: Stage) revision.record.Version {
28 return .{ .name = "accy-" ++ @tagName(self) ++ "-record", .version = 1 };
29 }
30 };
31
32 /// An opaque handle holds one reference to one stage record of the chosen
33 /// stage, together with the allocator that made the handle, so a caller can
34 /// pass it on or reopen it with the stage checked at compile time.
35 /// `fromRevision` returns `error.WrongStage` when the record belongs to another
36 /// stage, checks the record's schema gates, and then takes a reference. The
37 /// handle holds no compiler state and no cache, and those live in the job that
38 /// runs the stage. `deinit` drops the reference and frees the handle, and
39 /// `retain` makes a second handle over the same record. `plan` decodes the
40 /// stage's typed plan for the dispatch, memory, kernel and target stages, and
41 /// asking for it on another stage is a compile error.
42 pub fn Module(comptime stage: Stage) type {
43 return opaque {
44 const Self = @This();
45 const Data = struct { allocator: std.mem.Allocator, product: operation.Product };
46
47 pub fn fromRevision(
48 allocator: std.mem.Allocator,
49 published: *const revision.Revision,
50 ) !*Self {
51 if (!std.mem.eql(u8, published.address().stage, stage.name())) {
52 return error.WrongStage;
53 }
54 try published.requireGates(&.{ operation.schema_identity, stage.schema() });
55 const retained = try published.retain();
56 errdefer retained.release();
57 const data = try allocator.create(Data);
58 data.* = .{ .allocator = allocator, .product = .{ .revision = retained } };
59 return @ptrCast(data);
60 }
61
62 pub fn deinit(self: *Self) void {
63 const data: *Data = @ptrCast(@alignCast(self));
64 data.product.release();
65 data.allocator.destroy(data);
66 }
67
68 pub fn retain(self: *const Self, allocator: std.mem.Allocator) !*Self {
69 return fromRevision(allocator, self.record());
70 }
71
72 pub fn record(self: *const Self) *const revision.Revision {
73 const data: *const Data = @ptrCast(@alignCast(self));
74 return data.product.revision;
75 }
76
77 pub fn plan(
78 self: *const Self,
79 allocator: std.mem.Allocator,
80 limits: choir.bytecode.image.Limits,
81 ) !@import("root.zig").record.codec.Decoded(Plan(stage)) {
82 if (Plan(stage) == void) @compileError("this stage has no separate plan record");
83 const image = try self.open(allocator, limits);
84 defer image.destroy();
85 return @import("root.zig").record.codec.decode(
86 allocator,
87 Plan(stage),
88 stage,
89 image.stage(),
90 );
91 }
92
93 pub fn eql(self: *const Self, other: *const Self) bool {
94 return self.record().eql(other.record());
95 }
96
97 pub fn open(
98 self: *const Self,
99 allocator: std.mem.Allocator,
100 limits: choir.bytecode.image.Limits,
101 ) !*choir.product.entity.Image {
102 const data: *const Data = @ptrCast(@alignCast(self));
103 return data.product.open(allocator, limits);
104 }
105 };
106 }
107
108 /// Returns five fixed bytes, a version of 1 and the stage's number, for the
109 /// semantic, contract and tensor stages, and any other stage is a compile
110 /// error. The compiler writes these bytes as the stage payload of the first
111 /// three stages, and the checker compares against them. The first three stages
112 /// keep only their compiler image and store no plan and no cached analysis, so
113 /// these five bytes are their whole payload. The checker refuses a stage record
114 /// of one of those stages with `error.InvalidStageRecord` when its payload
115 /// differs from these bytes.
116 pub fn irRecord(comptime stage: Stage) [5]u8 {
117 comptime requireIrStage(stage);
118 return .{ 1, 0, 0, 0, @backingInt(stage) };
119 }
120
121 fn requireIrStage(comptime stage: Stage) void {
122 switch (stage) {
123 .semantic, .contract, .tensor => {},
124 else => @compileError("this stage requires its complete typed plan schema"),
125 }
126 }
127
128 pub fn Plan(comptime stage: Stage) type {
129 const records = @import("root.zig").record;
130 return switch (stage) {
131 .dispatch => records.dispatch.Record,
132 .memory => records.memory.Record,
133 .kernel => records.kernel.Record,
134 .target => records.target.Record,
135 else => void,
136 };
137 }