lib/accy/src/integration/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const choir = @import("../choir/root.zig");
5
6 const gating = @import("accy_validation_gating");
7 const artifact = @import("../artifact/root.zig");
8 const executable = @import("../executable/root.zig");
9 const kernel = @import("../kernel/root.zig");
10 const preparation = @import("../preparation/root.zig");
11 const target = @import("../target/root.zig");
12 const ir = choir.ir;
13 const dialect = choir.dialect;
14 const activation = choir.activation;
15 const semantics = choir.semantics;
16 const semantic = choir.semantic;
17 const shape = choir.shape;
18 const contract = choir.contract;
19 const tensor = choir.tensor;
20 const einsum = choir.einsum;
21 const dispatch = choir.dispatch;
22 const memory = choir.memory;
23 const AccyDialect = choir.AccyDialect;
24 const accyTensorType = choir.accyTensorType;
25 const tensor_type_name = choir.tensor_type_name;
26 const accy_package_extension = choir.accy_package_extension;
27 const accy_choir_package_extension = preparation.accy_choir_package_extension;
28 const registerAccyDialect = choir.registerAccyDialect;
29 const ActivationKind = choir.ActivationKind;
30 const semantic_product_name = choir.semantic_product_name;
31 const SemanticBuilder = choir.SemanticBuilder;
32 const SemanticModule = choir.SemanticModule;
33 const ShapeFamily = choir.ShapeFamily;
34 const ShapeFamilyBuilder = choir.ShapeFamilyBuilder;
35 const contract_product_name = choir.contract_product_name;
36 const ContractJob = choir.ContractJob;
37 const tensor_product_name = choir.tensor_product_name;
38 const TensorJob = choir.TensorJob;
39 const EinsumEquation = choir.EinsumEquation;
40 const EinsumPlan = choir.EinsumPlan;
41 const parseEinsumEquation = choir.parseEinsumEquation;
42 const createEinsumPlan = choir.createEinsumPlan;
43 const lowerEinsumPlan = choir.lowerEinsumPlan;
44 const dispatch_product_name = choir.dispatch_product_name;
45 const DispatchJob = choir.DispatchJob;
46 const memory_product_name = choir.memory_product_name;
47 const MemoryJob = choir.MemoryJob;
48 const kernel_product_name = choir.kernel_product_name;
49 const KernelJob = choir.KernelJob;
50 const KernelizationProduct = preparation.KernelizationProduct;
51 const target_product_name = target.product_name;
52 const TargetJob = target.TargetJob;
53 const artifact_product_name = artifact.product_name;
54 const ArtifactJob = artifact.ArtifactJob;
55 const exec_product_name = executable.product_name;
56 const CompiledFragment = executable.CompiledFragment;
57 const LaunchCandidateRecord = executable.LaunchCandidateRecord;
58 const LoadedFragment = executable.LoadedFragment;
59 const FragmentCompilerCache = executable.FragmentCompilerCache;
60 const FragmentCompilerCacheUpdate = executable.FragmentCompilerCacheUpdate;
61 const LaunchGraphLoopCarry = executable.LaunchGraphLoopCarry;
62 const LaunchGraphLoop = executable.LaunchGraphLoop;
63 const launchGraphLoopCarryFinalSlot = executable.launchGraphLoopCarryFinalSlot;
64 const ArtifactKernelSource = artifact.KernelSource;
65 const ArtifactKernelSummary = artifact.KernelSummary;
66 const ArtifactKernelSummaries = artifact.KernelSummaries;
67 const FragmentCompilerOptions = executable.FragmentCompilerOptions;
68 const Invocation = executable.Invocation;
69 const InvocationState = executable.InvocationState;
70 const InvocationError = executable.InvocationError;
71 const prepareInvocation = executable.prepareInvocation;
72 const invoke = executable.invoke;
73 const compileFragmentFromArtifactJob = executable.compileFragmentFromArtifactJob;
74 const compileFragmentFromArtifactJobWithLaunchPlan = executable.compileFragmentFromArtifactJobWithLaunchPlan;
75 const loadFragment = executable.loadFragment;
76 const compileFragmentFromPreparedJob = executable.compileFragmentFromPreparedJob;
77 const createArtifactJobFromPreparedJob = executable.createArtifactJobFromPreparedJob;
78 const prepareFragmentFromSemanticModule = executable.prepareFragmentFromSemanticModule;
79 const compileFragmentFromSemanticModule = executable.compileFragmentFromSemanticModule;
80 const BackendPreparationRun = preparation.BackendPreparationRun;
81 const BackendPreparedJob = preparation.BackendPreparedJob;
82 const BackendPreparationProductStamps = preparation.BackendPreparationProductStamps;
83 const BackendPreparationProductKeys = preparation.BackendPreparationProductKeys;
84 const BackendPreparationCache = preparation.BackendPreparationCache;
85 const BackendPreparationCacheUpdate = preparation.BackendPreparationCacheUpdate;
86 const BackendPreparationRunOptions = preparation.BackendPreparationRunOptions;
87 const BackendPreparationStats = preparation.BackendPreparationStats;
88 const BackendPreparationFailureKind = preparation.BackendPreparationFailureKind;
89 const BackendPreparationFailure = preparation.BackendPreparationFailure;
90 const BackendPreparationTiming = preparation.BackendPreparationTiming;
91 const GeneratedScheduleKind = preparation.GeneratedScheduleKind;
92 const GeneratedSchedule = preparation.GeneratedSchedule;
93 const GeneratedScanSchedule = preparation.target.GeneratedScanSchedule;
94 const GeneratedScanScheduleDecision = preparation.target.GeneratedScanScheduleDecision;
95 const generated_scan_schedule_version = preparation.kernelization.generated_scan_schedule_version;
96 const max_scan_schedule_candidates = preparation.kernelization.max_scan_schedule_candidates;
97 const scanScheduleCandidates = preparation.kernelization.scanScheduleCandidates;
98 const GeneratedRowPipelineSchedule = preparation.target.GeneratedRowPipelineSchedule;
99 const GeneratedRowPipelineScheduleDecision = preparation.target.GeneratedRowPipelineScheduleDecision;
100 const generated_row_pipeline_schedule_version = preparation.kernelization.generated_row_pipeline_schedule_version;
101 const max_row_pipeline_schedule_candidates = preparation.kernelization.max_row_pipeline_schedule_candidates;
102 const rowPipelineScheduleCandidates = preparation.kernelization.rowPipelineScheduleCandidates;
103 const GeneratedKernelProgram = preparation.GeneratedKernelProgram;
104 const GeneratedKernelSummary = preparation.GeneratedKernelSummary;
105 const GeneratedKernelSummaries = preparation.GeneratedKernelSummaries;
106 const KernelLibraryLowering = preparation.KernelLibraryLowering;
107 const ActivationLoweringOptions = preparation.ActivationLoweringOptions;
108 const EinsumLoweringOptions = preparation.EinsumLoweringOptions;
109 const IndexingLoweringOptions = preparation.IndexingLoweringOptions;
110 const TensorLoweringOptions = preparation.TensorLoweringOptions;
111 const canonicalization_pattern_entries = preparation.canonicalization_pattern_entries;
112 const contract_pass_plan = preparation.contract_pass_plan;
113 const contract_pipeline_name = preparation.contract_pipeline_name;
114 const tensor_pass_plan = preparation.tensor_pass_plan;
115 const tensor_pipeline_name = preparation.tensor_pipeline_name;
116 const dispatch_pass_plan = preparation.dispatch_pass_plan;
117 const dispatch_pipeline_name = preparation.dispatch_pipeline_name;
118 const memory_pass_plan = preparation.memory_pass_plan;
119 const memory_pipeline_name = preparation.memory_pipeline_name;
120 const kernel_pass_plan = preparation.kernel_pass_plan;
121 const kernel_pipeline_name = preparation.kernel_pipeline_name;
122 const target_pass_plan = preparation.target_pass_plan;
123 const target_pipeline_name = preparation.target_pipeline_name;
124 const contract_pass_count = preparation.contract_pass_count;
125 const contract_analysis_count = preparation.contract_analysis_count;
126 const tensor_pass_count = preparation.tensor_pass_count;
127 const tensor_analysis_count = preparation.tensor_analysis_count;
128 const dispatch_pass_count = preparation.dispatch_pass_count;
129 const dispatch_analysis_count = preparation.dispatch_analysis_count;
130 const memory_pass_count = preparation.memory_pass_count;
131 const memory_analysis_count = preparation.memory_analysis_count;
132 const kernel_pass_count = preparation.kernel_pass_count;
133 const kernel_analysis_count = preparation.kernel_analysis_count;
134 const target_pass_count = preparation.target_pass_count;
135 const contractPassName = preparation.contractPassName;
136 const contractAnalysisName = preparation.contractAnalysisName;
137 const tensorPassName = preparation.tensorPassName;
138 const tensorAnalysisName = preparation.tensorAnalysisName;
139 const dispatchPassName = preparation.dispatchPassName;
140 const dispatchAnalysisName = preparation.dispatchAnalysisName;
141 const memoryPassName = preparation.memoryPassName;
142 const memoryAnalysisName = preparation.memoryAnalysisName;
143 const kernelPassName = preparation.kernelPassName;
144 const kernelAnalysisName = preparation.kernelAnalysisName;
145 const target_analysis_count = preparation.target_analysis_count;
146 const targetPassName = preparation.targetPassName;
147 const targetAnalysisName = preparation.targetAnalysisName;
148 const prepareContractJobFromSemanticModule = preparation.prepareContractJobFromSemanticModule;
149 const prepareTensorJobFromContractJob = preparation.prepareTensorJobFromContractJob;
150 const prepareDispatchJobFromTensorJob = preparation.prepareDispatchJobFromTensorJob;
151 const prepareMemoryJobFromDispatchJob = preparation.prepareMemoryJobFromDispatchJob;
152 const prepareKernelJobFromMemoryJob = preparation.prepareKernelJobFromMemoryJob;
153 const prepareTargetJobFromKernelJob = preparation.prepareTargetJobFromKernelJob;
154 const prepareBackendJobFromSemanticModule = preparation.prepareBackendJobFromSemanticModule;
155 const activationLoweringPass = preparation.activationLoweringPass;
156 const activationLoweringPassWithOptions = preparation.activationLoweringPassWithOptions;
157 const einsumLoweringPass = preparation.einsumLoweringPass;
158 const einsumLoweringPassWithOptions = preparation.einsumLoweringPassWithOptions;
159 const indexingLoweringPass = preparation.indexingLoweringPass;
160 const indexingLoweringPassWithOptions = preparation.indexingLoweringPassWithOptions;
161 const createArtifactJobFromTargetJob = artifact.createArtifactJobFromTargetJob;
162 const runBackendPreparationPipelineFromSemanticModule = preparation.runBackendPreparationPipelineFromSemanticModule;
163 const PipelineError = preparation.PipelineError;
164
165 test {
166 _ = @import("../artifact/test.zig");
167 _ = @import("../eval/test.zig");
168 _ = @import("../executable/test.zig");
169 _ = @import("../choir/einsum/test.zig");
170 _ = @import("../kernel/root.zig");
171 _ = @import("../preparation/test.zig");
172 _ = @import("../choir/shape/test.zig");
173 _ = @import("../target/root.zig");
174 @import("test_discovery").discover(choir);
175 }
176
177 fn compileAndLoadSemanticModule(
178 allocator: std.mem.Allocator,
179 handle: gpu.BackendHandle,
180 module: *SemanticModule,
181 options: FragmentCompilerOptions,
182 ) !*LoadedFragment {
183 const compiled = try compileFragmentFromSemanticModule(allocator, handle, module, options);
184 return try loadFragment(allocator, handle, compiled, options);
185 }
186 const KernelCallAdd = kernel.library.elementwise.VectorAdd8F32;
187 const KernelCallGelu = kernel.library.elementwise.Gelu8F32;
188 const KernelCallRelu = kernel.library.elementwise.Relu8F32;
189 const KernelCallMatrixProduct = kernel.library.linalg.MatrixProduct4x16x8F32;
190 const KernelCallTranspose = kernel.library.layout.Transpose8x16F32;
191 const KernelCallDot = kernel.library.reduction.Dot8F32;
192
193 fn semanticKernelCallAddModule(allocator: std.mem.Allocator) !*SemanticModule {
194 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
195 errdefer builder.deinit();
196 const f32_8 = try builder.tensor(.f32, &.{8});
197 var fb = try builder.beginFunction("semantic_kernel_call_add", &.{ f32_8, f32_8 }, &.{f32_8});
198 const call = try fb.kernelCall(
199 &.{ fb.parameter(0), fb.parameter(1) },
200 &.{f32_8},
201 .{
202 .target = KernelCallAdd.target,
203 .operand_effects = &.{ .read, .read },
204 .result_aliases = &.{null},
205 },
206 );
207 try fb.return_(&.{call.getFirstResult()});
208 try fb.finish();
209 return try builder.finish();
210 }
211
212 fn semanticEinsumMatrixProductModule(allocator: std.mem.Allocator) !*SemanticModule {
213 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
214 errdefer builder.deinit();
215 const lhs_ty = try builder.tensor(.f32, &.{ 4, 8 });
216 const rhs_ty = try builder.tensor(.f32, &.{ 8, 16 });
217 const out_ty = try builder.tensor(.f32, &.{ 4, 16 });
218 var fb = try builder.beginFunction("semantic_einsum_kernel_call_matmul", &.{ lhs_ty, rhs_ty }, &.{out_ty});
219 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "ik,kj->ij");
220 try fb.return_(&.{out});
221 try fb.finish();
222 return try builder.finish();
223 }
224
225 fn semanticEinsumMatrixProductFamilyModule(allocator: std.mem.Allocator) !*SemanticModule {
226 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
227 errdefer builder.deinit();
228 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
229 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
230 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
231 var fb = try builder.beginFunction("semantic_einsum_kernel_call_matmul_family", &.{ lhs_ty, rhs_ty }, &.{out_ty});
232 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "ik,kj->ij");
233 try fb.return_(&.{out});
234 try fb.finish();
235 return try builder.finish();
236 }
237
238 fn semanticSegmentSumFamilyModule(allocator: std.mem.Allocator, family_target: []const u8) !*SemanticModule {
239 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
240 errdefer builder.deinit();
241 const data_ty = try builder.tensor(.f32, &.{12});
242 const offsets_ty = try builder.tensor(.i32, &.{4});
243 const out_ty = try builder.tensor(.f32, &.{3});
244 var fb = try builder.beginFunction("semantic_segment_sum_kernel_call_family", &.{ data_ty, offsets_ty }, &.{out_ty});
245 const call = try fb.kernelCall(
246 &.{ fb.parameter(0), fb.parameter(1) },
247 &.{out_ty},
248 .{
249 .target = family_target,
250 .version = kernel.library.segmented.segment_sum_family_version,
251 .operand_effects = &.{ .read, .read },
252 .result_aliases = &.{null},
253 },
254 );
255 try fb.return_(&.{call.getFirstResult()});
256 try fb.finish();
257 return try builder.finish();
258 }
259
260 fn semanticGatherFamilyModule(allocator: std.mem.Allocator) !*SemanticModule {
261 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
262 errdefer builder.deinit();
263 const data_ty = try builder.tensor(.f32, &.{ 2, 8, 3 });
264 const indices_ty = try builder.tensor(.i32, &.{5});
265 const out_ty = try builder.tensor(.f32, &.{ 2, 5, 3 });
266 var fb = try builder.beginFunction("semantic_gather_kernel_call_family", &.{ data_ty, indices_ty }, &.{out_ty});
267 const out = try fb.gather(fb.parameter(0), fb.parameter(1), out_ty, 1);
268 try fb.return_(&.{out});
269 try fb.finish();
270 return try builder.finish();
271 }
272
273 fn semanticRank2IotaGatherFusionModule(allocator: std.mem.Allocator) !*SemanticModule {
274 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
275 errdefer builder.deinit();
276 const flat_ty = try builder.tensor(.f32, &.{256});
277 const img_ty = try builder.tensor(.f32, &.{ 8, 32 });
278 const idx_ty = try builder.tensor(.i32, &.{ 8, 32 });
279 const scalar_ty = try builder.tensor(.f32, &.{});
280 var fb = try builder.beginFunction("semantic_rank2_iota_gather_fusion", &.{flat_ty}, &.{img_ty});
281 const indices = try fb.iota(idx_ty, 1);
282 const gathered = try fb.gather(fb.parameter(0), indices, img_ty, 0);
283 const zero: f32 = 0;
284 const scalar = try fb.constant(scalar_ty, std.mem.asBytes(&zero));
285 const zeros = try fb.broadcast(scalar, img_ty, &.{ 8, 32 });
286 const out = try fb.add(gathered, zeros);
287 try fb.return_(&.{out});
288 try fb.finish();
289 return try builder.finish();
290 }
291
292 fn semanticRowPipelineModule(allocator: std.mem.Allocator, rows: i64, cols: i64) !*SemanticModule {
293 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
294 errdefer builder.deinit();
295 const x_ty = try builder.tensor(.f32, &.{ rows, cols });
296 const row_ty = try builder.tensor(.f32, &.{rows});
297 const scalar_ty = try builder.tensor(.f32, &.{});
298 var fb = try builder.beginFunction("semantic_row_pipeline_schedule", &.{x_ty}, &.{x_ty});
299 const x = fb.parameter(0);
300 const lowest = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, -std.math.floatMax(f32))));
301 const row_max = try fb.reduce(x, lowest, row_ty, "max", &.{1});
302 const shifted = try fb.sub(x, try fb.broadcastInDim(row_max, x_ty, &.{ rows, cols }, &.{0}));
303 const exps = try fb.exp(shifted);
304 const zero = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
305 const row_sum = try fb.reduce(exps, zero, row_ty, "sum", &.{1});
306 const out = try fb.div(exps, try fb.broadcastInDim(row_sum, x_ty, &.{ rows, cols }, &.{0}));
307 try fb.return_(&.{out});
308 try fb.finish();
309 return try builder.finish();
310 }
311
312 fn semanticCumsumModule(allocator: std.mem.Allocator, total: i64) !*SemanticModule {
313 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
314 errdefer builder.deinit();
315 const vec_ty = try builder.tensor(.f32, &.{total});
316 var fb = try builder.beginFunction("semantic_cumsum_schedule", &.{vec_ty}, &.{vec_ty});
317 const out = try fb.cumsum(fb.parameter(0), vec_ty, 0);
318 try fb.return_(&.{out});
319 try fb.finish();
320 return try builder.finish();
321 }
322
323 fn semanticRank2ClampedGatherFusionModule(allocator: std.mem.Allocator) !*SemanticModule {
324 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
325 errdefer builder.deinit();
326 const rows: i64 = 8;
327 const cols: i64 = 8;
328 const flat_ty = try builder.tensor(.f32, &.{rows * cols});
329 const img_ty = try builder.tensor(.f32, &.{ rows, cols });
330 const idx_ty = try builder.tensor(.i32, &.{ rows, cols });
331 const scalar_ty = try builder.tensor(.f32, &.{});
332 var fb = try builder.beginFunction("semantic_rank2_clamped_gather_fusion", &.{flat_ty}, &.{img_ty});
333
334 const Consts = struct {
335 fb: *@TypeOf(fb),
336 img_ty: ir.Type,
337 scalar_ty: ir.Type,
338 rows: i64,
339 cols: i64,
340
341 fn splat(self: *@This(), value: f32) !*ir.Value {
342 const scalar = try self.fb.constant(self.scalar_ty, std.mem.asBytes(&value));
343 return self.fb.broadcast(scalar, self.img_ty, &.{ self.rows, self.cols });
344 }
345 };
346 var consts = Consts{ .fb = &fb, .img_ty = img_ty, .scalar_ty = scalar_ty, .rows = rows, .cols = cols };
347 const ix = try fb.iota(img_ty, 1);
348 const iy = try fb.iota(img_ty, 0);
349 const zero = try consts.splat(0);
350 const warped_x = try fb.add(try fb.mul(ix, try consts.splat(0.35)), try consts.splat(0.75));
351 const warped_y = try fb.sub(try fb.mul(iy, try consts.splat(1.15)), try consts.splat(0.25));
352 const cxx = try fb.min(try fb.max(warped_x, zero), try consts.splat(@floatFromInt(cols - 1)));
353 const cyy = try fb.min(try fb.max(warped_y, zero), try consts.splat(@floatFromInt(rows - 1)));
354 const x0 = try fb.min(try fb.floor(cxx), try consts.splat(@floatFromInt(cols - 2)));
355 const y0 = try fb.min(try fb.floor(cyy), try consts.splat(@floatFromInt(rows - 2)));
356 const width = try consts.splat(@floatFromInt(cols));
357 const one = try consts.splat(1);
358 const base = try fb.add(try fb.mul(y0, width), x0);
359 const idx = try fb.convert(try fb.add(try fb.add(base, width), one), idx_ty, .i32);
360 const gathered = try fb.gather(fb.parameter(0), idx, img_ty, 0);
361 const out = try fb.add(gathered, zero);
362 try fb.return_(&.{out});
363 try fb.finish();
364 return try builder.finish();
365 }
366
367 fn semanticRandomFamilyModule(
368 allocator: std.mem.Allocator,
369 family_target: []const u8,
370 version: u32,
371 dtype: choir_abi.DType,
372 count: u64,
373 ) !*SemanticModule {
374 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
375 errdefer builder.deinit();
376 const out_ty = try builder.tensor(dtype, &.{@intCast(count)});
377 var fb = try builder.beginFunction("semantic_random_kernel_call_family", &.{}, &.{out_ty});
378 const call = try fb.kernelCall(
379 &.{},
380 &.{out_ty},
381 .{
382 .target = family_target,
383 .version = version,
384 .operand_effects = &.{},
385 .result_aliases = &.{null},
386 },
387 );
388 try fb.return_(&.{call.getFirstResult()});
389 try fb.finish();
390 return try builder.finish();
391 }
392
393 test "compaction filter family AOT registry executes on live CUDA" {
394 const allocator = std.testing.allocator;
395
396 var state = try initCudaStateOrSkip(allocator);
397 defer state.deinit();
398
399 const compaction_library = kernel.library.compaction;
400 const instance = compaction_library.Filter{
401 .extent = 1000,
402 .threads = 64,
403 };
404 var call_artifact = try compaction_library.createFilterFamilyArtifact(
405 allocator,
406 state.handle(),
407 instance,
408 .{ .limits = .testing },
409 );
410 defer call_artifact.deinit();
411
412 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
413 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
414 defer allocator.free(encoded);
415 var decoded = try artifact.wire.decode(allocator, encoded);
416 defer decoded.deinit();
417
418 const family_target = try compaction_library.filterFamilyTarget(allocator, instance);
419 defer allocator.free(family_target);
420
421 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
422 var builder_owned = true;
423 errdefer if (builder_owned) builder.deinit();
424 const data_ty = try builder.tensor(.f32, &.{1000});
425 const packed_ty = try builder.tensor(.f32, &.{1016});
426 var fb = try builder.beginFunction("semantic_filter_kernel_call_family", &.{data_ty}, &.{packed_ty});
427 const call = try fb.kernelCall(
428 &.{fb.parameter(0)},
429 &.{packed_ty},
430 .{
431 .target = family_target,
432 .version = compaction_library.filter_family_version,
433 .operand_effects = &.{.none},
434 .result_aliases = &.{null},
435 },
436 );
437 try fb.return_(&.{call.getResult(0).?});
438 try fb.finish();
439 builder_owned = false;
440 const module = try builder.finish();
441
442 const registry = decoded.registry();
443 var fragment = try compileAndLoadSemanticModule(
444 allocator,
445 state.handle(),
446 module,
447 .{ .kernel_call_registry = ®istry },
448 );
449 defer fragment.deinit();
450
451 var data: [1000]f32 = undefined;
452 var seed: u32 = 0x9e3779b9;
453 for (&data, 0..) |*value, index| {
454 seed ^= seed << 13;
455 seed ^= seed >> 17;
456 seed ^= seed << 5;
457 value.* = if (seed % 4 == 0) 0.0 else @floatFromInt(index + 1);
458 }
459
460 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(data[0..])};
461 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
462 defer bindings.deinit();
463
464 const runtime_args = try compaction_library.filterRuntimeArguments(instance);
465 try bindings.launchWithOptions(allocator, .{
466 .runtime_scalar_arguments = runtime_args[0..],
467 });
468
469 var actual_dst = @as([1016]f32, @splat(-1.0));
470 try bindings.readOutput(0, std.mem.sliceAsBytes(actual_dst[0..]));
471
472 var expected_dst = @as([1016]f32, @splat(-1.0));
473 compaction_library.filterBlocksExpectedF32(data[0..], instance.threads, expected_dst[0..]);
474
475 try std.testing.expectEqualSlices(f32, expected_dst[1000..], actual_dst[1000..]);
476 for (0..16) |segment| {
477 const begin = segment * 64;
478 const survivors: usize = @intFromFloat(expected_dst[1000 + segment]);
479 try std.testing.expectEqualSlices(
480 f32,
481 expected_dst[begin .. begin + survivors],
482 actual_dst[begin .. begin + survivors],
483 );
484 }
485 }
486
487 test "compaction greater filter family AOT registry executes on live CUDA" {
488 const allocator = std.testing.allocator;
489
490 var state = try initCudaStateOrSkip(allocator);
491 defer state.deinit();
492
493 const compaction_library = kernel.library.compaction;
494 const instance = compaction_library.Filter{
495 .extent = 1000,
496 .predicate = .greater_than,
497 .threads = 64,
498 };
499 var call_artifact = try compaction_library.createFilterFamilyArtifact(
500 allocator,
501 state.handle(),
502 instance,
503 .{ .limits = .testing },
504 );
505 defer call_artifact.deinit();
506
507 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
508 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
509 defer allocator.free(encoded);
510 var decoded = try artifact.wire.decode(allocator, encoded);
511 defer decoded.deinit();
512
513 const family_target = try compaction_library.filterFamilyTarget(allocator, instance);
514 defer allocator.free(family_target);
515
516 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
517 var builder_owned = true;
518 errdefer if (builder_owned) builder.deinit();
519 const data_ty = try builder.tensor(.f32, &.{1000});
520 const packed_ty = try builder.tensor(.f32, &.{1016});
521 var fb = try builder.beginFunction("semantic_greater_filter_kernel_call_family", &.{data_ty}, &.{packed_ty});
522 const call = try fb.kernelCall(
523 &.{fb.parameter(0)},
524 &.{packed_ty},
525 .{
526 .target = family_target,
527 .version = compaction_library.filter_family_version,
528 .operand_effects = &.{.none},
529 .result_aliases = &.{null},
530 },
531 );
532 try fb.return_(&.{call.getResult(0).?});
533 try fb.finish();
534 builder_owned = false;
535 const module = try builder.finish();
536
537 const registry = decoded.registry();
538 var fragment = try compileAndLoadSemanticModule(
539 allocator,
540 state.handle(),
541 module,
542 .{ .kernel_call_registry = ®istry },
543 );
544 defer fragment.deinit();
545
546 var data: [1000]f32 = undefined;
547 var seed: u32 = 0x6c078965;
548 for (&data, 0..) |*value, index| {
549 seed ^= seed << 13;
550 seed ^= seed >> 17;
551 seed ^= seed << 5;
552 const magnitude: f32 = @floatFromInt(index + 1);
553 value.* = if (seed % 2 == 0) -magnitude else magnitude;
554 }
555
556 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(data[0..])};
557 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
558 defer bindings.deinit();
559
560 const threshold: f32 = 500.0;
561 const runtime_args = try compaction_library.filterGreaterRuntimeArguments(instance, .{ .f32 = threshold });
562 try bindings.launchWithOptions(allocator, .{
563 .runtime_scalar_arguments = runtime_args[0..],
564 });
565
566 var actual_dst = @as([1016]f32, @splat(-1.0));
567 try bindings.readOutput(0, std.mem.sliceAsBytes(actual_dst[0..]));
568
569 var expected_dst = @as([1016]f32, @splat(-1.0));
570 compaction_library.filterBlocksGreaterExpectedF32(data[0..], instance.threads, threshold, expected_dst[0..]);
571
572 try std.testing.expectEqualSlices(f32, expected_dst[1000..], actual_dst[1000..]);
573 for (0..16) |segment| {
574 const begin = segment * 64;
575 const survivors: usize = @intFromFloat(expected_dst[1000 + segment]);
576 for (actual_dst[begin .. begin + survivors]) |value| try std.testing.expect(value > threshold);
577 try std.testing.expectEqualSlices(
578 f32,
579 expected_dst[begin .. begin + survivors],
580 actual_dst[begin .. begin + survivors],
581 );
582 }
583 }
584
585 test "compaction composes filter prefix sum and gather into dense output on live CUDA" {
586 const allocator = std.testing.allocator;
587
588 var state = try initCudaStateOrSkip(allocator);
589 defer state.deinit();
590
591 const compaction_library = kernel.library.compaction;
592 const scan_library = kernel.library.scan;
593 const indexing_library = kernel.library.indexing;
594
595 const extent: usize = 1000;
596 const threads: usize = 64;
597 const segment_count: usize = 16;
598 const padded: usize = extent + segment_count;
599
600 var data: [extent]f32 = undefined;
601 var seed: u32 = 0x9e3779b9;
602 for (&data, 0..) |*value, index| {
603 seed ^= seed << 13;
604 seed ^= seed >> 17;
605 seed ^= seed << 5;
606 value.* = if (seed % 4 == 0) 0.0 else @floatFromInt(index + 1);
607 }
608
609 const filter_instance = compaction_library.Filter{ .extent = extent, .threads = threads };
610 var packed_dst = @as([padded]f32, @splat(-1.0));
611 {
612 var call_artifact = try compaction_library.createFilterFamilyArtifact(
613 allocator,
614 state.handle(),
615 filter_instance,
616 .{ .limits = .testing },
617 );
618 defer call_artifact.deinit();
619 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
620 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
621 defer allocator.free(encoded);
622 var decoded = try artifact.wire.decode(allocator, encoded);
623 defer decoded.deinit();
624
625 const family_target = try compaction_library.filterFamilyTarget(allocator, filter_instance);
626 defer allocator.free(family_target);
627
628 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
629 var builder_owned = true;
630 errdefer if (builder_owned) builder.deinit();
631 const data_ty = try builder.tensor(.f32, &.{extent});
632 const packed_ty = try builder.tensor(.f32, &.{padded});
633 var fb = try builder.beginFunction("composition_filter_stage", &.{data_ty}, &.{packed_ty});
634 const call = try fb.kernelCall(&.{fb.parameter(0)}, &.{packed_ty}, .{
635 .target = family_target,
636 .version = compaction_library.filter_family_version,
637 .operand_effects = &.{.none},
638 .result_aliases = &.{null},
639 });
640 try fb.return_(&.{call.getFirstResult()});
641 try fb.finish();
642 builder_owned = false;
643 const module = try builder.finish();
644
645 const registry = decoded.registry();
646 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
647 .kernel_call_registry = ®istry,
648 });
649 defer fragment.deinit();
650
651 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(data[0..])};
652 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
653 defer bindings.deinit();
654 const runtime_args = try compaction_library.filterRuntimeArguments(filter_instance);
655 try bindings.launchWithOptions(allocator, .{
656 .runtime_scalar_arguments = runtime_args[0..],
657 });
658 try bindings.readOutput(0, std.mem.sliceAsBytes(packed_dst[0..]));
659 }
660
661 const prefix_instance = scan_library.PrefixSum{ .extent = segment_count, .mode = .exclusive, .threads = 32 };
662 var prefix = @as([segment_count]f32, @splat(-1.0));
663 {
664 var call_artifact = try scan_library.createPrefixSumFamilyArtifact(
665 allocator,
666 state.handle(),
667 prefix_instance,
668 .{ .limits = .testing },
669 );
670 defer call_artifact.deinit();
671 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
672 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
673 defer allocator.free(encoded);
674 var decoded = try artifact.wire.decode(allocator, encoded);
675 defer decoded.deinit();
676
677 const family_target = try scan_library.prefixSumFamilyTarget(allocator, prefix_instance);
678 defer allocator.free(family_target);
679 const module = try semanticPrefixSumFamilyModule(allocator, family_target, segment_count);
680
681 const registry = decoded.registry();
682 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
683 .kernel_call_registry = ®istry,
684 });
685 defer fragment.deinit();
686
687 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(packed_dst[extent..])};
688 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
689 defer bindings.deinit();
690 const runtime_args = try scan_library.prefixSumRuntimeArguments(prefix_instance);
691 try bindings.launchWithOptions(allocator, .{
692 .runtime_scalar_arguments = runtime_args[0..],
693 });
694 try bindings.readOutput(0, std.mem.sliceAsBytes(prefix[0..]));
695 }
696
697 var indices: [extent]i32 = undefined;
698 var total: usize = 0;
699 for (0..segment_count) |segment| {
700 const begin: usize = @intFromFloat(prefix[segment]);
701 const survivors: usize = @intFromFloat(packed_dst[extent + segment]);
702 for (0..survivors) |offset| {
703 indices[begin + offset] = @intCast(segment * threads + offset);
704 }
705 total = begin + survivors;
706 }
707
708 var expected_dense: [extent]f32 = undefined;
709 var expected_total: usize = 0;
710 for (data) |value| {
711 if (value != 0.0) {
712 expected_dense[expected_total] = value;
713 expected_total += 1;
714 }
715 }
716 try std.testing.expectEqual(expected_total, total);
717
718 const gather_instance = indexing_library.Gather{
719 .axis_size = padded,
720 .gathered = total,
721 .threads = 64,
722 };
723 const dense = try allocator.alloc(f32, total);
724 defer allocator.free(dense);
725 {
726 var call_artifact = try indexing_library.createGatherFamilyArtifact(
727 allocator,
728 state.handle(),
729 gather_instance,
730 .{ .limits = .testing },
731 );
732 defer call_artifact.deinit();
733 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
734 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
735 defer allocator.free(encoded);
736 var decoded = try artifact.wire.decode(allocator, encoded);
737 defer decoded.deinit();
738
739 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
740 var builder_owned = true;
741 errdefer if (builder_owned) builder.deinit();
742 const source_ty = try builder.tensor(.f32, &.{padded});
743 const indices_ty = try builder.tensor(.i32, &.{@intCast(total)});
744 const out_ty = try builder.tensor(.f32, &.{@intCast(total)});
745 var fb = try builder.beginFunction("composition_gather_stage", &.{ source_ty, indices_ty }, &.{out_ty});
746 const out = try fb.gather(fb.parameter(0), fb.parameter(1), out_ty, 0);
747 try fb.return_(&.{out});
748 try fb.finish();
749 builder_owned = false;
750 const module = try builder.finish();
751
752 const registry = decoded.registry();
753 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
754 .kernel_call_registry = ®istry,
755 .gather_schedule = .{ .thread_blocks = 64 },
756 });
757 defer fragment.deinit();
758
759 const input_bytes = [_][]const u8{
760 std.mem.sliceAsBytes(packed_dst[0..]),
761 std.mem.sliceAsBytes(indices[0..total]),
762 };
763 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
764 defer bindings.deinit();
765 const runtime_args = try indexing_library.gatherRuntimeArguments(gather_instance);
766 try bindings.launchWithOptions(allocator, .{
767 .runtime_scalar_arguments = runtime_args[0..],
768 });
769 try bindings.readOutput(0, std.mem.sliceAsBytes(dense));
770 }
771
772 try std.testing.expectEqualSlices(f32, expected_dense[0..expected_total], dense);
773 }
774
775 test "squares family AOT registry executes on live CUDA" {
776 const allocator = std.testing.allocator;
777
778 var state = try initCudaStateOrSkip(allocator);
779 defer state.deinit();
780
781 const random_library = kernel.library.random;
782 const instance = random_library.Squares{
783 .count = 777,
784 .threads = 64,
785 .dtype = .i32,
786 };
787 var call_artifact = try random_library.createSquaresFamilyArtifact(
788 allocator,
789 state.handle(),
790 instance,
791 .{ .limits = .testing },
792 );
793 defer call_artifact.deinit();
794
795 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
796 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
797 defer allocator.free(encoded);
798 var decoded = try artifact.wire.decode(allocator, encoded);
799 defer decoded.deinit();
800
801 const family_target = try random_library.squaresFamilyTarget(allocator, instance);
802 defer allocator.free(family_target);
803 const module = try semanticRandomFamilyModule(
804 allocator,
805 family_target,
806 random_library.squares_family_version,
807 .i32,
808 instance.count,
809 );
810
811 const registry = decoded.registry();
812 var fragment = try compileAndLoadSemanticModule(
813 allocator,
814 state.handle(),
815 module,
816 .{ .kernel_call_registry = ®istry },
817 );
818 defer fragment.deinit();
819
820 const input_bytes = [_][]const u8{};
821 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
822 defer bindings.deinit();
823
824 const runtime_args = try random_library.squaresRuntimeArguments(instance);
825 try bindings.launchWithOptions(allocator, .{
826 .runtime_scalar_arguments = runtime_args[0..],
827 });
828
829 var actual = @as([777]i32, @splat(0));
830 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
831
832 for (0..777) |element| {
833 const expected = random_library.squaresBlock(@intCast(element), instance.key);
834 try std.testing.expectEqual(@as(i32, @bitCast(expected)), actual[element]);
835 }
836 }
837
838 test "philox family AOT registry executes on live CUDA" {
839 const allocator = std.testing.allocator;
840
841 var state = try initCudaStateOrSkip(allocator);
842 defer state.deinit();
843
844 const random_library = kernel.library.random;
845 const instance = random_library.Philox{
846 .count = 1000,
847 .threads = 64,
848 .seed = 0xdecafbad00c0ffee,
849 };
850 var call_artifact = try random_library.createPhiloxFamilyArtifact(
851 allocator,
852 state.handle(),
853 instance,
854 .{ .limits = .testing },
855 );
856 defer call_artifact.deinit();
857
858 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
859 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
860 defer allocator.free(encoded);
861 var decoded = try artifact.wire.decode(allocator, encoded);
862 defer decoded.deinit();
863
864 const family_target = try random_library.philoxFamilyTarget(allocator, instance);
865 defer allocator.free(family_target);
866 const module = try semanticRandomFamilyModule(
867 allocator,
868 family_target,
869 random_library.philox_family_version,
870 .f32,
871 instance.count,
872 );
873
874 const registry = decoded.registry();
875 var fragment = try compileAndLoadSemanticModule(
876 allocator,
877 state.handle(),
878 module,
879 .{ .kernel_call_registry = ®istry },
880 );
881 defer fragment.deinit();
882
883 const input_bytes = [_][]const u8{};
884 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
885 defer bindings.deinit();
886
887 const runtime_args = try random_library.philoxRuntimeArguments(instance);
888 try bindings.launchWithOptions(allocator, .{
889 .runtime_scalar_arguments = runtime_args[0..],
890 });
891
892 var actual = @as([1000]f32, @splat(-1.0));
893 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
894
895 for (0..250) |generator| {
896 const words = random_library.philoxBlock(
897 random_library.philox_default_rounds,
898 .{ @intCast(generator), 0, 0, 0 },
899 .{ instance.seedLo(), instance.seedHi() },
900 );
901 for (words, 0..) |word, lane| {
902 try std.testing.expectEqual(random_library.uniformFromBits(word), actual[generator * 4 + lane]);
903 }
904 }
905 }
906
907 test "threefry family AOT registry executes on live CUDA" {
908 const allocator = std.testing.allocator;
909
910 var state = try initCudaStateOrSkip(allocator);
911 defer state.deinit();
912
913 const random_library = kernel.library.random;
914 const instance = random_library.Threefry{
915 .count = 999,
916 .threads = 128,
917 .seed = 0x0123456789abcdef,
918 .dtype = .i32,
919 };
920 var call_artifact = try random_library.createThreefryFamilyArtifact(
921 allocator,
922 state.handle(),
923 instance,
924 .{ .limits = .testing },
925 );
926 defer call_artifact.deinit();
927
928 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
929 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
930 defer allocator.free(encoded);
931 var decoded = try artifact.wire.decode(allocator, encoded);
932 defer decoded.deinit();
933
934 const family_target = try random_library.threefryFamilyTarget(allocator, instance);
935 defer allocator.free(family_target);
936 const module = try semanticRandomFamilyModule(
937 allocator,
938 family_target,
939 random_library.threefry_family_version,
940 .i32,
941 instance.count,
942 );
943
944 const registry = decoded.registry();
945 var fragment = try compileAndLoadSemanticModule(
946 allocator,
947 state.handle(),
948 module,
949 .{ .kernel_call_registry = ®istry },
950 );
951 defer fragment.deinit();
952
953 const input_bytes = [_][]const u8{};
954 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
955 defer bindings.deinit();
956
957 const runtime_args = try random_library.threefryRuntimeArguments(instance);
958 try bindings.launchWithOptions(allocator, .{
959 .runtime_scalar_arguments = runtime_args[0..],
960 });
961
962 var actual = @as([999]i32, @splat(0));
963 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
964
965 for (0..500) |generator| {
966 const words = random_library.threefryBlock(
967 random_library.threefry_default_rounds,
968 .{ @intCast(generator), 0 },
969 .{ instance.seedLo(), instance.seedHi() },
970 );
971 for (words, 0..) |word, lane| {
972 const element = generator * 2 + lane;
973 if (element >= instance.count) continue;
974 try std.testing.expectEqual(@as(i32, @bitCast(word)), actual[element]);
975 }
976 }
977 }
978
979 fn semanticEinsumMatrixVectorProductModule(allocator: std.mem.Allocator) !*SemanticModule {
980 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
981 errdefer builder.deinit();
982 const matrix_ty = try builder.tensor(.f32, &.{ 4, 8 });
983 const vector_ty = try builder.tensor(.f32, &.{8});
984 const out_ty = try builder.tensor(.f32, &.{4});
985 var fb = try builder.beginFunction("semantic_einsum_kernel_call_matvec", &.{ matrix_ty, vector_ty }, &.{out_ty});
986 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,k->m");
987 try fb.return_(&.{out});
988 try fb.finish();
989 return try builder.finish();
990 }
991
992 fn semanticReluModule(allocator: std.mem.Allocator) !*SemanticModule {
993 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
994 errdefer builder.deinit();
995 const input_ty = try builder.tensor(.f32, &.{8});
996 const zeros = @as([8]f32, @splat(0.0));
997 var fb = try builder.beginFunction("semantic_relu_kernel_call", &.{input_ty}, &.{input_ty});
998 const zero = try fb.constant(input_ty, std.mem.sliceAsBytes(zeros[0..]));
999 const out = try fb.max(fb.parameter(0), zero);
1000 try fb.return_(&.{out});
1001 try fb.finish();
1002 return try builder.finish();
1003 }
1004
1005 fn semanticGeluActivationModule(allocator: std.mem.Allocator) !*SemanticModule {
1006 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1007 errdefer builder.deinit();
1008 const input_ty = try builder.tensor(.f32, &.{8});
1009 var fb = try builder.beginFunction("semantic_gelu_activation_kernel_call", &.{input_ty}, &.{input_ty});
1010 const out = try fb.activation(fb.parameter(0), .gelu);
1011 try fb.return_(&.{out});
1012 try fb.finish();
1013 return try builder.finish();
1014 }
1015
1016 fn semanticEinsumTransposeModule(allocator: std.mem.Allocator) !*SemanticModule {
1017 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1018 errdefer builder.deinit();
1019 const input_ty = try builder.tensor(.f32, &.{ 8, 16 });
1020 const out_ty = try builder.tensor(.f32, &.{ 16, 8 });
1021 var fb = try builder.beginFunction("semantic_einsum_kernel_call_transpose", &.{input_ty}, &.{out_ty});
1022 const out = try fb.einsum(&.{fb.parameter(0)}, out_ty, "ij->ji");
1023 try fb.return_(&.{out});
1024 try fb.finish();
1025 return try builder.finish();
1026 }
1027
1028 fn semanticEinsumDotModule(allocator: std.mem.Allocator) !*SemanticModule {
1029 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1030 errdefer builder.deinit();
1031 const input_ty = try builder.tensor(.f32, &.{8});
1032 const out_ty = try builder.tensor(.f32, &.{});
1033 var fb = try builder.beginFunction("semantic_einsum_kernel_call_dot", &.{ input_ty, input_ty }, &.{out_ty});
1034 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "i,i->");
1035 try fb.return_(&.{out});
1036 try fb.finish();
1037 return try builder.finish();
1038 }
1039
1040 fn semanticKernelCallRowSoftmaxModule(
1041 allocator: std.mem.Allocator,
1042 descriptor: kernel.library.CatalogDescriptor,
1043 ) !*SemanticModule {
1044 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1045 errdefer builder.deinit();
1046 const f32_2x4 = try builder.tensor(.f32, &.{ 2, 4 });
1047 var fb = try builder.beginFunction("semantic_kernel_call_row_softmax", &.{f32_2x4}, &.{f32_2x4});
1048 const call = try fb.kernelCall(
1049 &.{fb.parameter(0)},
1050 &.{f32_2x4},
1051 .{
1052 .target = descriptor.metadata.target,
1053 .version = descriptor.metadata.version,
1054 .operand_effects = &.{.read},
1055 .result_aliases = &.{null},
1056 },
1057 );
1058 try fb.return_(&.{call.getFirstResult()});
1059 try fb.finish();
1060 return try builder.finish();
1061 }
1062
1063 test "semantic kernel_call resolves through kernel library registry" {
1064 const allocator = std.testing.allocator;
1065
1066 var state = gpu.recording.BackendState{
1067 .allocator = allocator,
1068 .kind = .cuda,
1069 .format = .cuda_ptx,
1070 };
1071
1072 var call_artifact = try KernelCallAdd.createKernelCallArtifact(
1073 allocator,
1074 state.handle(),
1075 .{ .limits = KernelCallAdd.Limits.testing },
1076 );
1077 defer call_artifact.deinit();
1078
1079 const module = try semanticKernelCallAddModule(allocator);
1080
1081 const registry = call_artifact.registry();
1082 var fragment = try compileAndLoadSemanticModule(
1083 allocator,
1084 state.handle(),
1085 module,
1086 .{ .kernel_call_registry = ®istry },
1087 );
1088 defer fragment.deinit();
1089
1090 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1091 const summary = try fragment.kernelSummary(0);
1092 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
1093 try std.testing.expectEqualStrings(KernelCallAdd.name, summary.entry_name);
1094 try std.testing.expectEqual(@as(u32, 3), summary.compile_argument_count);
1095 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
1096 try std.testing.expect(summary.fixed_threadgroup);
1097 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
1098 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[0]);
1099 }
1100
1101 test "semantic einsum matrix product lowers through kernel library registry" {
1102 const allocator = std.testing.allocator;
1103
1104 var state = gpu.recording.BackendState{
1105 .allocator = allocator,
1106 .kind = .cuda,
1107 .format = .cuda_ptx,
1108 };
1109
1110 var call_artifact = try kernel.library.createKernelCallArtifact(allocator, state.handle(), .{
1111 .target = KernelCallMatrixProduct.target,
1112 .version = KernelCallMatrixProduct.version,
1113 .options = .{ .limits = KernelCallMatrixProduct.Limits.testing },
1114 });
1115 defer call_artifact.deinit();
1116
1117 const module = try semanticEinsumMatrixProductModule(allocator);
1118
1119 const registry = call_artifact.registry();
1120 var fragment = try compileAndLoadSemanticModule(
1121 allocator,
1122 state.handle(),
1123 module,
1124 .{ .kernel_call_registry = ®istry },
1125 );
1126 defer fragment.deinit();
1127
1128 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1129 const summary = try fragment.kernelSummary(0);
1130 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
1131 try std.testing.expectEqualStrings(KernelCallMatrixProduct.name, summary.entry_name);
1132 try std.testing.expectEqual(@as(u32, 3), summary.compile_argument_count);
1133 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
1134 try std.testing.expect(summary.fixed_threadgroup);
1135 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
1136 try std.testing.expectEqual(@as(u32, 1), summary.launch_geometry.grid[1]);
1137 try std.testing.expectEqual(@as(u32, 8), summary.launch_geometry.threadgroup[0]);
1138 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[1]);
1139 }
1140
1141 test "semantic einsum matrix product family lowers through AOT wire registry" {
1142 const allocator = std.testing.allocator;
1143
1144 var state = gpu.recording.BackendState{
1145 .allocator = allocator,
1146 .kind = .cuda,
1147 .format = .cuda_ptx,
1148 };
1149
1150 const instance = kernel.library.linalg.MatrixProduct{
1151 .m = 5,
1152 .n = 7,
1153 .k = 3,
1154 .threads = kernel.library.linalg.matrixProductThreadsForExtents(5, 7),
1155 };
1156 var call_artifact = try kernel.library.linalg.createMatrixProductFamilyArtifact(
1157 allocator,
1158 state.handle(),
1159 instance,
1160 .{ .limits = .testing },
1161 );
1162 defer call_artifact.deinit();
1163
1164 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1165 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1166 defer allocator.free(encoded);
1167 var decoded = try artifact.wire.decode(allocator, encoded);
1168 defer decoded.deinit();
1169
1170 const module = try semanticEinsumMatrixProductFamilyModule(allocator);
1171
1172 const registry = decoded.registry();
1173 var fragment = try compileAndLoadSemanticModule(
1174 allocator,
1175 state.handle(),
1176 module,
1177 .{ .kernel_call_registry = ®istry },
1178 );
1179 defer fragment.deinit();
1180
1181 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1182 const summary = try fragment.kernelSummary(0);
1183 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
1184 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_family_7x5_f32", summary.entry_name);
1185 try std.testing.expectEqual(@as(u32, 6), summary.compile_argument_count);
1186 try std.testing.expectEqual(gpu.DTypeSet.init(&.{ .f32, .i32 }).bits, summary.compile_required_dtype_bits);
1187 try std.testing.expectEqual(@as(u32, 3), summary.runtime_scalar_argument_count);
1188 try std.testing.expect(summary.fixed_threadgroup);
1189 try std.testing.expectEqual(@as(usize, 1), summary.launch_candidate_count);
1190 try std.testing.expectEqual(@as(u32, 7), summary.launch_geometry.threadgroup[0]);
1191 try std.testing.expectEqual(@as(u32, 5), summary.launch_geometry.threadgroup[1]);
1192 }
1193
1194 test "semantic einsum matrix product family lowers explicit schedule through registry" {
1195 const allocator = std.testing.allocator;
1196
1197 var state = gpu.recording.BackendState{
1198 .allocator = allocator,
1199 .kind = .cuda,
1200 .format = .cuda_ptx,
1201 };
1202
1203 const schedule = kernel.library.MatrixProductSchedule{ .thread_blocks = .{ .x = 4, .y = 2 } };
1204 const lhs_dims = [_]i64{ 5, 3 };
1205 const rhs_dims = [_]i64{ 3, 7 };
1206 const output_dims = [_]i64{ 5, 7 };
1207 var selected = (try kernel.library.selectOwned(allocator, .{ .matrix_product = .{
1208 .dtype = .f32,
1209 .lhs_indices = "mk",
1210 .rhs_indices = "kn",
1211 .output_indices = "mn",
1212 .lhs_dims = &lhs_dims,
1213 .rhs_dims = &rhs_dims,
1214 .output_dims = &output_dims,
1215 .schedule = schedule,
1216 } })) orelse return error.TestExpectedCatalogDescriptor;
1217 defer selected.deinit();
1218
1219 var call_artifact = try kernel.library.createOwnedKernelCallArtifact(
1220 allocator,
1221 state.handle(),
1222 selected,
1223 .{ .limits = .testing },
1224 );
1225 defer call_artifact.deinit();
1226 try std.testing.expectEqual(@as(u32, 4), call_artifact.entry().launch.derived.threadgroup[0]);
1227 try std.testing.expectEqual(@as(u32, 2), call_artifact.entry().launch.derived.threadgroup[1]);
1228
1229 const module = try semanticEinsumMatrixProductFamilyModule(allocator);
1230 const registry = call_artifact.registry();
1231 var fragment = try compileAndLoadSemanticModule(
1232 allocator,
1233 state.handle(),
1234 module,
1235 .{
1236 .kernel_call_registry = ®istry,
1237 .matrix_product_schedule = schedule,
1238 },
1239 );
1240 defer fragment.deinit();
1241
1242 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1243 const summary = try fragment.kernelSummary(0);
1244 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
1245 try std.testing.expectEqualStrings("accy_kernel_linalg_matmul_family_4x2_f32", summary.entry_name);
1246 try std.testing.expectEqual(artifact.PlannedKernelCompileLaunch.kernel_call, summary.compile_launch);
1247 try std.testing.expectEqual(@as(u32, 3), summary.runtime_scalar_argument_count);
1248 }
1249
1250 test "semantic einsum matrix product family AOT registry executes on live CUDA" {
1251 const allocator = std.testing.allocator;
1252
1253 var state = try initCudaStateOrSkip(allocator);
1254 defer state.deinit();
1255
1256 const instance = kernel.library.linalg.MatrixProduct{
1257 .m = 5,
1258 .n = 7,
1259 .k = 3,
1260 .threads = kernel.library.linalg.matrixProductThreadsForExtents(5, 7),
1261 };
1262 var call_artifact = try kernel.library.linalg.createMatrixProductFamilyArtifact(
1263 allocator,
1264 state.handle(),
1265 instance,
1266 .{ .limits = .testing },
1267 );
1268 defer call_artifact.deinit();
1269
1270 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1271 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1272 defer allocator.free(encoded);
1273 var decoded = try artifact.wire.decode(allocator, encoded);
1274 defer decoded.deinit();
1275
1276 const module = try semanticEinsumMatrixProductFamilyModule(allocator);
1277
1278 const registry = decoded.registry();
1279 var fragment = try compileAndLoadSemanticModule(
1280 allocator,
1281 state.handle(),
1282 module,
1283 .{ .kernel_call_registry = ®istry },
1284 );
1285 defer fragment.deinit();
1286
1287 var lhs: [15]f32 = undefined;
1288 for (&lhs, 0..) |*value, index| value.* = @floatFromInt(index + 1);
1289 var rhs: [21]f32 = undefined;
1290 for (&rhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index)) * 0.5 - 2.0;
1291 var expected: [35]f32 = undefined;
1292 fillMatrixProductExpected(expected[0..], lhs[0..], rhs[0..], instance.m, instance.n, instance.k);
1293
1294 const input_bytes = [_][]const u8{
1295 std.mem.sliceAsBytes(lhs[0..]),
1296 std.mem.sliceAsBytes(rhs[0..]),
1297 };
1298 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
1299 defer bindings.deinit();
1300
1301 const runtime_args = try kernel.library.linalg.matrixProductRuntimeArguments(instance);
1302 try bindings.launchWithOptions(allocator, .{
1303 .runtime_scalar_arguments = runtime_args[0..],
1304 });
1305
1306 var actual = @as([35]f32, @splat(0.0));
1307 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
1308
1309 for (expected, actual) |want, got| {
1310 try std.testing.expectApproxEqAbs(want, got, 0.0001);
1311 }
1312 }
1313
1314 test "semantic gather family lowers through AOT wire registry" {
1315 const allocator = std.testing.allocator;
1316
1317 var state = gpu.recording.BackendState{
1318 .allocator = allocator,
1319 .kind = .cuda,
1320 .format = .cuda_ptx,
1321 };
1322
1323 const instance = kernel.library.indexing.Gather{
1324 .outer = 2,
1325 .axis_size = 8,
1326 .gathered = 5,
1327 .inner = 3,
1328 .threads = 16,
1329 };
1330 var call_artifact = try kernel.library.indexing.createGatherFamilyArtifact(
1331 allocator,
1332 state.handle(),
1333 instance,
1334 .{ .limits = .testing },
1335 );
1336 defer call_artifact.deinit();
1337
1338 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1339 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1340 defer allocator.free(encoded);
1341 var decoded = try artifact.wire.decode(allocator, encoded);
1342 defer decoded.deinit();
1343
1344 const module = try semanticGatherFamilyModule(allocator);
1345
1346 const registry = decoded.registry();
1347 var fragment = try compileAndLoadSemanticModule(
1348 allocator,
1349 state.handle(),
1350 module,
1351 .{
1352 .kernel_call_registry = ®istry,
1353 .gather_schedule = .{ .thread_blocks = 16 },
1354 },
1355 );
1356 defer fragment.deinit();
1357
1358 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1359 const summary = try fragment.kernelSummary(0);
1360 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
1361 try std.testing.expectEqualStrings("accy_kernel_indexing_gather_family_16_f32", summary.entry_name);
1362 try std.testing.expectEqual(@as(u32, 8), summary.compile_argument_count);
1363 try std.testing.expectEqual(@as(u32, 5), summary.runtime_scalar_argument_count);
1364 try std.testing.expect(summary.fixed_threadgroup);
1365 try std.testing.expectEqual(@as(usize, 1), summary.launch_candidate_count);
1366 try std.testing.expectEqual(@as(u32, 16), summary.launch_geometry.threadgroup[0]);
1367 }
1368
1369 test "rank-2 iota gather fusion uses 2D generated launch geometry" {
1370 const allocator = std.testing.allocator;
1371
1372 var state = gpu.recording.BackendState{
1373 .allocator = allocator,
1374 .kind = .cuda,
1375 .format = .cuda_ptx,
1376 };
1377
1378 const module = try semanticRank2IotaGatherFusionModule(allocator);
1379 var prepared = try prepareFragmentFromSemanticModule(allocator, state.handle(), module, .{});
1380 defer prepared.deinit();
1381
1382 try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
1383 const summary = try prepared.generatedKernelSummary(0);
1384 try std.testing.expectEqual(GeneratedScheduleKind.matrix, summary.schedule.kind);
1385 const geometry = summary.launch_geometry orelse return error.TestExpectedLaunchGeometry;
1386 try std.testing.expectEqual(@as(u32, 1), geometry.grid[0]);
1387 try std.testing.expectEqual(@as(u32, 1), geometry.grid[1]);
1388 try std.testing.expectEqual(@as(u32, 32), geometry.threadgroup[0]);
1389 try std.testing.expectEqual(@as(u32, 8), geometry.threadgroup[1]);
1390
1391 const artifact_module = try semanticRank2IotaGatherFusionModule(allocator);
1392 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), artifact_module, .{});
1393 defer fragment.deinit();
1394
1395 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1396 const artifact_summary = try fragment.kernelSummary(0);
1397 try std.testing.expectEqual(ArtifactKernelSource.tensor, artifact_summary.source);
1398 try std.testing.expect(!artifact_summary.fixed_threadgroup);
1399 try std.testing.expect(artifact_summary.launch_candidate_count > 1);
1400 try std.testing.expectEqual(@as(u32, 1), artifact_summary.launch_geometry.grid[0]);
1401 try std.testing.expectEqual(@as(u32, 1), artifact_summary.launch_geometry.grid[1]);
1402 try std.testing.expectEqual(@as(u32, 32), artifact_summary.launch_geometry.threadgroup[0]);
1403 try std.testing.expectEqual(@as(u32, 8), artifact_summary.launch_geometry.threadgroup[1]);
1404
1405 var kernel_artifact = try fragment.copyKernelArtifact(allocator, 0);
1406 defer kernel_artifact.deinit();
1407 const ptx = switch (kernel_artifact.payload) {
1408 .text => |text| text,
1409 else => return error.TestExpectedTextPayload,
1410 };
1411 try std.testing.expect(std.mem.indexOf(u8, ptx, "max.s32") == null);
1412 try std.testing.expect(std.mem.indexOf(u8, ptx, "min.s32") == null);
1413 }
1414
1415 test "generated row pipeline schedule honors explicit and tuned selections" {
1416 const allocator = std.testing.allocator;
1417
1418 var state = gpu.recording.BackendState{
1419 .allocator = allocator,
1420 .kind = .cuda,
1421 .format = .cuda_ptx,
1422 };
1423
1424 const rows: u64 = 8;
1425 const cols: u64 = 2048;
1426
1427 {
1428 const module = try semanticRowPipelineModule(allocator, @intCast(rows), @intCast(cols));
1429 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{});
1430 defer fragment.deinit();
1431 const summary = try rowPipelineKernelSummary(fragment);
1432 try std.testing.expectEqual(@as(u32, 8), summary.launch_geometry.grid[0]);
1433 try std.testing.expectEqual(@as(u32, 256), summary.launch_geometry.threadgroup[0]);
1434 }
1435
1436 {
1437 const module = try semanticRowPipelineModule(allocator, @intCast(rows), @intCast(cols));
1438 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
1439 .generated_row_pipeline_schedule = .{ .threads = 512 },
1440 });
1441 defer fragment.deinit();
1442 const summary = try rowPipelineKernelSummary(fragment);
1443 try std.testing.expectEqual(@as(u32, 8), summary.launch_geometry.grid[0]);
1444 try std.testing.expectEqual(@as(u32, 512), summary.launch_geometry.threadgroup[0]);
1445 }
1446
1447 const caps = try state.handle().queryCapabilities();
1448 var candidate_buffer: [max_row_pipeline_schedule_candidates]GeneratedRowPipelineSchedule = undefined;
1449 const candidates = rowPipelineScheduleCandidates(cols, .cuda_ptx, &candidate_buffer);
1450 try std.testing.expectEqual(@as(usize, 3), candidates.len);
1451 const problem = executable.GeneratedRowPipelineScheduleTuningProblem{
1452 .format = .cuda_ptx,
1453 .rows = rows,
1454 .cols = cols,
1455 .dtype = .f32,
1456 .schedule_version = generated_row_pipeline_schedule_version,
1457 .candidates = candidates,
1458 };
1459 const tuned_record = executable.GeneratedRowPipelineScheduleTuningRecord{
1460 .key = try executable.GeneratedRowPipelineScheduleTuningKey.init(caps, problem),
1461 .selection = .{
1462 .schedule = .{ .threads = 128 },
1463 .winner_median_ns = 100,
1464 .runner_up_median_ns = 120,
1465 .sample_count = 10,
1466 },
1467 };
1468 const artifact_bytes = try executable.encodeGeneratedRowPipelineScheduleTuningArtifact(allocator, &.{tuned_record});
1469 defer allocator.free(artifact_bytes);
1470
1471 {
1472 const module = try semanticRowPipelineModule(allocator, @intCast(rows), @intCast(cols));
1473 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
1474 .generated_row_pipeline_schedule_tuning_artifact = artifact_bytes,
1475 });
1476 defer fragment.deinit();
1477 const summary = try rowPipelineKernelSummary(fragment);
1478 try std.testing.expectEqual(@as(u32, 8), summary.launch_geometry.grid[0]);
1479 try std.testing.expectEqual(@as(u32, 128), summary.launch_geometry.threadgroup[0]);
1480 }
1481 }
1482
1483 fn rowPipelineKernelSummary(fragment: *LoadedFragment) !artifact.KernelSummary {
1484 var kernel_index: usize = 0;
1485 while (kernel_index < fragment.kernelCount()) : (kernel_index += 1) {
1486 const summary = try fragment.kernelSummary(kernel_index);
1487 if (std.mem.indexOf(u8, summary.entry_name, "row_pipeline") != null) return summary;
1488 }
1489 return error.TestExpectedRowPipelineKernel;
1490 }
1491
1492 test "generated row pipeline schedules compute softmax on live CUDA" {
1493 const allocator = std.testing.allocator;
1494
1495 var state = try initCudaStateOrSkip(allocator);
1496 defer state.deinit();
1497
1498 const rows: usize = 8;
1499 const cols: usize = 2048;
1500
1501 const input = try allocator.alloc(f32, rows * cols);
1502 defer allocator.free(input);
1503 var seed: u32 = 0x243f6a88;
1504 for (input) |*value| {
1505 seed ^= seed << 13;
1506 seed ^= seed >> 17;
1507 seed ^= seed << 5;
1508 value.* = @as(f32, @floatFromInt(seed % 4096)) / 512.0 - 4.0;
1509 }
1510
1511 const expected = try allocator.alloc(f32, rows * cols);
1512 defer allocator.free(expected);
1513 for (0..rows) |row| {
1514 const row_values = input[row * cols .. (row + 1) * cols];
1515 var row_max = -std.math.floatMax(f32);
1516 for (row_values) |value| row_max = @max(row_max, value);
1517 var row_sum: f32 = 0.0;
1518 const expected_row = expected[row * cols .. (row + 1) * cols];
1519 for (row_values, expected_row) |value, *out| {
1520 out.* = @exp(value - row_max);
1521 row_sum += out.*;
1522 }
1523 for (expected_row) |*out| out.* /= row_sum;
1524 }
1525
1526 const schedules = [_]GeneratedRowPipelineSchedule{
1527 .{ .threads = 512 },
1528 .{ .threads = 256 },
1529 .{ .threads = 128 },
1530 };
1531 for (schedules) |schedule| {
1532 const module = try semanticRowPipelineModule(allocator, @intCast(rows), @intCast(cols));
1533 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
1534 .generated_row_pipeline_schedule = schedule,
1535 });
1536 defer fragment.deinit();
1537
1538 const summary = try rowPipelineKernelSummary(fragment);
1539 try std.testing.expectEqual(schedule.threads, summary.launch_geometry.threadgroup[0]);
1540
1541 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(input)};
1542 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
1543 defer bindings.deinit();
1544
1545 try bindings.launch(allocator);
1546
1547 const actual = try allocator.alloc(f32, rows * cols);
1548 defer allocator.free(actual);
1549 try bindings.readOutput(0, std.mem.sliceAsBytes(actual));
1550
1551 for (expected, actual) |want, got| {
1552 try std.testing.expectApproxEqAbs(want, got, 0.00001);
1553 }
1554 }
1555 }
1556
1557 test "generated scan schedule honors explicit and tuned selections" {
1558 const allocator = std.testing.allocator;
1559
1560 var state = gpu.recording.BackendState{
1561 .allocator = allocator,
1562 .kind = .cuda,
1563 .format = .cuda_ptx,
1564 };
1565
1566 const total: u64 = 16384;
1567
1568 {
1569 const module = try semanticCumsumModule(allocator, @intCast(total));
1570 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{});
1571 defer fragment.deinit();
1572 const summary = try scanKernelSummary(fragment);
1573 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
1574 try std.testing.expectEqual(@as(u32, 512), summary.launch_geometry.threadgroup[0]);
1575 }
1576
1577 {
1578 const module = try semanticCumsumModule(allocator, @intCast(total));
1579 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
1580 .generated_scan_schedule = .{ .threads = 256, .items = 16 },
1581 });
1582 defer fragment.deinit();
1583 const summary = try scanKernelSummary(fragment);
1584 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.grid[0]);
1585 try std.testing.expectEqual(@as(u32, 256), summary.launch_geometry.threadgroup[0]);
1586 }
1587
1588 const caps = try state.handle().queryCapabilities();
1589 var candidate_buffer: [max_scan_schedule_candidates]GeneratedScanSchedule = undefined;
1590 const candidates = scanScheduleCandidates(total, .cuda_ptx, &candidate_buffer);
1591 try std.testing.expectEqual(@as(usize, 2), candidates.len);
1592 const problem = executable.GeneratedScanScheduleTuningProblem{
1593 .format = .cuda_ptx,
1594 .total = total,
1595 .dtype = .f32,
1596 .schedule_version = generated_scan_schedule_version,
1597 .candidates = candidates,
1598 };
1599 const tuned_record = executable.GeneratedScanScheduleTuningRecord{
1600 .key = try executable.GeneratedScanScheduleTuningKey.init(caps, problem),
1601 .selection = .{
1602 .schedule = .{ .threads = 256, .items = 16 },
1603 .winner_median_ns = 100,
1604 .runner_up_median_ns = 120,
1605 .sample_count = 10,
1606 },
1607 };
1608 var foreign_key = tuned_record.key;
1609 foreign_key.total = total * 2;
1610 const foreign_record = executable.GeneratedScanScheduleTuningRecord{
1611 .key = foreign_key,
1612 .selection = .{
1613 .schedule = .{ .threads = 512, .items = 16 },
1614 .winner_median_ns = 90,
1615 .runner_up_median_ns = 110,
1616 .sample_count = 10,
1617 },
1618 };
1619 const artifact_bytes = try executable.encodeGeneratedScanScheduleTuningArtifact(allocator, &.{ tuned_record, foreign_record });
1620 defer allocator.free(artifact_bytes);
1621
1622 {
1623 const module = try semanticCumsumModule(allocator, @intCast(total));
1624 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{
1625 .generated_scan_schedule_tuning_artifact = artifact_bytes,
1626 });
1627 defer fragment.deinit();
1628 const summary = try scanKernelSummary(fragment);
1629 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.grid[0]);
1630 try std.testing.expectEqual(@as(u32, 256), summary.launch_geometry.threadgroup[0]);
1631 }
1632 }
1633
1634 fn scanKernelSummary(fragment: *LoadedFragment) !artifact.KernelSummary {
1635 var kernel_index: usize = 0;
1636 while (kernel_index < fragment.kernelCount()) : (kernel_index += 1) {
1637 const summary = try fragment.kernelSummary(kernel_index);
1638 if (std.mem.indexOf(u8, summary.entry_name, "scan_lookback") != null) return summary;
1639 }
1640 return error.TestExpectedScanKernel;
1641 }
1642
1643 test "rank-2 clamped gather fusion skips redundant signed gather clamp" {
1644 const allocator = std.testing.allocator;
1645
1646 var state = gpu.recording.BackendState{
1647 .allocator = allocator,
1648 .kind = .cuda,
1649 .format = .cuda_ptx,
1650 };
1651
1652 const module = try semanticRank2ClampedGatherFusionModule(allocator);
1653 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{});
1654 defer fragment.deinit();
1655
1656 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1657 var kernel_artifact = try fragment.copyKernelArtifact(allocator, 0);
1658 defer kernel_artifact.deinit();
1659 const ptx = switch (kernel_artifact.payload) {
1660 .text => |text| text,
1661 else => return error.TestExpectedTextPayload,
1662 };
1663 try std.testing.expect(std.mem.indexOf(u8, ptx, "max.s32") == null);
1664 try std.testing.expect(std.mem.indexOf(u8, ptx, "min.s32") == null);
1665 }
1666
1667 test "semantic gather family AOT registry executes on live CUDA" {
1668 const allocator = std.testing.allocator;
1669
1670 var state = try initCudaStateOrSkip(allocator);
1671 defer state.deinit();
1672
1673 const instance = kernel.library.indexing.Gather{
1674 .outer = 2,
1675 .axis_size = 8,
1676 .gathered = 5,
1677 .inner = 3,
1678 .threads = 16,
1679 };
1680 var call_artifact = try kernel.library.indexing.createGatherFamilyArtifact(
1681 allocator,
1682 state.handle(),
1683 instance,
1684 .{ .limits = .testing },
1685 );
1686 defer call_artifact.deinit();
1687
1688 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1689 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1690 defer allocator.free(encoded);
1691 var decoded = try artifact.wire.decode(allocator, encoded);
1692 defer decoded.deinit();
1693
1694 const module = try semanticGatherFamilyModule(allocator);
1695
1696 const registry = decoded.registry();
1697 var fragment = try compileAndLoadSemanticModule(
1698 allocator,
1699 state.handle(),
1700 module,
1701 .{
1702 .kernel_call_registry = ®istry,
1703 .gather_schedule = .{ .thread_blocks = 16 },
1704 },
1705 );
1706 defer fragment.deinit();
1707
1708 var data: [48]f32 = undefined;
1709 for (&data, 0..) |*value, index| value.* = @floatFromInt(index);
1710 const indices = [_]i32{ 3, 0, 7, -1, 9 };
1711 var expected: [30]f32 = undefined;
1712 fillGatherExpected(expected[0..], data[0..], indices[0..], 2, 8, 3);
1713
1714 const input_bytes = [_][]const u8{
1715 std.mem.sliceAsBytes(data[0..]),
1716 std.mem.sliceAsBytes(indices[0..]),
1717 };
1718 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
1719 defer bindings.deinit();
1720
1721 const runtime_args = try kernel.library.indexing.gatherRuntimeArguments(instance);
1722 try bindings.launchWithOptions(allocator, .{
1723 .runtime_scalar_arguments = runtime_args[0..],
1724 });
1725
1726 var actual = @as([30]f32, @splat(0.0));
1727 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
1728
1729 for (expected, actual) |want, got| {
1730 try std.testing.expectApproxEqAbs(want, got, 0.0001);
1731 }
1732 }
1733
1734 fn runSegmentSumFamilyOnLiveCuda(instance: kernel.library.segmented.SegmentSum) !void {
1735 const allocator = std.testing.allocator;
1736
1737 var state = try initCudaStateOrSkip(allocator);
1738 defer state.deinit();
1739
1740 var call_artifact = try kernel.library.segmented.createSegmentSumFamilyArtifact(
1741 allocator,
1742 state.handle(),
1743 instance,
1744 .{ .limits = .testing },
1745 );
1746 defer call_artifact.deinit();
1747
1748 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1749 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1750 defer allocator.free(encoded);
1751 var decoded = try artifact.wire.decode(allocator, encoded);
1752 defer decoded.deinit();
1753
1754 const family_target = try kernel.library.segmented.segmentSumFamilyTarget(allocator, instance);
1755 defer allocator.free(family_target);
1756 const module = try semanticSegmentSumFamilyModule(allocator, family_target);
1757
1758 const registry = decoded.registry();
1759 var fragment = try compileAndLoadSemanticModule(
1760 allocator,
1761 state.handle(),
1762 module,
1763 .{ .kernel_call_registry = ®istry },
1764 );
1765 defer fragment.deinit();
1766
1767 var data: [12]f32 = undefined;
1768 for (&data, 0..) |*value, index| value.* = @floatFromInt(index + 1);
1769 const offsets = [_]i32{ 0, 5, 5, 12 };
1770 var expected = [_]f32{ 0, 0, 0 };
1771 for (0..3) |segment| {
1772 const begin: usize = @intCast(offsets[segment]);
1773 const end: usize = @intCast(offsets[segment + 1]);
1774 for (begin..end) |element| expected[segment] += data[element];
1775 }
1776
1777 const input_bytes = [_][]const u8{
1778 std.mem.sliceAsBytes(data[0..]),
1779 std.mem.sliceAsBytes(offsets[0..]),
1780 };
1781 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
1782 defer bindings.deinit();
1783
1784 const runtime_args = try kernel.library.segmented.segmentSumRuntimeArguments(instance);
1785 try bindings.launchWithOptions(allocator, .{
1786 .runtime_scalar_arguments = runtime_args[0..],
1787 });
1788
1789 var actual = @as([3]f32, @splat(0.0));
1790 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
1791
1792 for (expected, actual) |want, got| {
1793 try std.testing.expectApproxEqAbs(want, got, 0.0001);
1794 }
1795 }
1796
1797 test "segment sum family AOT registry executes on live CUDA" {
1798 try runSegmentSumFamilyOnLiveCuda(.{
1799 .segments = 3,
1800 .total = 12,
1801 .threads = 4,
1802 });
1803 }
1804
1805 test "segment sum warp family AOT registry executes on live CUDA" {
1806 try runSegmentSumFamilyOnLiveCuda(.{
1807 .segments = 3,
1808 .total = 12,
1809 .granularity = .warp,
1810 .threads = 32,
1811 });
1812 }
1813
1814 fn semanticPrefixSumFamilyModule(allocator: std.mem.Allocator, family_target: []const u8, extent: i64) !*SemanticModule {
1815 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1816 errdefer builder.deinit();
1817 const data_ty = try builder.tensor(.f32, &.{extent});
1818 var fb = try builder.beginFunction("semantic_prefix_sum_kernel_call_family", &.{data_ty}, &.{data_ty});
1819 const call = try fb.kernelCall(
1820 &.{fb.parameter(0)},
1821 &.{data_ty},
1822 .{
1823 .target = family_target,
1824 .version = kernel.library.scan.prefix_sum_family_version,
1825 .operand_effects = &.{.read},
1826 .result_aliases = &.{null},
1827 },
1828 );
1829 try fb.return_(&.{call.getFirstResult()});
1830 try fb.finish();
1831 return try builder.finish();
1832 }
1833
1834 fn runPrefixSumFamilyOnLiveCuda(instance: kernel.library.scan.PrefixSum) !void {
1835 const allocator = std.testing.allocator;
1836
1837 var state = try initCudaStateOrSkip(allocator);
1838 defer state.deinit();
1839
1840 var call_artifact = try kernel.library.scan.createPrefixSumFamilyArtifact(
1841 allocator,
1842 state.handle(),
1843 instance,
1844 .{ .limits = .testing },
1845 );
1846 defer call_artifact.deinit();
1847
1848 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1849 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1850 defer allocator.free(encoded);
1851 var decoded = try artifact.wire.decode(allocator, encoded);
1852 defer decoded.deinit();
1853
1854 const family_target = try kernel.library.scan.prefixSumFamilyTarget(allocator, instance);
1855 defer allocator.free(family_target);
1856 const module = try semanticPrefixSumFamilyModule(allocator, family_target, 100);
1857
1858 const registry = decoded.registry();
1859 var fragment = try compileAndLoadSemanticModule(
1860 allocator,
1861 state.handle(),
1862 module,
1863 .{ .kernel_call_registry = ®istry },
1864 );
1865 defer fragment.deinit();
1866
1867 var data: [100]f32 = undefined;
1868 for (&data, 0..) |*value, index| value.* = @floatFromInt((index % 7) + 1);
1869 var expected: [100]f32 = undefined;
1870 var running: f32 = 0;
1871 for (data, 0..) |value, index| {
1872 switch (instance.mode) {
1873 .inclusive => {
1874 running += value;
1875 expected[index] = running;
1876 },
1877 .exclusive => {
1878 expected[index] = running;
1879 running += value;
1880 },
1881 }
1882 }
1883
1884 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(data[0..])};
1885 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
1886 defer bindings.deinit();
1887
1888 const runtime_args = try kernel.library.scan.prefixSumRuntimeArguments(instance);
1889 try bindings.launchWithOptions(allocator, .{
1890 .runtime_scalar_arguments = runtime_args[0..],
1891 });
1892
1893 var actual = @as([100]f32, @splat(0.0));
1894 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
1895
1896 for (expected, actual) |want, got| {
1897 try std.testing.expectApproxEqAbs(want, got, 0.001);
1898 }
1899 }
1900
1901 fn semanticScatterRoutingModule(allocator: std.mem.Allocator) !*SemanticModule {
1902 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1903 errdefer builder.deinit();
1904 const data_ty = try builder.tensor(.f32, &.{8});
1905 const indices_ty = try builder.tensor(.i32, &.{5});
1906 const updates_ty = try builder.tensor(.f32, &.{5});
1907 var fb = try builder.beginFunction("semantic_scatter_routing_family", &.{ data_ty, indices_ty, updates_ty }, &.{data_ty});
1908 const out = try fb.scatter(fb.parameter(0), fb.parameter(1), fb.parameter(2), data_ty, 0);
1909 try fb.return_(&.{out});
1910 try fb.finish();
1911 return try builder.finish();
1912 }
1913
1914 test "semantic scatter family routes and executes on live CUDA" {
1915 const allocator = std.testing.allocator;
1916
1917 var state = try initCudaStateOrSkip(allocator);
1918 defer state.deinit();
1919
1920 const instance = kernel.library.indexing.Scatter{
1921 .axis_size = 8,
1922 .updates = 5,
1923 .threads = 8,
1924 };
1925 var call_artifact = try kernel.library.indexing.createScatterFamilyArtifact(
1926 allocator,
1927 state.handle(),
1928 instance,
1929 .{ .limits = .testing },
1930 );
1931 defer call_artifact.deinit();
1932
1933 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
1934 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
1935 defer allocator.free(encoded);
1936 var decoded = try artifact.wire.decode(allocator, encoded);
1937 defer decoded.deinit();
1938
1939 const module = try semanticScatterRoutingModule(allocator);
1940
1941 const registry = decoded.registry();
1942 var fragment = try compileAndLoadSemanticModule(
1943 allocator,
1944 state.handle(),
1945 module,
1946 .{
1947 .kernel_call_registry = ®istry,
1948 .scatter_schedule = .{ .thread_blocks = 8 },
1949 },
1950 );
1951 defer fragment.deinit();
1952
1953 const data = [_]f32{ 10, 11, 12, 13, 14, 15, 16, 17 };
1954 const indices = [_]i32{ 3, 0, 3, 9, 1 };
1955 const updates = [_]f32{ 100, 200, 300, 400, 500 };
1956 const expected = [_]f32{ 200, 500, 12, 300, 14, 15, 16, 17 };
1957
1958 const input_bytes = [_][]const u8{
1959 std.mem.sliceAsBytes(data[0..]),
1960 std.mem.sliceAsBytes(indices[0..]),
1961 std.mem.sliceAsBytes(updates[0..]),
1962 };
1963 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
1964 defer bindings.deinit();
1965
1966 const runtime_args = try kernel.library.indexing.scatterRuntimeArguments(instance);
1967 try bindings.launchWithOptions(allocator, .{
1968 .runtime_scalar_arguments = runtime_args[0..],
1969 });
1970
1971 var actual = @as([8]f32, @splat(0.0));
1972 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
1973
1974 for (expected, actual) |want, got| {
1975 try std.testing.expectApproxEqAbs(want, got, 0.0001);
1976 }
1977 }
1978
1979 fn semanticScatterFamilyModule(allocator: std.mem.Allocator, family_target: []const u8) !*SemanticModule {
1980 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1981 errdefer builder.deinit();
1982 const data_ty = try builder.tensor(.f32, &.{8});
1983 const indices_ty = try builder.tensor(.i32, &.{4});
1984 const updates_ty = try builder.tensor(.f32, &.{4});
1985 const out_ty = try builder.tensor(.f32, &.{8});
1986 var fb = try builder.beginFunction("semantic_scatter_kernel_call_family", &.{ data_ty, indices_ty, updates_ty }, &.{out_ty});
1987 const call = try fb.kernelCall(
1988 &.{ fb.parameter(0), fb.parameter(1), fb.parameter(2) },
1989 &.{out_ty},
1990 .{
1991 .target = family_target,
1992 .version = kernel.library.indexing.scatter_family_version,
1993 .operand_effects = &.{ .read, .read, .read },
1994 .result_aliases = &.{null},
1995 },
1996 );
1997 try fb.return_(&.{call.getFirstResult()});
1998 try fb.finish();
1999 return try builder.finish();
2000 }
2001
2002 test "scatter family AOT registry executes on live CUDA" {
2003 const allocator = std.testing.allocator;
2004
2005 var state = try initCudaStateOrSkip(allocator);
2006 defer state.deinit();
2007
2008 const instance = kernel.library.indexing.Scatter{
2009 .axis_size = 8,
2010 .updates = 4,
2011 .threads = 8,
2012 };
2013 var call_artifact = try kernel.library.indexing.createScatterFamilyArtifact(
2014 allocator,
2015 state.handle(),
2016 instance,
2017 .{ .limits = .testing },
2018 );
2019 defer call_artifact.deinit();
2020
2021 const entries = [_]artifact.KernelCallArtifact{call_artifact.entry()};
2022 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{});
2023 defer allocator.free(encoded);
2024 var decoded = try artifact.wire.decode(allocator, encoded);
2025 defer decoded.deinit();
2026
2027 const family_target = try kernel.library.indexing.scatterFamilyTarget(allocator, instance);
2028 defer allocator.free(family_target);
2029 const module = try semanticScatterFamilyModule(allocator, family_target);
2030
2031 const registry = decoded.registry();
2032 var fragment = try compileAndLoadSemanticModule(
2033 allocator,
2034 state.handle(),
2035 module,
2036 .{ .kernel_call_registry = ®istry },
2037 );
2038 defer fragment.deinit();
2039
2040 const data = [_]f32{ 10, 11, 12, 13, 14, 15, 16, 17 };
2041 const indices = [_]i32{ 3, 0, 3, 9 };
2042 const updates = [_]f32{ 100, 200, 300, 400 };
2043 const expected = [_]f32{ 200, 11, 12, 300, 14, 15, 16, 17 };
2044
2045 const input_bytes = [_][]const u8{
2046 std.mem.sliceAsBytes(data[0..]),
2047 std.mem.sliceAsBytes(indices[0..]),
2048 std.mem.sliceAsBytes(updates[0..]),
2049 };
2050 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
2051 defer bindings.deinit();
2052
2053 const runtime_args = try kernel.library.indexing.scatterRuntimeArguments(instance);
2054 try bindings.launchWithOptions(allocator, .{
2055 .runtime_scalar_arguments = runtime_args[0..],
2056 });
2057
2058 var actual = @as([8]f32, @splat(0.0));
2059 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
2060
2061 for (expected, actual) |want, got| {
2062 try std.testing.expectApproxEqAbs(want, got, 0.0001);
2063 }
2064 }
2065
2066 test "semantic einsum matrix vector product family lowers through kernel library registry" {
2067 const allocator = std.testing.allocator;
2068
2069 var state = gpu.recording.BackendState{
2070 .allocator = allocator,
2071 .kind = .cuda,
2072 .format = .cuda_ptx,
2073 };
2074
2075 const matrix_dims = [_]i64{ 4, 8 };
2076 const vector_dims = [_]i64{8};
2077 const output_dims = [_]i64{4};
2078 var selected = (try kernel.library.selectOwned(allocator, .{ .matrix_vector_product = .{
2079 .dtype = .f32,
2080 .matrix_indices = "mk",
2081 .vector_indices = "k",
2082 .output_indices = "m",
2083 .matrix_dims = &matrix_dims,
2084 .vector_dims = &vector_dims,
2085 .output_dims = &output_dims,
2086 } })) orelse return error.TestExpectedCatalogDescriptor;
2087 defer selected.deinit();
2088
2089 var call_artifact = try kernel.library.createOwnedKernelCallArtifact(
2090 allocator,
2091 state.handle(),
2092 selected,
2093 .{ .limits = .testing },
2094 );
2095 defer call_artifact.deinit();
2096
2097 const module = try semanticEinsumMatrixVectorProductModule(allocator);
2098
2099 const registry = call_artifact.registry();
2100 var fragment = try compileAndLoadSemanticModule(
2101 allocator,
2102 state.handle(),
2103 module,
2104 .{ .kernel_call_registry = ®istry },
2105 );
2106 defer fragment.deinit();
2107
2108 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2109 const summary = try fragment.kernelSummary(0);
2110 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
2111 try std.testing.expectEqualStrings(selected.descriptor.name, summary.entry_name);
2112 try std.testing.expectEqual(@as(u32, 5), summary.compile_argument_count);
2113 try std.testing.expectEqual(gpu.DTypeSet.init(&.{ .f32, .i32 }).bits, summary.compile_required_dtype_bits);
2114 try std.testing.expectEqual(@as(u32, 2), summary.runtime_scalar_argument_count);
2115 try std.testing.expect(summary.fixed_threadgroup);
2116 try std.testing.expectEqual(@as(usize, 1), summary.launch_candidate_count);
2117 try std.testing.expectEqual(@as(u32, 1), summary.launch_geometry.grid[0]);
2118 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[0]);
2119 }
2120
2121 fn initCudaStateOrSkip(allocator: std.mem.Allocator) !gpu.cuda.State {
2122 try gating.skipIfBuildFlagDisabled(.cuda);
2123 if (!gpu.cuda.platformSupported()) return gating.skip(.cuda, .unsupported_platform);
2124 return gpu.cuda.State.initDevice(allocator, 0) catch |err| switch (err) {
2125 error.RuntimeUnavailable => return gating.skip(.cuda, .cuda_device_missing),
2126 else => return err,
2127 };
2128 }
2129
2130 fn fillMatrixProductExpected(
2131 out: []f32,
2132 lhs: []const f32,
2133 rhs: []const f32,
2134 m: u64,
2135 n: u64,
2136 k: u64,
2137 ) void {
2138 const rows: usize = @intCast(m);
2139 const cols: usize = @intCast(n);
2140 const reduction: usize = @intCast(k);
2141 for (0..rows) |row| {
2142 for (0..cols) |col| {
2143 var sum: f32 = 0.0;
2144 for (0..reduction) |offset| {
2145 sum += lhs[row * reduction + offset] * rhs[offset * cols + col];
2146 }
2147 out[row * cols + col] = sum;
2148 }
2149 }
2150 }
2151
2152 fn fillGatherExpected(
2153 out: []f32,
2154 data: []const f32,
2155 indices: []const i32,
2156 outer: usize,
2157 axis_size: usize,
2158 inner: usize,
2159 ) void {
2160 const gathered = indices.len;
2161 for (0..outer) |outer_index| {
2162 for (0..gathered) |gathered_index| {
2163 var index = indices[gathered_index];
2164 if (index < 0) index = 0;
2165 const limit: i32 = @intCast(axis_size - 1);
2166 if (index > limit) index = limit;
2167 const src_base = (outer_index * axis_size + @as(usize, @intCast(index))) * inner;
2168 const dst_base = (outer_index * gathered + gathered_index) * inner;
2169 for (0..inner) |inner_index| {
2170 out[dst_base + inner_index] = data[src_base + inner_index];
2171 }
2172 }
2173 }
2174 }
2175
2176 test "semantic relu lowers through kernel library registry" {
2177 const allocator = std.testing.allocator;
2178
2179 var state = gpu.recording.BackendState{
2180 .allocator = allocator,
2181 .kind = .cuda,
2182 .format = .cuda_ptx,
2183 };
2184
2185 var call_artifact = try kernel.library.createKernelCallArtifact(allocator, state.handle(), .{
2186 .target = KernelCallRelu.target,
2187 .version = KernelCallRelu.version,
2188 .options = .{ .limits = KernelCallRelu.Limits.testing },
2189 });
2190 defer call_artifact.deinit();
2191
2192 const module = try semanticReluModule(allocator);
2193
2194 const registry = call_artifact.registry();
2195 var fragment = try compileAndLoadSemanticModule(
2196 allocator,
2197 state.handle(),
2198 module,
2199 .{ .kernel_call_registry = ®istry },
2200 );
2201 defer fragment.deinit();
2202
2203 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2204 const summary = try fragment.kernelSummary(0);
2205 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
2206 try std.testing.expectEqualStrings(KernelCallRelu.name, summary.entry_name);
2207 try std.testing.expectEqual(@as(u32, 2), summary.compile_argument_count);
2208 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
2209 try std.testing.expect(summary.fixed_threadgroup);
2210 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
2211 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[0]);
2212 }
2213
2214 test "semantic relu prepares without kernel library registry" {
2215 const allocator = std.testing.allocator;
2216
2217 var state = gpu.recording.BackendState{
2218 .allocator = allocator,
2219 .kind = .cuda,
2220 .format = .cuda_ptx,
2221 };
2222
2223 const module = try semanticReluModule(allocator);
2224
2225 var fragment = try compileAndLoadSemanticModule(
2226 allocator,
2227 state.handle(),
2228 module,
2229 .{},
2230 );
2231 defer fragment.deinit();
2232
2233 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2234 const summary = try fragment.kernelSummary(0);
2235 try std.testing.expectEqual(ArtifactKernelSource.tensor, summary.source);
2236 }
2237
2238 test "semantic gelu activation lowers through kernel library registry" {
2239 const allocator = std.testing.allocator;
2240
2241 var state = gpu.recording.BackendState{
2242 .allocator = allocator,
2243 .kind = .cuda,
2244 .format = .cuda_ptx,
2245 };
2246
2247 var call_artifact = try kernel.library.createKernelCallArtifact(allocator, state.handle(), .{
2248 .target = KernelCallGelu.target,
2249 .version = KernelCallGelu.version,
2250 .options = .{ .limits = KernelCallGelu.Limits.testing },
2251 });
2252 defer call_artifact.deinit();
2253
2254 const module = try semanticGeluActivationModule(allocator);
2255
2256 const registry = call_artifact.registry();
2257 var fragment = try compileAndLoadSemanticModule(
2258 allocator,
2259 state.handle(),
2260 module,
2261 .{ .kernel_call_registry = ®istry },
2262 );
2263 defer fragment.deinit();
2264
2265 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2266 const summary = try fragment.kernelSummary(0);
2267 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
2268 try std.testing.expectEqualStrings(KernelCallGelu.name, summary.entry_name);
2269 try std.testing.expectEqual(@as(u32, 2), summary.compile_argument_count);
2270 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
2271 try std.testing.expect(summary.fixed_threadgroup);
2272 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
2273 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[0]);
2274 }
2275
2276 test "semantic gelu activation prepares without kernel library registry" {
2277 const allocator = std.testing.allocator;
2278
2279 var state = gpu.recording.BackendState{
2280 .allocator = allocator,
2281 .kind = .cuda,
2282 .format = .cuda_ptx,
2283 };
2284
2285 const module = try semanticGeluActivationModule(allocator);
2286
2287 var fragment = try compileAndLoadSemanticModule(
2288 allocator,
2289 state.handle(),
2290 module,
2291 .{},
2292 );
2293 defer fragment.deinit();
2294
2295 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2296 const summary = try fragment.kernelSummary(0);
2297 try std.testing.expectEqual(ArtifactKernelSource.tensor, summary.source);
2298 }
2299
2300 test "semantic einsum transpose lowers through kernel library registry" {
2301 const allocator = std.testing.allocator;
2302
2303 var state = gpu.recording.BackendState{
2304 .allocator = allocator,
2305 .kind = .cuda,
2306 .format = .cuda_ptx,
2307 };
2308
2309 var call_artifact = try kernel.library.createKernelCallArtifact(allocator, state.handle(), .{
2310 .target = KernelCallTranspose.target,
2311 .version = KernelCallTranspose.version,
2312 .options = .{ .limits = KernelCallTranspose.Limits.testing },
2313 });
2314 defer call_artifact.deinit();
2315
2316 const module = try semanticEinsumTransposeModule(allocator);
2317
2318 const registry = call_artifact.registry();
2319 var fragment = try compileAndLoadSemanticModule(
2320 allocator,
2321 state.handle(),
2322 module,
2323 .{ .kernel_call_registry = ®istry },
2324 );
2325 defer fragment.deinit();
2326
2327 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2328 const summary = try fragment.kernelSummary(0);
2329 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
2330 try std.testing.expectEqualStrings(KernelCallTranspose.name, summary.entry_name);
2331 try std.testing.expectEqual(@as(u32, 2), summary.compile_argument_count);
2332 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
2333 try std.testing.expect(summary.fixed_threadgroup);
2334 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[0]);
2335 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.grid[1]);
2336 try std.testing.expectEqual(@as(u32, 8), summary.launch_geometry.threadgroup[0]);
2337 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[1]);
2338 }
2339
2340 test "semantic einsum dot product lowers through kernel library registry" {
2341 const allocator = std.testing.allocator;
2342
2343 var state = gpu.recording.BackendState{
2344 .allocator = allocator,
2345 .kind = .cuda,
2346 .format = .cuda_ptx,
2347 };
2348
2349 var call_artifact = try kernel.library.createKernelCallArtifact(allocator, state.handle(), .{
2350 .target = KernelCallDot.target,
2351 .version = KernelCallDot.version,
2352 .options = .{ .limits = KernelCallDot.Limits.testing },
2353 });
2354 defer call_artifact.deinit();
2355
2356 const module = try semanticEinsumDotModule(allocator);
2357
2358 const registry = call_artifact.registry();
2359 var fragment = try compileAndLoadSemanticModule(
2360 allocator,
2361 state.handle(),
2362 module,
2363 .{ .kernel_call_registry = ®istry },
2364 );
2365 defer fragment.deinit();
2366
2367 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2368 const summary = try fragment.kernelSummary(0);
2369 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
2370 try std.testing.expectEqualStrings(KernelCallDot.name, summary.entry_name);
2371 try std.testing.expectEqual(@as(u32, 3), summary.compile_argument_count);
2372 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
2373 try std.testing.expect(summary.fixed_threadgroup);
2374 try std.testing.expectEqual(@as(u32, 1), summary.launch_geometry.grid[0]);
2375 try std.testing.expectEqual(@as(u32, 1), summary.launch_geometry.threadgroup[0]);
2376 }
2377
2378 test "semantic einsum transpose prepares without kernel library registry" {
2379 const allocator = std.testing.allocator;
2380
2381 var state = gpu.recording.BackendState{
2382 .allocator = allocator,
2383 .kind = .cuda,
2384 .format = .cuda_ptx,
2385 };
2386
2387 const module = try semanticEinsumTransposeModule(allocator);
2388
2389 var fragment = try compileAndLoadSemanticModule(
2390 allocator,
2391 state.handle(),
2392 module,
2393 .{},
2394 );
2395 defer fragment.deinit();
2396
2397 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2398 const summary = try fragment.kernelSummary(0);
2399 try std.testing.expectEqual(ArtifactKernelSource.tensor, summary.source);
2400 }
2401
2402 test "semantic row normalization selection resolves through kernel library registry" {
2403 const allocator = std.testing.allocator;
2404
2405 var state = gpu.recording.BackendState{
2406 .allocator = allocator,
2407 .kind = .cuda,
2408 .format = .cuda_ptx,
2409 };
2410
2411 const selected = kernel.logical.selectRowNormalizationCatalog(.{
2412 .dtype = .f32,
2413 .kind = .softmax,
2414 .rows = 2,
2415 .cols = 4,
2416 }) orelse return error.TestExpectedRowNormalizationSelection;
2417
2418 var call_artifact = try kernel.library.createKernelCallArtifact(allocator, state.handle(), .{
2419 .target = selected.descriptor.metadata.target,
2420 .version = selected.descriptor.metadata.version,
2421 .options = .{ .limits = kernel.Limits.testing },
2422 });
2423 defer call_artifact.deinit();
2424
2425 const module = try semanticKernelCallRowSoftmaxModule(allocator, selected.descriptor);
2426
2427 const registry = call_artifact.registry();
2428 var fragment = try compileAndLoadSemanticModule(
2429 allocator,
2430 state.handle(),
2431 module,
2432 .{ .kernel_call_registry = ®istry },
2433 );
2434 defer fragment.deinit();
2435
2436 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2437 const summary = try fragment.kernelSummary(0);
2438 try std.testing.expectEqual(ArtifactKernelSource.kernel_call, summary.source);
2439 try std.testing.expectEqualStrings(selected.descriptor.name, summary.entry_name);
2440 try std.testing.expectEqual(@as(u32, 2), summary.compile_argument_count);
2441 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, summary.compile_required_dtype_bits);
2442 try std.testing.expect(summary.fixed_threadgroup);
2443 try std.testing.expectEqual(@as(u32, 1), summary.launch_geometry.grid[0]);
2444 try std.testing.expectEqual(@as(u32, 1), summary.launch_geometry.grid[1]);
2445 try std.testing.expectEqual(@as(u32, 4), summary.launch_geometry.threadgroup[0]);
2446 try std.testing.expectEqual(@as(u32, 2), summary.launch_geometry.threadgroup[1]);
2447 }
2448
2449 fn semanticShiftPadModule(allocator: std.mem.Allocator) !*SemanticModule {
2450 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
2451 errdefer builder.deinit();
2452 const vec_ty = try builder.tensor(.f32, &.{8});
2453 const scalar_ty = try builder.tensor(.f32, &.{});
2454 var fb = try builder.beginFunction("semantic_shift_pad", &.{vec_ty}, &.{vec_ty});
2455 const zero: f32 = 0;
2456 const pad_value = try fb.constant(scalar_ty, std.mem.asBytes(&zero));
2457 const out = try fb.pad(fb.parameter(0), pad_value, vec_ty, &.{1}, &.{-1}, &.{0});
2458 try fb.return_(&.{out});
2459 try fb.finish();
2460 return try builder.finish();
2461 }
2462
2463 test "negative edge pad lowers through the cuda executable path" {
2464 const allocator = std.testing.allocator;
2465 var state = gpu.recording.BackendState{
2466 .allocator = allocator,
2467 .kind = .cuda,
2468 .format = .cuda_ptx,
2469 };
2470
2471 const module = try semanticShiftPadModule(allocator);
2472 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{});
2473 defer fragment.deinit();
2474 try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
2475 }
2476
2477 test "negative edge pad shifts elements on live CUDA" {
2478 const allocator = std.testing.allocator;
2479
2480 var state = try initCudaStateOrSkip(allocator);
2481 defer state.deinit();
2482
2483 const module = try semanticShiftPadModule(allocator);
2484 var fragment = try compileAndLoadSemanticModule(allocator, state.handle(), module, .{});
2485 defer fragment.deinit();
2486
2487 var data: [8]f32 = undefined;
2488 for (&data, 0..) |*value, index| value.* = @floatFromInt(index + 1);
2489 const expected = [_]f32{ 0, 1, 2, 3, 4, 5, 6, 7 };
2490
2491 const input_bytes = [_][]const u8{std.mem.sliceAsBytes(data[0..])};
2492 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
2493 defer bindings.deinit();
2494
2495 try bindings.launchWithOptions(allocator, .{});
2496
2497 var actual = @as([8]f32, @splat(0.0));
2498 try bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
2499
2500 for (expected, actual) |want, got| {
2501 try std.testing.expectApproxEqAbs(want, got, 0.0001);
2502 }
2503 }
2504
2505 fn semanticEinsumTunedMatmulModule(allocator: std.mem.Allocator) !*SemanticModule {
2506 var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
2507 errdefer builder.deinit();
2508 const lhs_ty = try builder.tensor(.f32, &.{ 64, 32 });
2509 const rhs_ty = try builder.tensor(.f32, &.{ 32, 64 });
2510 const out_ty = try builder.tensor(.f32, &.{ 64, 64 });
2511 var fb = try builder.beginFunction("semantic_einsum_family_tuning_round_trip", &.{ lhs_ty, rhs_ty }, &.{out_ty});
2512 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,kn->mn");
2513 try fb.return_(&.{out});
2514 try fb.finish();
2515 return try builder.finish();
2516 }
2517
2518 test "family tuning round trip selects the measured winner on live CUDA" {
2519 const allocator = std.testing.allocator;
2520
2521 var state = try initCudaStateOrSkip(allocator);
2522 defer state.deinit();
2523
2524 const linalg = kernel.library.linalg;
2525 const tuning = kernel.library.tuning;
2526 const probe = linalg.MatrixProduct{ .m = 64, .n = 64, .k = 32 };
2527 const caps = try state.handle().queryCapabilities();
2528 const device = tuning.deviceFingerprint(caps);
2529
2530 const thread_candidates = linalg.matrixProductThreadCandidatesForExtents(probe.m, probe.n);
2531 try std.testing.expect(thread_candidates.slice().len >= 2);
2532
2533 var artifacts = std.ArrayListUnmanaged(kernel.OwnedKernelCallArtifact).empty;
2534 defer {
2535 for (artifacts.items) |*owned| owned.deinit();
2536 artifacts.deinit(allocator);
2537 }
2538 for (thread_candidates.slice()) |threads| {
2539 var instance = probe;
2540 instance.threads = threads;
2541 const call_artifact = try linalg.createMatrixProductFamilyArtifact(
2542 allocator,
2543 state.handle(),
2544 instance,
2545 .{ .limits = .testing },
2546 );
2547 try artifacts.append(allocator, call_artifact);
2548 }
2549
2550 var entries = std.ArrayListUnmanaged(artifact.KernelCallArtifact).empty;
2551 defer entries.deinit(allocator);
2552 for (artifacts.items) |*owned| {
2553 try entries.append(allocator, owned.entry());
2554 }
2555 const encoded = try artifact.wire.encode(allocator, entries.items, &.{});
2556 defer allocator.free(encoded);
2557 var decoded = try artifact.wire.decode(allocator, encoded);
2558 defer decoded.deinit();
2559 const registry = decoded.registry();
2560
2561 var lhs: [64 * 32]f32 = undefined;
2562 for (lhs[0..], 0..) |*value, index| value.* = @floatFromInt((index % 13) + 1);
2563 var rhs: [32 * 64]f32 = undefined;
2564 for (rhs[0..], 0..) |*value, index| value.* = @floatFromInt((index % 7) + 2);
2565 const input_bytes = [_][]const u8{
2566 std.mem.sliceAsBytes(lhs[0..]),
2567 std.mem.sliceAsBytes(rhs[0..]),
2568 };
2569
2570 var accumulator = tuning.FamilyMeasurementAccumulator.init(allocator);
2571 defer accumulator.deinit();
2572
2573 for (thread_candidates.slice()) |threads| {
2574 var instance = probe;
2575 instance.threads = threads;
2576 const measured_target = try linalg.matrixProductFamilyTarget(allocator, instance);
2577 defer allocator.free(measured_target);
2578 const entry_name = try linalg.matrixProductFamilyEntryName(allocator, instance);
2579 defer allocator.free(entry_name);
2580
2581 const module = try semanticEinsumTunedMatmulModule(allocator);
2582 var fragment = try compileAndLoadSemanticModule(
2583 allocator,
2584 state.handle(),
2585 module,
2586 .{
2587 .kernel_call_registry = ®istry,
2588 .matrix_product_schedule = .{ .thread_blocks = threads },
2589 },
2590 );
2591 defer fragment.deinit();
2592
2593 const bindings = try prepareInvocation(fragment, allocator, input_bytes[0..]);
2594 defer bindings.deinit();
2595
2596 var kernel_index: usize = 0;
2597 var family_kernel_index: ?usize = null;
2598 while (kernel_index < fragment.kernelCount()) : (kernel_index += 1) {
2599 const summary = try fragment.kernelSummary(kernel_index);
2600 if (std.mem.eql(u8, summary.entry_name, entry_name)) {
2601 family_kernel_index = kernel_index;
2602 break;
2603 }
2604 }
2605 const measured_index = family_kernel_index orelse return error.TestExpectedFamilyKernel;
2606
2607 const runtime_arguments = try linalg.matrixProductRuntimeArguments(instance);
2608 const records = try bindings.measureLaunchCandidateRecords(
2609 allocator,
2610 allocator,
2611 measured_index,
2612 .{
2613 .warmup = 2,
2614 .samples = 8,
2615 .base_options = .{ .runtime_scalar_arguments = runtime_arguments[0..] },
2616 .synchronize = .device,
2617 },
2618 );
2619 defer allocator.free(records);
2620 try std.testing.expect(records.len >= 1);
2621
2622 var best = records[0];
2623 for (records[1..]) |record| {
2624 if (record.median_ns < best.median_ns) best = record;
2625 }
2626 const key = try kernel.library.linalg.matrixProductFamilyTuningKey(allocator, device, instance);
2627 try accumulator.append(key, measured_target, best.median_ns, best.sample_count);
2628 }
2629
2630 var winners = try accumulator.selectWinners(allocator, 0);
2631 defer winners.deinit();
2632 try std.testing.expectEqual(@as(usize, 1), winners.records.len);
2633 const winner = winners.records[0];
2634
2635 const reader = tuning.FamilyTuningReader.init(caps, winners.table());
2636 const tuned_module = try semanticEinsumTunedMatmulModule(allocator);
2637 var tuned_fragment = try compileAndLoadSemanticModule(
2638 allocator,
2639 state.handle(),
2640 tuned_module,
2641 .{
2642 .kernel_call_registry = ®istry,
2643 .family_tuning = &reader,
2644 },
2645 );
2646 defer tuned_fragment.deinit();
2647
2648 var winner_entry_seen = false;
2649 var kernel_index: usize = 0;
2650 while (kernel_index < tuned_fragment.kernelCount()) : (kernel_index += 1) {
2651 const summary = try tuned_fragment.kernelSummary(kernel_index);
2652 for (thread_candidates.slice()) |threads| {
2653 var instance = probe;
2654 instance.threads = threads;
2655 const candidate_entry = try linalg.matrixProductFamilyEntryName(allocator, instance);
2656 defer allocator.free(candidate_entry);
2657 if (!std.mem.eql(u8, summary.entry_name, candidate_entry)) continue;
2658 const candidate_target = try linalg.matrixProductFamilyTarget(allocator, instance);
2659 defer allocator.free(candidate_target);
2660 try std.testing.expectEqualStrings(winner.target, candidate_target);
2661 winner_entry_seen = true;
2662 }
2663 }
2664 try std.testing.expect(winner_entry_seen);
2665
2666 const tuned_bindings = try prepareInvocation(tuned_fragment, allocator, input_bytes[0..]);
2667 defer tuned_bindings.deinit();
2668 const tuned_runtime_arguments = try linalg.matrixProductRuntimeArguments(probe);
2669 try tuned_bindings.launchWithOptions(allocator, .{
2670 .runtime_scalar_arguments = tuned_runtime_arguments[0..],
2671 });
2672
2673 var actual = @as([(64 * 64)]f32, @splat(0));
2674 try tuned_bindings.readOutput(0, std.mem.sliceAsBytes(actual[0..]));
2675
2676 var expected = @as([(64 * 64)]f32, @splat(0));
2677 for (0..64) |row| {
2678 for (0..64) |col| {
2679 var accumulated: f32 = 0;
2680 for (0..32) |inner| {
2681 accumulated += lhs[row * 32 + inner] * rhs[inner * 64 + col];
2682 }
2683 expected[row * 64 + col] = accumulated;
2684 }
2685 }
2686 for (expected, actual) |want, got| {
2687 try std.testing.expectApproxEqAbs(want, got, 0.001);
2688 }
2689 }
2690
2691 test "scatter add family kernel accumulates exactly on live CUDA" {
2692 const allocator = std.testing.allocator;
2693
2694 var state = try initCudaStateOrSkip(allocator);
2695 defer state.deinit();
2696 const handle = state.handle();
2697
2698 const indexing_library = kernel.library.indexing;
2699 const instance = indexing_library.ScatterAdd{ .axis_size = 8, .updates = 5, .threads = 32 };
2700 const entry_name = try indexing_library.scatterAddFamilyEntryName(allocator, instance);
2701 defer allocator.free(entry_name);
2702
2703 var graph = try indexing_library.ScatterAddRuntimeFamilyI32.buildNamed(allocator, indexing_library.ScatterAddRuntimeFamilyI32.Limits.testing, entry_name, instance);
2704 defer graph.deinit();
2705 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
2706 .artifact_format = .cuda_ptx,
2707 .authored_kernel_diagnostic_id = "choir/kernel/scatter-add-live",
2708 });
2709 defer kernel_artifact.deinit();
2710 const loaded = try handle.loadArtifact(&kernel_artifact);
2711
2712 var dst = [_]i32{ 5, 5, 5, 5, 5, 5, 5, 5 };
2713 var indices = [_]i32{ 3, 0, 3, 9, 1 };
2714 var updates = [_]i32{ 100, 200, 300, 400, 500 };
2715 const expected = [_]i32{ 205, 505, 5, 405, 5, 5, 5, 5 };
2716
2717 const host_views = [_][]u8{
2718 std.mem.sliceAsBytes(dst[0..]),
2719 std.mem.sliceAsBytes(dst[0..]),
2720 std.mem.sliceAsBytes(indices[0..]),
2721 std.mem.sliceAsBytes(updates[0..]),
2722 };
2723 const accesses = [_]gpu.BufferAccess{ .read_write, .read_only, .read_only, .read_only };
2724
2725 var bindings: [4]gpu.BufferBinding = undefined;
2726 for (host_views, accesses, 0..) |bytes, access, slot| {
2727 const buffer = try handle.allocateBuffer(.{
2728 .byte_size = bytes.len,
2729 .alignment = 256,
2730 .dtype = .i32,
2731 .element_count = bytes.len / @sizeOf(i32),
2732 });
2733 try handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });
2734 bindings[slot] = .{
2735 .handle = buffer,
2736 .access = access,
2737 .ownership = buffer.ownership,
2738 .byte_size = buffer.byte_size,
2739 };
2740 }
2741
2742 const runtime_arguments = try indexing_library.scatterAddRuntimeArguments(instance);
2743 try handle.launch(.{
2744 .artifact = &kernel_artifact,
2745 .loaded_artifact = loaded,
2746 .buffers = bindings[0..],
2747 .scalar_arguments = runtime_arguments[0..],
2748 .geometry = .{
2749 .grid = .{ 1, 1, 1 },
2750 .threadgroup = .{ instance.threads, 1, 1 },
2751 },
2752 });
2753 try handle.synchronize(.{ .scope = .device });
2754
2755 var actual = @as([8]i32, @splat(0));
2756 try handle.readBuffer(.{ .handle = bindings[0].handle, .bytes = std.mem.sliceAsBytes(actual[0..]) });
2757 try std.testing.expectEqualSlices(i32, expected[0..], actual[0..]);
2758 }
2759
2760 test "scatter add f32 family kernel accumulates within tolerance on live CUDA" {
2761 const allocator = std.testing.allocator;
2762
2763 var state = try initCudaStateOrSkip(allocator);
2764 defer state.deinit();
2765 const handle = state.handle();
2766
2767 const indexing_library = kernel.library.indexing;
2768 const bin_count: usize = 16;
2769 const update_count: usize = 4096;
2770 const instance = indexing_library.ScatterAdd{
2771 .axis_size = bin_count,
2772 .updates = update_count,
2773 .dtype = .f32,
2774 .threads = 128,
2775 };
2776 const entry_name = try indexing_library.scatterAddFamilyEntryName(allocator, instance);
2777 defer allocator.free(entry_name);
2778
2779 var graph = try indexing_library.ScatterAddRuntimeFamilyF32.buildNamed(allocator, indexing_library.ScatterAddRuntimeFamilyF32.Limits.testing, entry_name, instance);
2780 defer graph.deinit();
2781 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
2782 .artifact_format = .cuda_ptx,
2783 .authored_kernel_diagnostic_id = "choir/kernel/scatter-add-f32-live",
2784 });
2785 defer kernel_artifact.deinit();
2786 const loaded = try handle.loadArtifact(&kernel_artifact);
2787
2788 var dst = @as([bin_count]f32, @splat(0));
2789 var indices: [update_count]i32 = undefined;
2790 var updates: [update_count]f32 = undefined;
2791 for (0..update_count) |u| {
2792 indices[u] = @intCast((u * 7 + 3) % bin_count);
2793 updates[u] = 0.5 + @as(f32, @floatFromInt((u * 13) % 100)) * 0.01;
2794 }
2795
2796 var expected = @as([bin_count]f32, @splat(0));
2797 for (0..update_count) |u| {
2798 expected[@intCast(indices[u])] += updates[u];
2799 }
2800
2801 const host_views = [_][]u8{
2802 std.mem.sliceAsBytes(dst[0..]),
2803 std.mem.sliceAsBytes(dst[0..]),
2804 std.mem.sliceAsBytes(indices[0..]),
2805 std.mem.sliceAsBytes(updates[0..]),
2806 };
2807 const dtypes = [_]choir_abi.DType{ .f32, .f32, .i32, .f32 };
2808 const accesses = [_]gpu.BufferAccess{ .read_write, .read_only, .read_only, .read_only };
2809
2810 var bindings: [4]gpu.BufferBinding = undefined;
2811 for (host_views, dtypes, accesses, 0..) |bytes, buffer_dtype, access, slot| {
2812 const buffer = try handle.allocateBuffer(.{
2813 .byte_size = bytes.len,
2814 .alignment = 256,
2815 .dtype = buffer_dtype,
2816 .element_count = bytes.len / 4,
2817 });
2818 try handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });
2819 bindings[slot] = .{
2820 .handle = buffer,
2821 .access = access,
2822 .ownership = buffer.ownership,
2823 .byte_size = buffer.byte_size,
2824 };
2825 }
2826
2827 const runtime_arguments = try indexing_library.scatterAddRuntimeArguments(instance);
2828 try handle.launch(.{
2829 .artifact = &kernel_artifact,
2830 .loaded_artifact = loaded,
2831 .buffers = bindings[0..],
2832 .scalar_arguments = runtime_arguments[0..],
2833 .geometry = .{
2834 .grid = .{ @intCast((update_count + 127) / 128), 1, 1 },
2835 .threadgroup = .{ instance.threads, 1, 1 },
2836 },
2837 });
2838 try handle.synchronize(.{ .scope = .device });
2839
2840 var actual = @as([bin_count]f32, @splat(0));
2841 try handle.readBuffer(.{ .handle = bindings[0].handle, .bytes = std.mem.sliceAsBytes(actual[0..]) });
2842 for (expected, actual) |want, got| {
2843 try std.testing.expectApproxEqAbs(want, got, 0.05);
2844 }
2845 }
2846
2847 test "scatter add shared bins variant accumulates exactly on live CUDA" {
2848 const allocator = std.testing.allocator;
2849
2850 var state = try initCudaStateOrSkip(allocator);
2851 defer state.deinit();
2852 const handle = state.handle();
2853
2854 const indexing_library = kernel.library.indexing;
2855 const bin_count: usize = 64;
2856 const update_count: usize = 4096;
2857 const instance = indexing_library.ScatterAdd{
2858 .axis_size = bin_count,
2859 .updates = update_count,
2860 .variant = .shared_bins,
2861 .threads = 128,
2862 };
2863 const entry_name = try indexing_library.scatterAddFamilyEntryName(allocator, instance);
2864 defer allocator.free(entry_name);
2865
2866 var graph = try indexing_library.ScatterAddRuntimeFamilyI32.buildNamed(allocator, indexing_library.ScatterAddRuntimeFamilyI32.Limits.testing, entry_name, instance);
2867 defer graph.deinit();
2868 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
2869 .artifact_format = .cuda_ptx,
2870 .authored_kernel_diagnostic_id = "choir/kernel/scatter-add-shared-live",
2871 });
2872 defer kernel_artifact.deinit();
2873 const loaded = try handle.loadArtifact(&kernel_artifact);
2874
2875 var dst = @as([bin_count]i32, @splat(0));
2876 var indices: [update_count]i32 = undefined;
2877 var updates: [update_count]i32 = undefined;
2878 for (0..update_count) |u| {
2879 indices[u] = @intCast((u * 11 + 5) % bin_count);
2880 updates[u] = @intCast((u % 9) + 1);
2881 }
2882
2883 var expected = @as([bin_count]i32, @splat(0));
2884 for (indices, updates) |index, update| {
2885 expected[@intCast(index)] += update;
2886 }
2887
2888 const host_views = [_][]u8{
2889 std.mem.sliceAsBytes(dst[0..]),
2890 std.mem.sliceAsBytes(dst[0..]),
2891 std.mem.sliceAsBytes(indices[0..]),
2892 std.mem.sliceAsBytes(updates[0..]),
2893 };
2894 const accesses = [_]gpu.BufferAccess{ .read_write, .read_only, .read_only, .read_only };
2895
2896 var bindings: [4]gpu.BufferBinding = undefined;
2897 for (host_views, accesses, 0..) |bytes, access, slot| {
2898 const buffer = try handle.allocateBuffer(.{
2899 .byte_size = bytes.len,
2900 .alignment = 256,
2901 .dtype = .i32,
2902 .element_count = bytes.len / @sizeOf(i32),
2903 });
2904 try handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });
2905 bindings[slot] = .{
2906 .handle = buffer,
2907 .access = access,
2908 .ownership = buffer.ownership,
2909 .byte_size = buffer.byte_size,
2910 };
2911 }
2912
2913 const runtime_arguments = try indexing_library.scatterAddRuntimeArguments(instance);
2914 try handle.launch(.{
2915 .artifact = &kernel_artifact,
2916 .loaded_artifact = loaded,
2917 .buffers = bindings[0..],
2918 .scalar_arguments = runtime_arguments[0..],
2919 .geometry = .{
2920 .grid = .{ @intCast((update_count + 127) / 128), 1, 1 },
2921 .threadgroup = .{ instance.threads, 1, 1 },
2922 },
2923 });
2924 try handle.synchronize(.{ .scope = .device });
2925
2926 var actual = @as([bin_count]i32, @splat(0));
2927 try handle.readBuffer(.{ .handle = bindings[0].handle, .bytes = std.mem.sliceAsBytes(actual[0..]) });
2928 try std.testing.expectEqualSlices(i32, expected[0..], actual[0..]);
2929 }
2930 fn expectDeviceScanPipelineOnLiveCuda(
2931 allocator: std.mem.Allocator,
2932 handle: gpu.BackendHandle,
2933 extent: usize,
2934 threads: u32,
2935 mode: kernel.library.scan.PrefixSumMode,
2936 ) !void {
2937 const scan_library = kernel.library.scan;
2938 const instance = scan_library.DeviceScan{ .extent = extent, .threads = threads, .mode = mode };
2939 const stages = try scan_library.deviceScanStages(instance);
2940 const expected_blocks: u32 = @intCast(scan_library.deviceScanBlockCount(extent, threads));
2941 try std.testing.expectEqual(expected_blocks, stages.block_count);
2942
2943 var artifacts = try scan_library.createDeviceScanPipelineArtifacts(allocator, handle, instance, .{
2944 .limits = kernel.Limits.testing,
2945 });
2946 defer artifacts.deinit();
2947 const entries = artifacts.entries();
2948
2949 var owned_pipeline = try scan_library.deviceScanPipeline(allocator, instance);
2950 defer owned_pipeline.deinit();
2951
2952 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{owned_pipeline.value});
2953 defer allocator.free(encoded);
2954 var decoded = try artifact.wire.decode(allocator, encoded);
2955 defer decoded.deinit();
2956 const product_target = try scan_library.deviceScanFamilyTarget(allocator, instance);
2957 defer allocator.free(product_target);
2958 const found = artifact.findPipeline(
2959 decoded.pipelines,
2960 product_target,
2961 scan_library.device_scan_family_version,
2962 ) orelse return error.TestExpectedPipeline;
2963
2964 const data = try allocator.alloc(f32, extent);
2965 defer allocator.free(data);
2966 for (data, 0..) |*value, index| value.* = @floatFromInt((index % 7) + 1);
2967 const zero_dst = try allocator.alloc(f32, extent);
2968 defer allocator.free(zero_dst);
2969 @memset(zero_dst, 0);
2970
2971 const host_views = [_][]const u8{
2972 std.mem.sliceAsBytes(data),
2973 std.mem.sliceAsBytes(zero_dst),
2974 };
2975 var bindings: [2]gpu.BufferBinding = undefined;
2976 for (host_views, 0..) |bytes, slot| {
2977 const buffer = try handle.allocateBuffer(.{
2978 .byte_size = bytes.len,
2979 .alignment = 256,
2980 .dtype = .f32,
2981 .element_count = bytes.len / @sizeOf(f32),
2982 });
2983 try handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });
2984 bindings[slot] = .{
2985 .handle = buffer,
2986 .access = if (slot == 0) .read_only else .read_write,
2987 .ownership = buffer.ownership,
2988 .byte_size = buffer.byte_size,
2989 };
2990 }
2991
2992 try executable.launchPipeline(allocator, handle, .{
2993 .pipeline = found,
2994 .registry = decoded.registry(),
2995 .format = .cuda_ptx,
2996 .operands = bindings[0..1],
2997 .results = bindings[1..2],
2998 .runtime_scalar_arguments = &.{.{ .u32 = @intCast(extent) }},
2999 });
3000
3001 const actual = try allocator.alloc(f32, extent);
3002 defer allocator.free(actual);
3003 try handle.readBuffer(.{ .handle = bindings[1].handle, .bytes = std.mem.sliceAsBytes(actual) });
3004
3005 var running: f32 = 0;
3006 for (data, actual) |value, got| {
3007 switch (mode) {
3008 .inclusive => {
3009 running += value;
3010 try std.testing.expectEqual(running, got);
3011 },
3012 .exclusive => {
3013 try std.testing.expectEqual(running, got);
3014 running += value;
3015 },
3016 }
3017 }
3018 }
3019
3020 test "scan device-wide pipeline launches three stages through the registry on live CUDA" {
3021 const allocator = std.testing.allocator;
3022
3023 var state = try initCudaStateOrSkip(allocator);
3024 defer state.deinit();
3025 const handle = state.handle();
3026
3027 try expectDeviceScanPipelineOnLiveCuda(allocator, handle, 5000, 64, .inclusive);
3028 try expectDeviceScanPipelineOnLiveCuda(allocator, handle, 5000, 64, .exclusive);
3029 try expectDeviceScanPipelineOnLiveCuda(allocator, handle, 33, 32, .inclusive);
3030 }
3031
3032 test "scan device-wide pipeline reaches the million-element cap on live CUDA" {
3033 const allocator = std.testing.allocator;
3034
3035 var state = try initCudaStateOrSkip(allocator);
3036 defer state.deinit();
3037 const handle = state.handle();
3038
3039 try expectDeviceScanPipelineOnLiveCuda(allocator, handle, 1 << 20, 1024, .inclusive);
3040 }
3041
3042 test "scan query selects dispatches and launches the device scan pipeline on live CUDA" {
3043 const allocator = std.testing.allocator;
3044
3045 var state = try initCudaStateOrSkip(allocator);
3046 defer state.deinit();
3047 const handle = state.handle();
3048
3049 const library = kernel.library;
3050 const extent: usize = 5000;
3051
3052 var owned_descriptor = (try library.catalog.selectOwned(allocator, .{ .scan = .{
3053 .dtype = .f32,
3054 .kind = .prefix_sum,
3055 .extent = extent,
3056 } })) orelse return error.TestExpectedCatalogDescriptor;
3057 defer owned_descriptor.deinit();
3058
3059 var package = (try library.createOwnedKernelCallPipelinePackage(allocator, handle, owned_descriptor, .{ .limits = .testing })) orelse {
3060 return error.TestExpectedPipelinePackage;
3061 };
3062 defer package.deinit();
3063
3064 const encoded = try artifact.wire.encode(allocator, package.entries, &.{package.pipeline.value});
3065 defer allocator.free(encoded);
3066 var decoded = try artifact.wire.decode(allocator, encoded);
3067 defer decoded.deinit();
3068 const found = artifact.findPipeline(
3069 decoded.pipelines,
3070 package.pipeline.value.target,
3071 library.scan.device_scan_family_version,
3072 ) orelse return error.TestExpectedPipeline;
3073
3074 const data = try allocator.alloc(f32, extent);
3075 defer allocator.free(data);
3076 for (data, 0..) |*value, index| value.* = @floatFromInt((index % 7) + 1);
3077 const zero_dst = try allocator.alloc(f32, extent);
3078 defer allocator.free(zero_dst);
3079 @memset(zero_dst, 0);
3080
3081 const host_views = [_][]const u8{
3082 std.mem.sliceAsBytes(data),
3083 std.mem.sliceAsBytes(zero_dst),
3084 };
3085 var bindings: [2]gpu.BufferBinding = undefined;
3086 for (host_views, 0..) |bytes, slot| {
3087 const buffer = try handle.allocateBuffer(.{
3088 .byte_size = bytes.len,
3089 .alignment = 256,
3090 .dtype = .f32,
3091 .element_count = bytes.len / @sizeOf(f32),
3092 });
3093 try handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });
3094 bindings[slot] = .{
3095 .handle = buffer,
3096 .access = if (slot == 0) .read_only else .read_write,
3097 .ownership = buffer.ownership,
3098 .byte_size = buffer.byte_size,
3099 };
3100 }
3101
3102 try executable.launchPipeline(allocator, handle, .{
3103 .pipeline = found,
3104 .registry = decoded.registry(),
3105 .format = .cuda_ptx,
3106 .operands = bindings[0..1],
3107 .results = bindings[1..2],
3108 .runtime_scalar_arguments = &.{.{ .u32 = @intCast(extent) }},
3109 });
3110
3111 const actual = try allocator.alloc(f32, extent);
3112 defer allocator.free(actual);
3113 try handle.readBuffer(.{ .handle = bindings[1].handle, .bytes = std.mem.sliceAsBytes(actual) });
3114
3115 var running: f32 = 0;
3116 for (data, actual) |value, got| {
3117 running += value;
3118 try std.testing.expectEqual(running, got);
3119 }
3120 }
3121
3122 test "sort radix split kernels partition stably on live CUDA" {
3123 const allocator = std.testing.allocator;
3124
3125 var state = try initCudaStateOrSkip(allocator);
3126 defer state.deinit();
3127 const handle = state.handle();
3128
3129 const sort_library = kernel.library.sort;
3130 const extent: usize = 5000;
3131 const bit: u32 = 3;
3132 const instance = sort_library.RadixSplit{ .extent = extent, .threads = 64 };
3133 const blocks: u32 = @intCast(sort_library.radixSplitBlockCount(extent, instance.threads));
3134
3135 const keys = try allocator.alloc(i32, extent);
3136 defer allocator.free(keys);
3137 var seed: u32 = 0x2545f491;
3138 for (keys, 0..) |*key, index| {
3139 seed ^= seed << 13;
3140 seed ^= seed >> 17;
3141 seed ^= seed << 5;
3142 key.* = @intCast((seed >> 8) % 100000 * 10 + index % 10);
3143 }
3144
3145 const keys_buffer = try handle.allocateBuffer(.{
3146 .byte_size = extent * @sizeOf(i32),
3147 .alignment = 256,
3148 .dtype = .i32,
3149 .element_count = extent,
3150 });
3151 try handle.writeBuffer(.{ .handle = keys_buffer, .bytes = std.mem.sliceAsBytes(keys) });
3152 const flags_buffer = try handle.allocateBuffer(.{
3153 .byte_size = extent * @sizeOf(f32),
3154 .alignment = 256,
3155 .dtype = .f32,
3156 .element_count = extent,
3157 });
3158 const scanned_buffer = try handle.allocateBuffer(.{
3159 .byte_size = extent * @sizeOf(f32),
3160 .alignment = 256,
3161 .dtype = .f32,
3162 .element_count = extent,
3163 });
3164 const dst_buffer = try handle.allocateBuffer(.{
3165 .byte_size = extent * @sizeOf(i32),
3166 .alignment = 256,
3167 .dtype = .i32,
3168 .element_count = extent,
3169 });
3170
3171 const keys_binding = gpu.BufferBinding{
3172 .handle = keys_buffer,
3173 .access = .read_only,
3174 .ownership = keys_buffer.ownership,
3175 .byte_size = keys_buffer.byte_size,
3176 };
3177 const flags_binding = gpu.BufferBinding{
3178 .handle = flags_buffer,
3179 .access = .read_write,
3180 .ownership = flags_buffer.ownership,
3181 .byte_size = flags_buffer.byte_size,
3182 };
3183 const scanned_binding = gpu.BufferBinding{
3184 .handle = scanned_buffer,
3185 .access = .read_only,
3186 .ownership = scanned_buffer.ownership,
3187 .byte_size = scanned_buffer.byte_size,
3188 };
3189 const dst_binding = gpu.BufferBinding{
3190 .handle = dst_buffer,
3191 .access = .read_write,
3192 .ownership = dst_buffer.ownership,
3193 .byte_size = dst_buffer.byte_size,
3194 };
3195
3196 {
3197 const entry_name = try sort_library.radixSplitFlagsFamilyEntryName(allocator, instance);
3198 defer allocator.free(entry_name);
3199 var graph = try sort_library.RadixSplitFlagsRuntimeFamilyI32.buildNamed(allocator, sort_library.RadixSplitFlagsRuntimeFamilyI32.Limits.testing, entry_name, instance);
3200 defer graph.deinit();
3201 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
3202 .artifact_format = .cuda_ptx,
3203 .authored_kernel_diagnostic_id = "choir/kernel/radix-split-flags-live",
3204 });
3205 defer kernel_artifact.deinit();
3206 const loaded = try handle.loadArtifact(&kernel_artifact);
3207 var stage_bindings = [_]gpu.BufferBinding{ flags_binding, keys_binding };
3208 const runtime_arguments = try sort_library.radixSplitFlagsRuntimeArguments(instance, bit);
3209 try handle.launch(.{
3210 .artifact = &kernel_artifact,
3211 .loaded_artifact = loaded,
3212 .buffers = stage_bindings[0..],
3213 .scalar_arguments = runtime_arguments[0..],
3214 .geometry = .{
3215 .grid = .{ blocks, 1, 1 },
3216 .threadgroup = .{ instance.threads, 1, 1 },
3217 },
3218 });
3219 try handle.synchronize(.{ .scope = .device });
3220 }
3221
3222 const flags = try allocator.alloc(f32, extent);
3223 defer allocator.free(flags);
3224 try handle.readBuffer(.{ .handle = flags_buffer, .bytes = std.mem.sliceAsBytes(flags) });
3225
3226 const scanned = try allocator.alloc(f32, extent);
3227 defer allocator.free(scanned);
3228 var running: f32 = 0;
3229 for (flags, scanned) |flag, *value| {
3230 value.* = running;
3231 running += flag;
3232 }
3233 try handle.writeBuffer(.{ .handle = scanned_buffer, .bytes = std.mem.sliceAsBytes(scanned) });
3234
3235 {
3236 const entry_name = try sort_library.radixSplitScatterFamilyEntryName(allocator, instance);
3237 defer allocator.free(entry_name);
3238 var graph = try sort_library.RadixSplitScatterRuntimeFamilyI32.buildNamed(allocator, sort_library.RadixSplitScatterRuntimeFamilyI32.Limits.testing, entry_name, instance);
3239 defer graph.deinit();
3240 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
3241 .artifact_format = .cuda_ptx,
3242 .authored_kernel_diagnostic_id = "choir/kernel/radix-split-scatter-live",
3243 });
3244 defer kernel_artifact.deinit();
3245 const loaded = try handle.loadArtifact(&kernel_artifact);
3246 var stage_bindings = [_]gpu.BufferBinding{ dst_binding, keys_binding, flags_binding, scanned_binding };
3247 const runtime_arguments = try sort_library.radixSplitScatterRuntimeArguments(instance);
3248 try handle.launch(.{
3249 .artifact = &kernel_artifact,
3250 .loaded_artifact = loaded,
3251 .buffers = stage_bindings[0..],
3252 .scalar_arguments = runtime_arguments[0..],
3253 .geometry = .{
3254 .grid = .{ blocks, 1, 1 },
3255 .threadgroup = .{ instance.threads, 1, 1 },
3256 },
3257 });
3258 try handle.synchronize(.{ .scope = .device });
3259 }
3260
3261 const actual = try allocator.alloc(i32, extent);
3262 defer allocator.free(actual);
3263 try handle.readBuffer(.{ .handle = dst_buffer, .bytes = std.mem.sliceAsBytes(actual) });
3264
3265 const expected = try allocator.alloc(i32, extent);
3266 defer allocator.free(expected);
3267 var count: usize = 0;
3268 const shift: u5 = @intCast(bit);
3269 for (keys) |key| {
3270 if ((key >> shift) & 1 == 0) {
3271 expected[count] = key;
3272 count += 1;
3273 }
3274 }
3275 for (keys) |key| {
3276 if ((key >> shift) & 1 == 1) {
3277 expected[count] = key;
3278 count += 1;
3279 }
3280 }
3281 try std.testing.expectEqualSlices(i32, expected, actual);
3282 }
3283
3284 test "sort radix split pipeline sorts five thousand keys on live CUDA" {
3285 const allocator = std.testing.allocator;
3286
3287 var state = try initCudaStateOrSkip(allocator);
3288 defer state.deinit();
3289 const handle = state.handle();
3290
3291 const sort_library = kernel.library.sort;
3292 const extent: usize = 5000;
3293 const instance = sort_library.RadixSplit{ .extent = extent, .threads = 64 };
3294
3295 var artifacts = try sort_library.createRadixSplitPipelineArtifacts(allocator, handle, instance, .{ .limits = .testing });
3296 defer artifacts.deinit();
3297 const entries = artifacts.entries();
3298 var owned_pipeline = try sort_library.radixSplitPipeline(allocator, instance);
3299 defer owned_pipeline.deinit();
3300
3301 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{owned_pipeline.value});
3302 defer allocator.free(encoded);
3303 var decoded = try artifact.wire.decode(allocator, encoded);
3304 defer decoded.deinit();
3305 const pipeline_target = try sort_library.radixSplitPipelineTarget(allocator, instance);
3306 defer allocator.free(pipeline_target);
3307 const found = artifact.findPipeline(
3308 decoded.pipelines,
3309 pipeline_target,
3310 sort_library.radix_split_family_version,
3311 ) orelse return error.TestExpectedPipeline;
3312
3313 const keys = try allocator.alloc(i32, extent);
3314 defer allocator.free(keys);
3315 var seed: u32 = 0x2545f491;
3316 for (keys) |*key| {
3317 seed ^= seed << 13;
3318 seed ^= seed >> 17;
3319 seed ^= seed << 5;
3320 const magnitude: i32 = @intCast(seed % 1000000);
3321 key.* = if (seed & 1 == 1) -magnitude else magnitude;
3322 }
3323
3324 var buffers: [2]gpu.BufferHandle = undefined;
3325 for (&buffers) |*buffer| {
3326 buffer.* = try handle.allocateBuffer(.{
3327 .byte_size = extent * @sizeOf(i32),
3328 .alignment = 256,
3329 .dtype = .i32,
3330 .element_count = extent,
3331 });
3332 }
3333 try handle.writeBuffer(.{ .handle = buffers[0], .bytes = std.mem.sliceAsBytes(keys) });
3334
3335 var bindings: [2]gpu.BufferBinding = undefined;
3336 for (buffers, 0..) |buffer, index| {
3337 bindings[index] = .{
3338 .handle = buffer,
3339 .access = .read_write,
3340 .ownership = buffer.ownership,
3341 .byte_size = buffer.byte_size,
3342 };
3343 }
3344
3345 const pass_zero_arguments = try sort_library.radixSplitPipelineRuntimeArguments(instance, 0);
3346 const pooled = try executable.allocatePipelineIntermediates(allocator, handle, found, pass_zero_arguments[0..]);
3347 defer executable.deinitPipelineIntermediates(allocator, handle, pooled);
3348 var artifact_pool = try executable.loadPipelineArtifacts(allocator, handle, found, decoded.registry(), .cuda_ptx);
3349 defer artifact_pool.deinit();
3350
3351 var source: usize = 0;
3352 var bit: u32 = 0;
3353 while (bit < sort_library.radix_split_key_bits) : (bit += 1) {
3354 const runtime_arguments = try sort_library.radixSplitPipelineRuntimeArguments(instance, bit);
3355 try executable.launchPipeline(allocator, handle, .{
3356 .pipeline = found,
3357 .registry = decoded.registry(),
3358 .format = .cuda_ptx,
3359 .operands = bindings[source .. source + 1],
3360 .results = bindings[1 - source .. 2 - source],
3361 .runtime_scalar_arguments = runtime_arguments[0..],
3362 .intermediates = pooled,
3363 .artifacts = &artifact_pool,
3364 });
3365 source = 1 - source;
3366 }
3367 try handle.synchronize(.{ .scope = .device });
3368
3369 const actual = try allocator.alloc(i32, extent);
3370 defer allocator.free(actual);
3371 try handle.readBuffer(.{ .handle = bindings[source].handle, .bytes = std.mem.sliceAsBytes(actual) });
3372
3373 const expected = try allocator.alloc(i32, extent);
3374 defer allocator.free(expected);
3375 @memcpy(expected, keys);
3376 std.mem.sort(i32, expected, {}, std.sort.asc(i32));
3377 try std.testing.expectEqualSlices(i32, expected, actual);
3378 }
3379
3380 test "warp ballot rank counts same-predicate lower lanes on live CUDA" {
3381 const allocator = std.testing.allocator;
3382
3383 var state = try initCudaStateOrSkip(allocator);
3384 defer state.deinit();
3385 const handle = state.handle();
3386
3387 var b = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "ballot_rank_live_i32", &.{
3388 kernel.dynamicBuffer(.i32),
3389 kernel.dynamicBuffer(.i32),
3390 });
3391 errdefer b.deinit();
3392 const out = b.argument(0);
3393 const keys = b.argument(1);
3394 const index = try b.globalId(.x);
3395 const key = try b.load(keys, index);
3396 const one = try b.constantInt(.i32, 1);
3397 const low_bit = try b.and_(key, one);
3398 const is_odd = try b.compare(.eq, low_bit, one);
3399 const mask = try b.ballotSync(is_odd);
3400 const lane = try b.laneId();
3401 const lane_i32 = try b.cast(lane, .i32);
3402 const lane_bit = try b.shl(one, lane_i32);
3403 const lower = try b.sub(lane_bit, one);
3404 const below = try b.and_(mask, lower);
3405 const rank = try b.popcount(below);
3406 try b.store(rank, out, index);
3407 try b.return_();
3408 var graph = try b.finish();
3409 defer graph.deinit();
3410
3411 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
3412 .artifact_format = .cuda_ptx,
3413 .authored_kernel_diagnostic_id = "choir/kernel/ballot-rank-live",
3414 });
3415 defer kernel_artifact.deinit();
3416 const loaded = try handle.loadArtifact(&kernel_artifact);
3417
3418 var keys_data: [64]i32 = undefined;
3419 var seed: u32 = 0xdecafbad;
3420 for (&keys_data) |*value| {
3421 seed ^= seed << 13;
3422 seed ^= seed >> 17;
3423 seed ^= seed << 5;
3424 value.* = @intCast(seed % 1000);
3425 }
3426
3427 var zero_out = @as([64]i32, @splat(-1));
3428 const host_views = [_][]const u8{
3429 std.mem.sliceAsBytes(zero_out[0..]),
3430 std.mem.sliceAsBytes(keys_data[0..]),
3431 };
3432 var bindings: [2]gpu.BufferBinding = undefined;
3433 for (host_views, 0..) |bytes, slot| {
3434 const buffer = try handle.allocateBuffer(.{
3435 .byte_size = bytes.len,
3436 .alignment = 256,
3437 .dtype = .i32,
3438 .element_count = bytes.len / @sizeOf(i32),
3439 });
3440 try handle.writeBuffer(.{ .handle = buffer, .bytes = bytes });
3441 bindings[slot] = .{
3442 .handle = buffer,
3443 .access = if (slot == 0) .read_write else .read_only,
3444 .ownership = buffer.ownership,
3445 .byte_size = buffer.byte_size,
3446 };
3447 }
3448
3449 try handle.launch(.{
3450 .artifact = &kernel_artifact,
3451 .loaded_artifact = loaded,
3452 .buffers = bindings[0..],
3453 .geometry = .{ .grid = .{ 2, 1, 1 }, .threadgroup = .{ 32, 1, 1 } },
3454 });
3455 try handle.synchronize(.{ .scope = .device });
3456
3457 var actual = @as([64]i32, @splat(-1));
3458 try handle.readBuffer(.{ .handle = bindings[0].handle, .bytes = std.mem.sliceAsBytes(actual[0..]) });
3459
3460 for (0..64) |element| {
3461 const warp_base = (element / 32) * 32;
3462 var expected: i32 = 0;
3463 for (warp_base..element) |prior| {
3464 if (@mod(keys_data[prior], 2) == 1) expected += 1;
3465 }
3466 try std.testing.expectEqual(expected, actual[element]);
3467 }
3468 }
3469
3470 test "sort radix digit pass partitions stably on live CUDA" {
3471 const allocator = std.testing.allocator;
3472
3473 var state = try initCudaStateOrSkip(allocator);
3474 defer state.deinit();
3475 const handle = state.handle();
3476
3477 const sort_library = kernel.library.sort;
3478 const extent: usize = 5000;
3479 const shift: u32 = 4;
3480 const instance = sort_library.RadixSplit{ .extent = extent, .threads = 64 };
3481 const blocks: u32 = @intCast(sort_library.radixSplitBlockCount(extent, instance.threads));
3482 const cell_count: usize = sort_library.radix_digit_bins * blocks;
3483
3484 const keys = try allocator.alloc(i32, extent);
3485 defer allocator.free(keys);
3486 var seed: u32 = 0x2545f491;
3487 for (keys, 0..) |*key, index| {
3488 seed ^= seed << 13;
3489 seed ^= seed >> 17;
3490 seed ^= seed << 5;
3491 key.* = @intCast((seed >> 8) % 100000 * 10 + index % 10);
3492 }
3493
3494 const keys_buffer = try handle.allocateBuffer(.{
3495 .byte_size = extent * @sizeOf(i32),
3496 .alignment = 256,
3497 .dtype = .i32,
3498 .element_count = extent,
3499 });
3500 try handle.writeBuffer(.{ .handle = keys_buffer, .bytes = std.mem.sliceAsBytes(keys) });
3501 const counts_buffer = try handle.allocateBuffer(.{
3502 .byte_size = cell_count * @sizeOf(f32),
3503 .alignment = 256,
3504 .dtype = .f32,
3505 .element_count = cell_count,
3506 });
3507 const scanned_buffer = try handle.allocateBuffer(.{
3508 .byte_size = cell_count * @sizeOf(f32),
3509 .alignment = 256,
3510 .dtype = .f32,
3511 .element_count = cell_count,
3512 });
3513 const dst_buffer = try handle.allocateBuffer(.{
3514 .byte_size = extent * @sizeOf(i32),
3515 .alignment = 256,
3516 .dtype = .i32,
3517 .element_count = extent,
3518 });
3519
3520 const keys_binding = gpu.BufferBinding{
3521 .handle = keys_buffer,
3522 .access = .read_only,
3523 .ownership = keys_buffer.ownership,
3524 .byte_size = keys_buffer.byte_size,
3525 };
3526 const counts_binding = gpu.BufferBinding{
3527 .handle = counts_buffer,
3528 .access = .read_write,
3529 .ownership = counts_buffer.ownership,
3530 .byte_size = counts_buffer.byte_size,
3531 };
3532 const scanned_binding = gpu.BufferBinding{
3533 .handle = scanned_buffer,
3534 .access = .read_only,
3535 .ownership = scanned_buffer.ownership,
3536 .byte_size = scanned_buffer.byte_size,
3537 };
3538 const dst_binding = gpu.BufferBinding{
3539 .handle = dst_buffer,
3540 .access = .read_write,
3541 .ownership = dst_buffer.ownership,
3542 .byte_size = dst_buffer.byte_size,
3543 };
3544
3545 {
3546 const entry_name = try sort_library.radixDigitHistogramFamilyEntryName(allocator, instance);
3547 defer allocator.free(entry_name);
3548 var graph = try sort_library.RadixDigitHistogramRuntimeFamilyI32.buildNamed(allocator, sort_library.RadixDigitHistogramRuntimeFamilyI32.Limits.testing, entry_name, instance);
3549 defer graph.deinit();
3550 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
3551 .artifact_format = .cuda_ptx,
3552 .authored_kernel_diagnostic_id = "choir/kernel/radix-digit-histogram-live",
3553 });
3554 defer kernel_artifact.deinit();
3555 const loaded = try handle.loadArtifact(&kernel_artifact);
3556 var stage_bindings = [_]gpu.BufferBinding{ counts_binding, keys_binding };
3557 const runtime_arguments = try sort_library.radixDigitHistogramRuntimeArguments(instance, shift);
3558 try handle.launch(.{
3559 .artifact = &kernel_artifact,
3560 .loaded_artifact = loaded,
3561 .buffers = stage_bindings[0..],
3562 .scalar_arguments = runtime_arguments[0..],
3563 .geometry = .{
3564 .grid = .{ blocks, 1, 1 },
3565 .threadgroup = .{ instance.threads, 1, 1 },
3566 },
3567 });
3568 try handle.synchronize(.{ .scope = .device });
3569 }
3570
3571 const counts = try allocator.alloc(f32, cell_count);
3572 defer allocator.free(counts);
3573 try handle.readBuffer(.{ .handle = counts_buffer, .bytes = std.mem.sliceAsBytes(counts) });
3574 const scanned = try allocator.alloc(f32, cell_count);
3575 defer allocator.free(scanned);
3576 var running: f32 = 0;
3577 for (counts, scanned) |count, *value| {
3578 value.* = running;
3579 running += count;
3580 }
3581 try handle.writeBuffer(.{ .handle = scanned_buffer, .bytes = std.mem.sliceAsBytes(scanned) });
3582
3583 {
3584 const entry_name = try sort_library.radixDigitRankScatterFamilyEntryName(allocator, instance);
3585 defer allocator.free(entry_name);
3586 var graph = try sort_library.RadixDigitRankScatterRuntimeFamilyI32.buildNamed(allocator, sort_library.RadixDigitRankScatterRuntimeFamilyI32.Limits.testing, entry_name, instance);
3587 defer graph.deinit();
3588 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
3589 .artifact_format = .cuda_ptx,
3590 .authored_kernel_diagnostic_id = "choir/kernel/radix-digit-rank-scatter-live",
3591 });
3592 defer kernel_artifact.deinit();
3593 const loaded = try handle.loadArtifact(&kernel_artifact);
3594 var stage_bindings = [_]gpu.BufferBinding{ dst_binding, keys_binding, scanned_binding };
3595 const runtime_arguments = try sort_library.radixDigitHistogramRuntimeArguments(instance, shift);
3596 try handle.launch(.{
3597 .artifact = &kernel_artifact,
3598 .loaded_artifact = loaded,
3599 .buffers = stage_bindings[0..],
3600 .scalar_arguments = runtime_arguments[0..],
3601 .geometry = .{
3602 .grid = .{ blocks, 1, 1 },
3603 .threadgroup = .{ instance.threads, 1, 1 },
3604 },
3605 });
3606 try handle.synchronize(.{ .scope = .device });
3607 }
3608
3609 const actual = try allocator.alloc(i32, extent);
3610 defer allocator.free(actual);
3611 try handle.readBuffer(.{ .handle = dst_buffer, .bytes = std.mem.sliceAsBytes(actual) });
3612
3613 const expected = try allocator.alloc(i32, extent);
3614 defer allocator.free(expected);
3615 var count: usize = 0;
3616 const shift_amount: u5 = @intCast(shift);
3617 var digit: i32 = 0;
3618 while (digit < sort_library.radix_digit_bins) : (digit += 1) {
3619 for (keys) |key| {
3620 if ((key >> shift_amount) & (sort_library.radix_digit_bins - 1) == digit) {
3621 expected[count] = key;
3622 count += 1;
3623 }
3624 }
3625 }
3626 try std.testing.expectEqualSlices(i32, expected, actual);
3627 }
3628
3629 test "sort radix digit pipeline sorts in eight passes on live CUDA" {
3630 const allocator = std.testing.allocator;
3631
3632 var state = try initCudaStateOrSkip(allocator);
3633 defer state.deinit();
3634 const handle = state.handle();
3635
3636 const sort_library = kernel.library.sort;
3637 const extent: usize = 5000;
3638 const instance = sort_library.RadixSplit{ .extent = extent, .threads = 64 };
3639
3640 var artifacts = try sort_library.createRadixDigitPipelineArtifacts(allocator, handle, instance, .{ .limits = .testing });
3641 defer artifacts.deinit();
3642 const entries = artifacts.entries();
3643 var owned_pipeline = try sort_library.radixDigitPipeline(allocator, instance);
3644 defer owned_pipeline.deinit();
3645
3646 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{owned_pipeline.value});
3647 defer allocator.free(encoded);
3648 var decoded = try artifact.wire.decode(allocator, encoded);
3649 defer decoded.deinit();
3650 const pipeline_target = try sort_library.radixDigitPipelineTarget(allocator, instance);
3651 defer allocator.free(pipeline_target);
3652 const found = artifact.findPipeline(
3653 decoded.pipelines,
3654 pipeline_target,
3655 sort_library.radix_split_family_version,
3656 ) orelse return error.TestExpectedPipeline;
3657
3658 const keys = try allocator.alloc(i32, extent);
3659 defer allocator.free(keys);
3660 var seed: u32 = 0x2545f491;
3661 for (keys) |*key| {
3662 seed ^= seed << 13;
3663 seed ^= seed >> 17;
3664 seed ^= seed << 5;
3665 const magnitude: i32 = @intCast(seed % 1000000);
3666 key.* = if (seed & 1 == 1) -magnitude else magnitude;
3667 }
3668
3669 var buffers: [2]gpu.BufferHandle = undefined;
3670 for (&buffers) |*buffer| {
3671 buffer.* = try handle.allocateBuffer(.{
3672 .byte_size = extent * @sizeOf(i32),
3673 .alignment = 256,
3674 .dtype = .i32,
3675 .element_count = extent,
3676 });
3677 }
3678 try handle.writeBuffer(.{ .handle = buffers[0], .bytes = std.mem.sliceAsBytes(keys) });
3679
3680 var bindings: [2]gpu.BufferBinding = undefined;
3681 for (buffers, 0..) |buffer, index| {
3682 bindings[index] = .{
3683 .handle = buffer,
3684 .access = .read_write,
3685 .ownership = buffer.ownership,
3686 .byte_size = buffer.byte_size,
3687 };
3688 }
3689
3690 const pass_zero = [_]choir_abi.ScalarArgument{ .{ .u32 = @intCast(extent) }, .{ .u32 = 0 }, .{ .u32 = 0 } };
3691 const pooled = try executable.allocatePipelineIntermediates(allocator, handle, found, pass_zero[0..]);
3692 defer executable.deinitPipelineIntermediates(allocator, handle, pooled);
3693 var artifact_pool = try executable.loadPipelineArtifacts(allocator, handle, found, decoded.registry(), .cuda_ptx);
3694 defer artifact_pool.deinit();
3695
3696 var source: usize = 0;
3697 var shift: u32 = 0;
3698 while (shift < sort_library.radix_split_key_bits) : (shift += sort_library.radix_digit_bits) {
3699 const runtime_arguments = try sort_library.radixDigitHistogramRuntimeArguments(instance, shift);
3700 try executable.launchPipeline(allocator, handle, .{
3701 .pipeline = found,
3702 .registry = decoded.registry(),
3703 .format = .cuda_ptx,
3704 .operands = bindings[source .. source + 1],
3705 .results = bindings[1 - source .. 2 - source],
3706 .runtime_scalar_arguments = runtime_arguments[0..],
3707 .intermediates = pooled,
3708 .artifacts = &artifact_pool,
3709 });
3710 source = 1 - source;
3711 }
3712 try handle.synchronize(.{ .scope = .device });
3713
3714 const actual = try allocator.alloc(i32, extent);
3715 defer allocator.free(actual);
3716 try handle.readBuffer(.{ .handle = bindings[source].handle, .bytes = std.mem.sliceAsBytes(actual) });
3717
3718 const expected = try allocator.alloc(i32, extent);
3719 defer allocator.free(expected);
3720 @memcpy(expected, keys);
3721 std.mem.sort(i32, expected, {}, std.sort.asc(i32));
3722 try std.testing.expectEqualSlices(i32, expected, actual);
3723 }
3724
3725 test "sort radix digit pairs pipeline computes a stable argsort on live CUDA" {
3726 const allocator = std.testing.allocator;
3727
3728 var state = try initCudaStateOrSkip(allocator);
3729 defer state.deinit();
3730 const handle = state.handle();
3731
3732 const sort_library = kernel.library.sort;
3733 const extent: usize = 5000;
3734 const instance = sort_library.RadixSplit{ .extent = extent, .threads = 64 };
3735
3736 var artifacts = try sort_library.createRadixDigitPairsPipelineArtifacts(allocator, handle, instance, .{ .limits = .testing });
3737 defer artifacts.deinit();
3738 const entries = artifacts.entries();
3739 var owned_pipeline = try sort_library.radixDigitPairsPipeline(allocator, instance);
3740 defer owned_pipeline.deinit();
3741
3742 const encoded = try artifact.wire.encode(allocator, entries[0..], &.{owned_pipeline.value});
3743 defer allocator.free(encoded);
3744 var decoded = try artifact.wire.decode(allocator, encoded);
3745 defer decoded.deinit();
3746 const pipeline_target = try sort_library.radixDigitPairsPipelineTarget(allocator, instance);
3747 defer allocator.free(pipeline_target);
3748 const found = artifact.findPipeline(
3749 decoded.pipelines,
3750 pipeline_target,
3751 sort_library.radix_split_family_version,
3752 ) orelse return error.TestExpectedPipeline;
3753
3754 const keys = try allocator.alloc(i32, extent);
3755 defer allocator.free(keys);
3756 const iota = try allocator.alloc(i32, extent);
3757 defer allocator.free(iota);
3758 var seed: u32 = 0x2545f491;
3759 for (keys, iota, 0..) |*key, *value, index| {
3760 seed ^= seed << 13;
3761 seed ^= seed >> 17;
3762 seed ^= seed << 5;
3763 const magnitude: i32 = @intCast(seed % 1000);
3764 key.* = if (seed & 1 == 1) -magnitude else magnitude;
3765 value.* = @intCast(index);
3766 }
3767
3768 var key_buffers: [2]gpu.BufferHandle = undefined;
3769 var value_buffers: [2]gpu.BufferHandle = undefined;
3770 for (&key_buffers, &value_buffers) |*key_buffer, *value_buffer| {
3771 key_buffer.* = try handle.allocateBuffer(.{
3772 .byte_size = extent * @sizeOf(i32),
3773 .alignment = 256,
3774 .dtype = .i32,
3775 .element_count = extent,
3776 });
3777 value_buffer.* = try handle.allocateBuffer(.{
3778 .byte_size = extent * @sizeOf(i32),
3779 .alignment = 256,
3780 .dtype = .i32,
3781 .element_count = extent,
3782 });
3783 }
3784 try handle.writeBuffer(.{ .handle = key_buffers[0], .bytes = std.mem.sliceAsBytes(keys) });
3785 try handle.writeBuffer(.{ .handle = value_buffers[0], .bytes = std.mem.sliceAsBytes(iota) });
3786
3787 var key_bindings: [2]gpu.BufferBinding = undefined;
3788 var value_bindings: [2]gpu.BufferBinding = undefined;
3789 for (key_buffers, value_buffers, 0..) |key_buffer, value_buffer, index| {
3790 key_bindings[index] = .{
3791 .handle = key_buffer,
3792 .access = .read_write,
3793 .ownership = key_buffer.ownership,
3794 .byte_size = key_buffer.byte_size,
3795 };
3796 value_bindings[index] = .{
3797 .handle = value_buffer,
3798 .access = .read_write,
3799 .ownership = value_buffer.ownership,
3800 .byte_size = value_buffer.byte_size,
3801 };
3802 }
3803
3804 const pass_zero = try sort_library.radixDigitHistogramRuntimeArguments(instance, 0);
3805 const pooled = try executable.allocatePipelineIntermediates(allocator, handle, found, pass_zero[0..]);
3806 defer executable.deinitPipelineIntermediates(allocator, handle, pooled);
3807 var artifact_pool = try executable.loadPipelineArtifacts(allocator, handle, found, decoded.registry(), .cuda_ptx);
3808 defer artifact_pool.deinit();
3809
3810 var source: usize = 0;
3811 var shift: u32 = 0;
3812 while (shift < sort_library.radix_split_key_bits) : (shift += sort_library.radix_digit_bits) {
3813 const runtime_arguments = try sort_library.radixDigitHistogramRuntimeArguments(instance, shift);
3814 const operands = [_]gpu.BufferBinding{ key_bindings[source], value_bindings[source] };
3815 const results = [_]gpu.BufferBinding{ key_bindings[1 - source], value_bindings[1 - source] };
3816 try executable.launchPipeline(allocator, handle, .{
3817 .pipeline = found,
3818 .registry = decoded.registry(),
3819 .format = .cuda_ptx,
3820 .operands = operands[0..],
3821 .results = results[0..],
3822 .runtime_scalar_arguments = runtime_arguments[0..],
3823 .intermediates = pooled,
3824 .artifacts = &artifact_pool,
3825 });
3826 source = 1 - source;
3827 }
3828 try handle.synchronize(.{ .scope = .device });
3829
3830 const sorted_keys = try allocator.alloc(i32, extent);
3831 defer allocator.free(sorted_keys);
3832 try handle.readBuffer(.{ .handle = key_bindings[source].handle, .bytes = std.mem.sliceAsBytes(sorted_keys) });
3833 const argsort = try allocator.alloc(i32, extent);
3834 defer allocator.free(argsort);
3835 try handle.readBuffer(.{ .handle = value_bindings[source].handle, .bytes = std.mem.sliceAsBytes(argsort) });
3836
3837 const expected = try allocator.alloc(i32, extent);
3838 defer allocator.free(expected);
3839 @memcpy(expected, keys);
3840 std.mem.sort(i32, expected, {}, std.sort.asc(i32));
3841 try std.testing.expectEqualSlices(i32, expected, sorted_keys);
3842
3843 var previous_index: i32 = -1;
3844 var previous_key: i32 = std.math.minInt(i32);
3845 for (sorted_keys, argsort) |key, original_index| {
3846 try std.testing.expectEqual(keys[@intCast(original_index)], key);
3847 if (key == previous_key) {
3848 try std.testing.expect(original_index > previous_index);
3849 }
3850 previous_index = original_index;
3851 previous_key = key;
3852 }
3853 }
3854
3855 test "spatial uniform grid builds on live CUDA through composed pipelines" {
3856 const allocator = std.testing.allocator;
3857
3858 var state = try initCudaStateOrSkip(allocator);
3859 defer state.deinit();
3860 const handle = state.handle();
3861
3862 const spatial = kernel.library.spatial;
3863 const sort_library = kernel.library.sort;
3864 const scan_library = kernel.library.scan;
3865
3866 const count: usize = 5000;
3867 const geometry = spatial.GridGeometry{
3868 .origin_x = -1.0,
3869 .origin_y = -1.0,
3870 .inv_cell_size = 4.0,
3871 .dims_x = 8,
3872 .dims_y = 8,
3873 };
3874 const plan = try spatial.gridBuildPlan(count, geometry, 64);
3875 const blocks: u32 = @intCast(spatial.gridCellsBlockCount(count, plan.cells_instance.threads));
3876 const cell_count: usize = plan.cells_total;
3877 const count_blocks: u32 = blocks;
3878 const counts_extent: usize = @intCast(plan.offsets_scan.extent);
3879
3880 const xs = try allocator.alloc(f32, count);
3881 defer allocator.free(xs);
3882 const ys = try allocator.alloc(f32, count);
3883 defer allocator.free(ys);
3884 var seed: u32 = 0x2545f491;
3885 for (xs, ys) |*x, *y| {
3886 seed ^= seed << 13;
3887 seed ^= seed >> 17;
3888 seed ^= seed << 5;
3889 x.* = @as(f32, @floatFromInt(seed % 1000)) / 250.0 - 2.0;
3890 seed ^= seed << 13;
3891 seed ^= seed >> 17;
3892 seed ^= seed << 5;
3893 y.* = @as(f32, @floatFromInt(seed % 1000)) / 250.0 - 2.0;
3894 }
3895
3896 const alloc_f32 = struct {
3897 fn run(h: gpu.BackendHandle, n: usize) !gpu.BufferHandle {
3898 return h.allocateBuffer(.{
3899 .byte_size = n * @sizeOf(f32),
3900 .alignment = 256,
3901 .dtype = .f32,
3902 .element_count = n,
3903 });
3904 }
3905 }.run;
3906 const alloc_i32 = struct {
3907 fn run(h: gpu.BackendHandle, n: usize) !gpu.BufferHandle {
3908 return h.allocateBuffer(.{
3909 .byte_size = n * @sizeOf(i32),
3910 .alignment = 256,
3911 .dtype = .i32,
3912 .element_count = n,
3913 });
3914 }
3915 }.run;
3916 const bind = struct {
3917 fn run(buffer: gpu.BufferHandle, access: gpu.BufferAccess) gpu.BufferBinding {
3918 return .{
3919 .handle = buffer,
3920 .access = access,
3921 .ownership = buffer.ownership,
3922 .byte_size = buffer.byte_size,
3923 };
3924 }
3925 }.run;
3926
3927 const x_buffer = try alloc_f32(handle, count);
3928 const y_buffer = try alloc_f32(handle, count);
3929 try handle.writeBuffer(.{ .handle = x_buffer, .bytes = std.mem.sliceAsBytes(xs) });
3930 try handle.writeBuffer(.{ .handle = y_buffer, .bytes = std.mem.sliceAsBytes(ys) });
3931
3932 var id_buffers: [2]gpu.BufferHandle = undefined;
3933 var payload_buffers: [2]gpu.BufferHandle = undefined;
3934 for (&id_buffers, &payload_buffers) |*ids, *payload| {
3935 ids.* = try alloc_i32(handle, count);
3936 payload.* = try alloc_i32(handle, count);
3937 }
3938 const counts_buffer = try alloc_f32(handle, counts_extent);
3939 const scanned_buffer = try alloc_f32(handle, counts_extent);
3940
3941 {
3942 const entry_name = try spatial.gridCellsFamilyEntryName(allocator, plan.cells_instance);
3943 defer allocator.free(entry_name);
3944 var graph = try spatial.GridCellsRuntimeFamilyF32.buildNamed(allocator, spatial.GridCellsRuntimeFamilyF32.Limits.testing, entry_name, plan.cells_instance);
3945 defer graph.deinit();
3946 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
3947 .artifact_format = .cuda_ptx,
3948 .authored_kernel_diagnostic_id = "choir/kernel/grid-cells-live",
3949 });
3950 defer kernel_artifact.deinit();
3951 const loaded = try handle.loadArtifact(&kernel_artifact);
3952 var stage_bindings = [_]gpu.BufferBinding{
3953 bind(id_buffers[0], .read_write),
3954 bind(x_buffer, .read_only),
3955 bind(y_buffer, .read_only),
3956 };
3957 const runtime_arguments = try spatial.gridCellsRuntimeArguments(plan.cells_instance, geometry);
3958 try handle.launch(.{
3959 .artifact = &kernel_artifact,
3960 .loaded_artifact = loaded,
3961 .buffers = stage_bindings[0..],
3962 .scalar_arguments = runtime_arguments[0..],
3963 .geometry = .{
3964 .grid = .{ blocks, 1, 1 },
3965 .threadgroup = .{ plan.cells_instance.threads, 1, 1 },
3966 },
3967 });
3968 try handle.synchronize(.{ .scope = .device });
3969 }
3970
3971 const iota = try allocator.alloc(i32, count);
3972 defer allocator.free(iota);
3973 for (iota, 0..) |*value, index| value.* = @intCast(index);
3974 try handle.writeBuffer(.{ .handle = payload_buffers[0], .bytes = std.mem.sliceAsBytes(iota) });
3975
3976 var sort_artifacts = try sort_library.createRadixDigitPairsPipelineArtifacts(allocator, handle, plan.sort_instance, .{ .limits = .testing });
3977 defer sort_artifacts.deinit();
3978 const sort_entries = sort_artifacts.entries();
3979 var sort_pipeline = try sort_library.radixDigitPairsPipeline(allocator, plan.sort_instance);
3980 defer sort_pipeline.deinit();
3981 const sort_registry = artifact.KernelCallRegistry{ .entries = sort_entries[0..] };
3982
3983 const sort_args0 = try sort_library.radixDigitHistogramRuntimeArguments(plan.sort_instance, 0);
3984 const sort_pool = try executable.allocatePipelineIntermediates(allocator, handle, sort_pipeline.value, sort_args0[0..]);
3985 defer executable.deinitPipelineIntermediates(allocator, handle, sort_pool);
3986 var sort_artifact_pool = try executable.loadPipelineArtifacts(allocator, handle, sort_pipeline.value, sort_registry, .cuda_ptx);
3987 defer sort_artifact_pool.deinit();
3988
3989 var source: usize = 0;
3990 var pass: u32 = 0;
3991 while (pass < plan.sort_passes) : (pass += 1) {
3992 const shift = pass * sort_library.radix_digit_bits;
3993 const runtime_arguments = try sort_library.radixDigitHistogramRuntimeArguments(plan.sort_instance, shift);
3994 const operands = [_]gpu.BufferBinding{
3995 bind(id_buffers[source], .read_write),
3996 bind(payload_buffers[source], .read_write),
3997 };
3998 const results = [_]gpu.BufferBinding{
3999 bind(id_buffers[1 - source], .read_write),
4000 bind(payload_buffers[1 - source], .read_write),
4001 };
4002 try executable.launchPipeline(allocator, handle, .{
4003 .pipeline = sort_pipeline.value,
4004 .registry = sort_registry,
4005 .format = .cuda_ptx,
4006 .operands = operands[0..],
4007 .results = results[0..],
4008 .runtime_scalar_arguments = runtime_arguments[0..],
4009 .intermediates = sort_pool,
4010 .artifacts = &sort_artifact_pool,
4011 });
4012 source = 1 - source;
4013 }
4014 try handle.synchronize(.{ .scope = .device });
4015
4016 {
4017 const entry_name = try spatial.gridCountFamilyEntryName(allocator, plan.count_instance);
4018 defer allocator.free(entry_name);
4019 var graph = try spatial.GridCountRuntimeFamilyI32.buildNamed(allocator, spatial.GridCountRuntimeFamilyI32.Limits.testing, entry_name, plan.count_instance);
4020 defer graph.deinit();
4021 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
4022 .artifact_format = .cuda_ptx,
4023 .authored_kernel_diagnostic_id = "choir/kernel/grid-count-live",
4024 });
4025 defer kernel_artifact.deinit();
4026 const loaded = try handle.loadArtifact(&kernel_artifact);
4027 var stage_bindings = [_]gpu.BufferBinding{
4028 bind(counts_buffer, .read_write),
4029 bind(id_buffers[source], .read_only),
4030 };
4031 const runtime_arguments = try spatial.gridCountRuntimeArguments(plan.count_instance, plan.cells_total);
4032 try handle.launch(.{
4033 .artifact = &kernel_artifact,
4034 .loaded_artifact = loaded,
4035 .buffers = stage_bindings[0..],
4036 .scalar_arguments = runtime_arguments[0..],
4037 .geometry = .{
4038 .grid = .{ count_blocks, 1, 1 },
4039 .threadgroup = .{ plan.count_instance.threads, 1, 1 },
4040 },
4041 });
4042 try handle.synchronize(.{ .scope = .device });
4043 }
4044
4045 var scan_artifacts = try scan_library.createDeviceScanPipelineArtifacts(allocator, handle, plan.offsets_scan, .{
4046 .limits = kernel.Limits.testing,
4047 });
4048 defer scan_artifacts.deinit();
4049 const scan_entries = scan_artifacts.entries();
4050 var scan_pipeline = try scan_library.deviceScanPipeline(allocator, plan.offsets_scan);
4051 defer scan_pipeline.deinit();
4052 const scan_registry = artifact.KernelCallRegistry{ .entries = scan_entries[0..] };
4053 {
4054 const operands = [_]gpu.BufferBinding{bind(counts_buffer, .read_only)};
4055 const results = [_]gpu.BufferBinding{bind(scanned_buffer, .read_write)};
4056 try executable.launchPipeline(allocator, handle, .{
4057 .pipeline = scan_pipeline.value,
4058 .registry = scan_registry,
4059 .format = .cuda_ptx,
4060 .operands = operands[0..],
4061 .results = results[0..],
4062 .runtime_scalar_arguments = &.{.{ .u32 = @intCast(counts_extent) }},
4063 });
4064 }
4065
4066 const sorted_ids = try allocator.alloc(i32, count);
4067 defer allocator.free(sorted_ids);
4068 try handle.readBuffer(.{ .handle = id_buffers[source], .bytes = std.mem.sliceAsBytes(sorted_ids) });
4069 const permutation = try allocator.alloc(i32, count);
4070 defer allocator.free(permutation);
4071 try handle.readBuffer(.{ .handle = payload_buffers[source], .bytes = std.mem.sliceAsBytes(permutation) });
4072 const scanned = try allocator.alloc(f32, counts_extent);
4073 defer allocator.free(scanned);
4074 try handle.readBuffer(.{ .handle = scanned_buffer, .bytes = std.mem.sliceAsBytes(scanned) });
4075
4076 const host_ids = try allocator.alloc(i32, count);
4077 defer allocator.free(host_ids);
4078 for (xs, ys, host_ids) |x, y, *id| {
4079 const fx = (x - geometry.origin_x) * geometry.inv_cell_size;
4080 const fy = (y - geometry.origin_y) * geometry.inv_cell_size;
4081 const cx = std.math.clamp(@as(i32, @intFromFloat(fx)), 0, @as(i32, @intCast(geometry.dims_x - 1)));
4082 const cy = std.math.clamp(@as(i32, @intFromFloat(fy)), 0, @as(i32, @intCast(geometry.dims_y - 1)));
4083 id.* = cy * @as(i32, @intCast(geometry.dims_x)) + cx;
4084 }
4085
4086 var previous_id: i32 = -1;
4087 var previous_perm: i32 = -1;
4088 for (sorted_ids, permutation) |id, original| {
4089 try std.testing.expectEqual(host_ids[@intCast(original)], id);
4090 try std.testing.expect(id >= previous_id);
4091 if (id == previous_id) try std.testing.expect(original > previous_perm);
4092 previous_id = id;
4093 previous_perm = original;
4094 }
4095
4096 const host_counts = try allocator.alloc(u32, cell_count);
4097 defer allocator.free(host_counts);
4098 @memset(host_counts, 0);
4099 for (host_ids) |id| host_counts[@intCast(id)] += 1;
4100 var running: u32 = 0;
4101 for (0..cell_count) |cell| {
4102 const offset: u32 = @intFromFloat(scanned[cell * count_blocks]);
4103 try std.testing.expectEqual(running, offset);
4104 running += host_counts[cell];
4105 }
4106
4107 const radius: f32 = 0.2;
4108 const neighbors_buffer = try alloc_i32(handle, count);
4109 {
4110 const neighbor_instance = spatial.GridNeighborCount{
4111 .count = count,
4112 .threads = plan.cells_instance.threads,
4113 };
4114 const entry_name = try spatial.gridNeighborCountFamilyEntryName(allocator, neighbor_instance);
4115 defer allocator.free(entry_name);
4116 var graph = try spatial.GridNeighborCountRuntimeFamilyF32.buildNamed(allocator, spatial.GridNeighborCountRuntimeFamilyF32.Limits.testing, entry_name, neighbor_instance);
4117 defer graph.deinit();
4118 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
4119 .artifact_format = .cuda_ptx,
4120 .authored_kernel_diagnostic_id = "choir/kernel/grid-neighbor-live",
4121 });
4122 defer kernel_artifact.deinit();
4123 const loaded = try handle.loadArtifact(&kernel_artifact);
4124 var stage_bindings = [_]gpu.BufferBinding{
4125 bind(neighbors_buffer, .read_write),
4126 bind(x_buffer, .read_only),
4127 bind(y_buffer, .read_only),
4128 bind(payload_buffers[source], .read_only),
4129 bind(scanned_buffer, .read_only),
4130 };
4131 const runtime_arguments = try spatial.gridNeighborCountRuntimeArguments(
4132 neighbor_instance,
4133 geometry,
4134 count_blocks,
4135 radius,
4136 );
4137 try handle.launch(.{
4138 .artifact = &kernel_artifact,
4139 .loaded_artifact = loaded,
4140 .buffers = stage_bindings[0..],
4141 .scalar_arguments = runtime_arguments[0..],
4142 .geometry = .{
4143 .grid = .{ blocks, 1, 1 },
4144 .threadgroup = .{ neighbor_instance.threads, 1, 1 },
4145 },
4146 });
4147 try handle.synchronize(.{ .scope = .device });
4148 }
4149
4150 const device_neighbors = try allocator.alloc(i32, count);
4151 defer allocator.free(device_neighbors);
4152 try handle.readBuffer(.{ .handle = neighbors_buffer, .bytes = std.mem.sliceAsBytes(device_neighbors) });
4153
4154 const radius2 = radius * radius;
4155 var any_neighbors = false;
4156 for (0..count) |a| {
4157 var expected_neighbors: i32 = 0;
4158 for (0..count) |b| {
4159 if (a == b) continue;
4160 const dx = xs[a] - xs[b];
4161 const dy = ys[a] - ys[b];
4162 if (dx * dx + dy * dy <= radius2) expected_neighbors += 1;
4163 }
4164 if (expected_neighbors > 0) any_neighbors = true;
4165 try std.testing.expectEqual(expected_neighbors, device_neighbors[a]);
4166 }
4167 try std.testing.expect(any_neighbors);
4168 }
4169
4170 test "linalg batched cholesky factors and solves compose on live CUDA" {
4171 const allocator = std.testing.allocator;
4172
4173 var state = try initCudaStateOrSkip(allocator);
4174 defer state.deinit();
4175 const handle = state.handle();
4176
4177 const factor_library = kernel.library.factor;
4178 const n: usize = 3;
4179 const batch: usize = 64;
4180 const cholesky_instance = factor_library.BatchedCholesky{ .batch = batch, .n = n, .threads = 32 };
4181 const solve_instance = factor_library.BatchedCholeskySolve{ .batch = batch, .n = n, .threads = 32 };
4182 const blocks: u32 = @intCast(factor_library.batchedCholeskyBlockCount(batch, cholesky_instance.threads));
4183
4184 const a = try allocator.alloc(f32, batch * n * n);
4185 defer allocator.free(a);
4186 const rhs = try allocator.alloc(f32, batch * n);
4187 defer allocator.free(rhs);
4188 var seed: u32 = 0x0c0de53d;
4189 for (0..batch) |b| {
4190 var mixed: [n * n]f32 = undefined;
4191 for (&mixed) |*value| {
4192 seed ^= seed << 13;
4193 seed ^= seed >> 17;
4194 seed ^= seed << 5;
4195 value.* = @as(f32, @floatFromInt(seed % 1000)) / 1000.0;
4196 }
4197 for (0..n) |i| {
4198 for (0..n) |j| {
4199 var sum: f32 = 0;
4200 for (0..n) |c| {
4201 sum += mixed[i * n + c] * mixed[j * n + c];
4202 }
4203 if (i == j) sum += @floatFromInt(n);
4204 a[b * n * n + i * n + j] = sum;
4205 }
4206 }
4207 for (rhs[b * n ..][0..n]) |*value| {
4208 seed ^= seed << 13;
4209 seed ^= seed >> 17;
4210 seed ^= seed << 5;
4211 value.* = @as(f32, @floatFromInt(seed % 1000)) / 500.0 - 1.0;
4212 }
4213 }
4214
4215 const alloc_f32 = struct {
4216 fn run(h: gpu.BackendHandle, count: usize) !gpu.BufferHandle {
4217 return h.allocateBuffer(.{
4218 .byte_size = count * @sizeOf(f32),
4219 .alignment = 256,
4220 .dtype = .f32,
4221 .element_count = count,
4222 });
4223 }
4224 }.run;
4225 const bind = struct {
4226 fn run(buffer: gpu.BufferHandle, access: gpu.BufferAccess) gpu.BufferBinding {
4227 return .{
4228 .handle = buffer,
4229 .access = access,
4230 .ownership = buffer.ownership,
4231 .byte_size = buffer.byte_size,
4232 };
4233 }
4234 }.run;
4235
4236 const a_buffer = try alloc_f32(handle, batch * n * n);
4237 const l_buffer = try alloc_f32(handle, batch * n * n);
4238 const rhs_buffer = try alloc_f32(handle, batch * n);
4239 const x_buffer = try alloc_f32(handle, batch * n);
4240 try handle.writeBuffer(.{ .handle = a_buffer, .bytes = std.mem.sliceAsBytes(a) });
4241 try handle.writeBuffer(.{ .handle = rhs_buffer, .bytes = std.mem.sliceAsBytes(rhs) });
4242
4243 {
4244 const entry_name = try factor_library.batchedCholeskyFamilyEntryName(allocator, cholesky_instance);
4245 defer allocator.free(entry_name);
4246 var graph = try factor_library.BatchedCholeskyRuntimeFamilyF32.buildNamed(allocator, factor_library.BatchedCholeskyRuntimeFamilyF32.Limits.testing, entry_name, cholesky_instance);
4247 defer graph.deinit();
4248 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
4249 .artifact_format = .cuda_ptx,
4250 .authored_kernel_diagnostic_id = "choir/kernel/batched-cholesky-live",
4251 });
4252 defer kernel_artifact.deinit();
4253 const loaded = try handle.loadArtifact(&kernel_artifact);
4254 var stage_bindings = [_]gpu.BufferBinding{
4255 bind(l_buffer, .read_write),
4256 bind(a_buffer, .read_only),
4257 };
4258 const runtime_arguments = try factor_library.batchedCholeskyRuntimeArguments(cholesky_instance);
4259 try handle.launch(.{
4260 .artifact = &kernel_artifact,
4261 .loaded_artifact = loaded,
4262 .buffers = stage_bindings[0..],
4263 .scalar_arguments = runtime_arguments[0..],
4264 .geometry = .{
4265 .grid = .{ blocks, 1, 1 },
4266 .threadgroup = .{ cholesky_instance.threads, 1, 1 },
4267 },
4268 });
4269 try handle.synchronize(.{ .scope = .device });
4270 }
4271
4272 {
4273 const entry_name = try factor_library.batchedCholeskySolveFamilyEntryName(allocator, solve_instance);
4274 defer allocator.free(entry_name);
4275 var graph = try factor_library.BatchedCholeskySolveRuntimeFamilyF32.buildNamed(allocator, factor_library.BatchedCholeskySolveRuntimeFamilyF32.Limits.testing, entry_name, solve_instance);
4276 defer graph.deinit();
4277 var kernel_artifact = try kernel.createKernelArtifact(allocator, handle, &graph, .{
4278 .artifact_format = .cuda_ptx,
4279 .authored_kernel_diagnostic_id = "choir/kernel/batched-cholesky-solve-live",
4280 });
4281 defer kernel_artifact.deinit();
4282 const loaded = try handle.loadArtifact(&kernel_artifact);
4283 var stage_bindings = [_]gpu.BufferBinding{
4284 bind(x_buffer, .read_write),
4285 bind(l_buffer, .read_only),
4286 bind(rhs_buffer, .read_only),
4287 };
4288 const runtime_arguments = try factor_library.batchedCholeskySolveRuntimeArguments(solve_instance);
4289 try handle.launch(.{
4290 .artifact = &kernel_artifact,
4291 .loaded_artifact = loaded,
4292 .buffers = stage_bindings[0..],
4293 .scalar_arguments = runtime_arguments[0..],
4294 .geometry = .{
4295 .grid = .{ blocks, 1, 1 },
4296 .threadgroup = .{ solve_instance.threads, 1, 1 },
4297 },
4298 });
4299 try handle.synchronize(.{ .scope = .device });
4300 }
4301
4302 const x = try allocator.alloc(f32, batch * n);
4303 defer allocator.free(x);
4304 try handle.readBuffer(.{ .handle = x_buffer, .bytes = std.mem.sliceAsBytes(x) });
4305
4306 for (0..batch) |b| {
4307 for (0..n) |i| {
4308 var residual: f32 = -rhs[b * n + i];
4309 for (0..n) |j| {
4310 residual += a[b * n * n + i * n + j] * x[b * n + j];
4311 }
4312 try std.testing.expect(@abs(residual) <= 0.0001);
4313 }
4314 }
4315 }
4316
4317 test "accy integration declaration coverage" {
4318 std.testing.refAllDecls(choir);
4319 }