lib/accy/src/kernel/library/catalog/artifact/pipeline.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const artifact_product = @import("../../../../artifact/model/root.zig");
4 const kernel = @import("../../../root.zig");
5 const descriptor_mod = @import("../root.zig");
6 const entry = @import("../../root.zig").entry;
7 const query_mod = descriptor_mod;
8 const scan = @import("../../root.zig").scan;
9 const sort_mod = @import("../../root.zig").sort;
10
11 const model = @import("model.zig");
12
13 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
14 const OwnedKernelCallPipelinePackage = model.OwnedKernelCallPipelinePackage;
15
16 pub fn createOwnedKernelCallPipelinePackage(
17 allocator: std.mem.Allocator,
18 handle: kernel.BackendHandle,
19 descriptor: OwnedDescriptor,
20 options: entry.ArtifactOptions,
21 ) !?OwnedKernelCallPipelinePackage {
22 const metadata = descriptor.descriptor.metadata;
23 return switch (metadata.category) {
24 .scan => try createDeviceScanPipelinePackage(allocator, handle, metadata, options),
25 .sort => try createRadixSortPipelinePackage(allocator, handle, metadata, options),
26 else => null,
27 };
28 }
29
30 fn createDeviceScanPipelinePackage(
31 allocator: std.mem.Allocator,
32 handle: kernel.BackendHandle,
33 metadata: entry.Metadata,
34 options: entry.ArtifactOptions,
35 ) !?OwnedKernelCallPipelinePackage {
36 if (metadata.version != scan.device_scan_family_version) return null;
37 const instance = scan.deviceScanInstanceFromSpecialization(metadata.specialization) orelse return null;
38 const target = try scan.deviceScanFamilyTarget(allocator, instance);
39 defer allocator.free(target);
40 if (!std.mem.eql(u8, metadata.target, target)) return null;
41
42 var stage_artifacts = try scan.createDeviceScanPipelineArtifacts(allocator, handle, instance, options);
43 var stage_artifacts_owned = true;
44 errdefer if (stage_artifacts_owned) stage_artifacts.deinit();
45
46 var pipeline = try scan.deviceScanPipeline(allocator, instance);
47 errdefer pipeline.deinit();
48
49 const artifacts = allocator.alloc(kernel.OwnedKernelCallArtifact, 3) catch return error.OutOfMemory;
50 errdefer allocator.free(artifacts);
51 artifacts[0] = stage_artifacts.block_scan;
52 artifacts[1] = stage_artifacts.sums_scan;
53 artifacts[2] = stage_artifacts.add_base;
54 stage_artifacts_owned = false;
55
56 return try finishPipelinePackage(allocator, artifacts, &pipeline);
57 }
58
59 fn createRadixSortPipelinePackage(
60 allocator: std.mem.Allocator,
61 handle: kernel.BackendHandle,
62 metadata: entry.Metadata,
63 options: entry.ArtifactOptions,
64 ) !?OwnedKernelCallPipelinePackage {
65 if (metadata.specialization.structureIs(sort_mod.bitonic_block_structure_name)) return null;
66 if (metadata.version != sort_mod.radix_split_family_version) return null;
67 const instance = sort_mod.radixSplitInstanceFromSpecialization(metadata.specialization) orelse return null;
68 if (metadata.specialization.structureIs(@tagName(query_mod.SortStructure.radix_digit))) {
69 return try createRadixDigitPackage(allocator, handle, metadata, instance, options);
70 }
71 return try createRadixSplitPackage(allocator, handle, metadata, instance, options);
72 }
73
74 fn createRadixSplitPackage(
75 allocator: std.mem.Allocator,
76 handle: kernel.BackendHandle,
77 metadata: entry.Metadata,
78 instance: sort_mod.RadixSplit,
79 options: entry.ArtifactOptions,
80 ) !?OwnedKernelCallPipelinePackage {
81 const target = try sort_mod.radixSplitPipelineTarget(allocator, instance);
82 defer allocator.free(target);
83 if (!std.mem.eql(u8, metadata.target, target)) return null;
84
85 var stage_artifacts = try sort_mod.createRadixSplitPipelineArtifacts(allocator, handle, instance, options);
86 var stage_artifacts_owned = true;
87 errdefer if (stage_artifacts_owned) stage_artifacts.deinit();
88
89 var pipeline = try sort_mod.radixSplitPipeline(allocator, instance);
90 errdefer pipeline.deinit();
91
92 const artifacts = allocator.alloc(kernel.OwnedKernelCallArtifact, 5) catch return error.OutOfMemory;
93 errdefer allocator.free(artifacts);
94 artifacts[0] = stage_artifacts.flags;
95 artifacts[1] = stage_artifacts.block_scan;
96 artifacts[2] = stage_artifacts.sums_scan;
97 artifacts[3] = stage_artifacts.add_base;
98 artifacts[4] = stage_artifacts.scatter;
99 stage_artifacts_owned = false;
100
101 return try finishPipelinePackage(allocator, artifacts, &pipeline);
102 }
103
104 fn createRadixDigitPackage(
105 allocator: std.mem.Allocator,
106 handle: kernel.BackendHandle,
107 metadata: entry.Metadata,
108 instance: sort_mod.RadixSplit,
109 options: entry.ArtifactOptions,
110 ) !?OwnedKernelCallPipelinePackage {
111 const target = try sort_mod.radixDigitPipelineTarget(allocator, instance);
112 defer allocator.free(target);
113 if (!std.mem.eql(u8, metadata.target, target)) return null;
114
115 var stage_artifacts = try sort_mod.createRadixDigitPipelineArtifacts(allocator, handle, instance, options);
116 var stage_artifacts_owned = true;
117 errdefer if (stage_artifacts_owned) stage_artifacts.deinit();
118
119 var pipeline = try sort_mod.radixDigitPipeline(allocator, instance);
120 errdefer pipeline.deinit();
121
122 const artifacts = allocator.alloc(kernel.OwnedKernelCallArtifact, 5) catch return error.OutOfMemory;
123 errdefer allocator.free(artifacts);
124 artifacts[0] = stage_artifacts.histogram;
125 artifacts[1] = stage_artifacts.block_scan;
126 artifacts[2] = stage_artifacts.sums_scan;
127 artifacts[3] = stage_artifacts.add_base;
128 artifacts[4] = stage_artifacts.rank_scatter;
129 stage_artifacts_owned = false;
130
131 return try finishPipelinePackage(allocator, artifacts, &pipeline);
132 }
133
134 fn finishPipelinePackage(
135 allocator: std.mem.Allocator,
136 artifacts: []kernel.OwnedKernelCallArtifact,
137 pipeline: *artifact_product.OwnedKernelCallPipeline,
138 ) !?OwnedKernelCallPipelinePackage {
139 errdefer for (artifacts) |*artifact| artifact.deinit();
140
141 const entries = allocator.alloc(artifact_product.KernelCallArtifact, artifacts.len) catch return error.OutOfMemory;
142 errdefer allocator.free(entries);
143 for (artifacts, entries) |*owned_artifact, *registry_entry| registry_entry.* = owned_artifact.entry();
144
145 const index = try artifact_product.buildKernelCallRegistryIndex(allocator, entries);
146 errdefer artifact_product.deinitKernelCallRegistryIndex(allocator, index);
147
148 const registry_view = artifact_product.KernelCallRegistry{ .entries = entries, .index = index };
149 try pipeline.value.validate(registry_view, entries[0].format);
150
151 return .{
152 .allocator = allocator,
153 .artifacts = artifacts,
154 .entries = entries,
155 .index = index,
156 .pipeline = pipeline.*,
157 };
158 }