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