lib/accy/src/choir/dispatch.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const choir = @import("choir");
 3 const contract = @import("contract.zig");
 4 const semantic = @import("semantic.zig");
 5 const tensor = @import("tensor.zig");
 6 
 7 const ir = choir.ir;
 8 const passes = choir.passes;
 9 
10 pub const product_name = "accy.dispatch";
11 
12 pub const DispatchModule = @import("root.zig").publication.Module(.dispatch);
13 
14 pub const DispatchJob = struct {
15     allocator: std.mem.Allocator,
16     tensor_module: *tensor.TensorJob,
17     choir_module: *ir.Operation,
18     observed_fingerprint: u64,
19 
20     pub fn init(
21         allocator: std.mem.Allocator,
22         tensor_module: *tensor.TensorJob,
23     ) !*DispatchJob {
24         const module = try allocator.create(DispatchJob);
25         errdefer allocator.destroy(module);
26         module.* = .{
27             .allocator = allocator,
28             .tensor_module = tensor_module,
29             .choir_module = tensor_module.choir_module,
30             .observed_fingerprint = 0,
31         };
32         try module.verify();
33         module.observed_fingerprint = try choir.operationFingerprint(allocator, module.choir_module);
34         return module;
35     }
36 
37     pub fn context(self: *DispatchJob) *ir.Context {
38         return self.tensor_module.context();
39     }
40 
41     pub fn deinit(self: *DispatchJob) void {
42         self.tensor_module.deinit();
43         const allocator = self.allocator;
44         allocator.destroy(self);
45     }
46 
47     pub fn verify(self: *DispatchJob) !void {
48         try self.tensor_module.verify();
49     }
50 
51     pub fn fingerprint(self: *const DispatchJob) u64 {
52         return self.observed_fingerprint;
53     }
54 
55     pub fn passContext(self: *DispatchJob) passes.PassContext {
56         return self.tensor_module.passContext();
57     }
58 };
59 
60 test "dispatch module owns tensor product" {
61     const allocator = std.testing.allocator;
62 
63     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
64     defer builder.deinit();
65 
66     const ty = try builder.tensor(.f32, &.{4});
67     var function = try builder.beginFunction("dispatch_add", &.{ ty, ty }, &.{ty});
68     const sum = try function.add(function.parameter(0), function.parameter(1));
69     try function.return_(&.{sum});
70     try function.finish();
71 
72     const semantic_module = try builder.finish();
73     var contract_module = try contract.ContractJob.init(allocator, semantic_module);
74     var contract_owned = true;
75     errdefer if (contract_owned) contract_module.deinit();
76 
77     var analysis_cache = passes.AnalysisCache.init(allocator, null);
78     var cache_owned = true;
79     errdefer if (cache_owned) analysis_cache.deinit();
80 
81     var tensor_module = try tensor.TensorJob.init(allocator, contract_module, analysis_cache);
82     contract_owned = false;
83     cache_owned = false;
84     var tensor_owned = true;
85     errdefer if (tensor_owned) tensor_module.deinit();
86 
87     var dispatch_module = try DispatchJob.init(allocator, tensor_module);
88     tensor_owned = false;
89     defer dispatch_module.deinit();
90 
91     try std.testing.expectEqualStrings(product_name, "accy.dispatch");
92     try dispatch_module.verify();
93     try std.testing.expectEqual(try choir.operationFingerprint(allocator, dispatch_module.choir_module), dispatch_module.fingerprint());
94 }