lib/accy/src/target/module.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const accy_choir = @import("../choir/root.zig");
4 const kernelization = @import("../preparation/kernelization/model/root.zig");
5 const contract = accy_choir.contract;
6 const dispatch = accy_choir.dispatch;
7 const kernel_product = accy_choir.gpu;
8 const memory = accy_choir.memory;
9 const semantic = accy_choir.semantic;
10 const tensor = accy_choir.tensor;
11
12 const ir = choir.ir;
13 const passes = choir.passes;
14
15 pub const product_name = "accy.target";
16
17 pub const TargetModule = accy_choir.publication.Module(.target);
18
19 pub const TargetJob = struct {
20 allocator: std.mem.Allocator,
21 kernel_module: *kernel_product.KernelJob,
22 choir_module: *ir.Operation,
23 analysis_cache: passes.AnalysisCache,
24 kernelization_product: *const kernelization.KernelizationAnalysis,
25 observed_fingerprint: u64,
26
27 pub fn init(
28 allocator: std.mem.Allocator,
29 kernel_module: *kernel_product.KernelJob,
30 choir_module: *ir.Operation,
31 analysis_cache: passes.AnalysisCache,
32 kernelization_product: *const kernelization.KernelizationAnalysis,
33 ) !*TargetJob {
34 const target_module = try allocator.create(TargetJob);
35 errdefer allocator.destroy(target_module);
36 target_module.* = .{
37 .allocator = allocator,
38 .kernel_module = kernel_module,
39 .choir_module = choir_module,
40 .analysis_cache = analysis_cache,
41 .kernelization_product = kernelization_product,
42 .observed_fingerprint = 0,
43 };
44 try target_module.verify();
45 target_module.observed_fingerprint = try choir.operationFingerprint(allocator, target_module.choir_module);
46 return target_module;
47 }
48
49 pub fn context(self: *TargetJob) *ir.Context {
50 return self.kernel_module.context();
51 }
52
53 pub fn deinit(self: *TargetJob) void {
54 self.analysis_cache.deinit();
55 if (self.choir_module != self.kernel_module.choir_module) self.choir_module.erase();
56 self.kernel_module.deinit();
57 const allocator = self.allocator;
58 allocator.destroy(self);
59 }
60
61 pub fn verify(self: *TargetJob) !void {
62 try self.kernel_module.verify();
63 try ir.verifyOperation(self.choir_module, ir.verify.default_options);
64 }
65
66 pub fn fingerprint(self: *const TargetJob) u64 {
67 return self.observed_fingerprint;
68 }
69
70 pub fn kernelizationProduct(self: *const TargetJob) *const kernelization.KernelizationAnalysis {
71 return self.kernelization_product;
72 }
73
74 pub fn passContext(self: *TargetJob) passes.PassContext {
75 return passes.PassContext.init(self.choir_module, self.context(), self.allocator, &self.analysis_cache);
76 }
77 };
78
79 test "target module owns kernel product" {
80 const allocator = std.testing.allocator;
81
82 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
83 defer builder.deinit();
84
85 const ty = try builder.tensor(.f32, &.{4});
86 var function = try builder.beginFunction("target_add", &.{ ty, ty }, &.{ty});
87 const sum = try function.add(function.parameter(0), function.parameter(1));
88 try function.return_(&.{sum});
89 try function.finish();
90
91 const semantic_module = try builder.finish();
92 var contract_module = try contract.ContractJob.init(allocator, semantic_module);
93 var contract_owned = true;
94 errdefer if (contract_owned) contract_module.deinit();
95
96 var analysis_cache = passes.AnalysisCache.init(allocator, null);
97 var cache_owned = true;
98 errdefer if (cache_owned) analysis_cache.deinit();
99
100 var tensor_module = try tensor.TensorJob.init(allocator, contract_module, analysis_cache);
101 contract_owned = false;
102 cache_owned = false;
103 var tensor_owned = true;
104 errdefer if (tensor_owned) tensor_module.deinit();
105
106 var dispatch_module = try dispatch.DispatchJob.init(allocator, tensor_module);
107 tensor_owned = false;
108 var dispatch_owned = true;
109 errdefer if (dispatch_owned) dispatch_module.deinit();
110
111 var memory_module = try memory.MemoryJob.init(allocator, dispatch_module);
112 dispatch_owned = false;
113 var memory_owned = true;
114 errdefer if (memory_owned) memory_module.deinit();
115
116 var kernel_module = try kernel_product.KernelJob.init(allocator, memory_module);
117 memory_owned = false;
118 var kernel_owned = true;
119 errdefer if (kernel_owned) kernel_module.deinit();
120
121 var kernelization_product = try kernelization.KernelizationAnalysis.init(
122 allocator,
123 kernel_module.choir_module.context.capacity.asLimits(),
124 );
125 defer kernelization_product.deinit();
126
127 const target_choir_module = try kernel_module.choir_module.clone();
128 var target_choir_module_owned = true;
129 errdefer if (target_choir_module_owned) target_choir_module.erase();
130
131 var target_analysis_cache = passes.AnalysisCache.init(allocator, null);
132 var target_cache_owned = true;
133 errdefer if (target_cache_owned) target_analysis_cache.deinit();
134
135 var target_module = try TargetJob.init(
136 allocator,
137 kernel_module,
138 target_choir_module,
139 target_analysis_cache,
140 &kernelization_product,
141 );
142 kernel_owned = false;
143 target_choir_module_owned = false;
144 target_cache_owned = false;
145 defer target_module.deinit();
146
147 try std.testing.expectEqualStrings(product_name, "accy.target");
148 try target_module.verify();
149 try std.testing.expect(target_module.choir_module != target_module.kernel_module.choir_module);
150 try std.testing.expectEqual(try choir.operationFingerprint(allocator, target_module.choir_module), target_module.fingerprint());
151 try std.testing.expectEqual(
152 target_module.kernel_module.fingerprint(),
153 try choir.operationFingerprint(allocator, target_module.kernel_module.choir_module),
154 );
155 try std.testing.expectEqual(@as(usize, 0), target_module.kernelizationProduct().kernelCount());
156 }