lib/accy/src/choir/gpu.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

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