lib/choir/src/backends/gpu/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("../../root.zig");
3 const backends_gpu = @import("root.zig");
4 const gpu = @import("../../dialects/gpu/root.zig");
5 const spirv = backends_gpu.spirv;
6 const nvptx = backends_gpu.nvptx;
7 const lowering = backends_gpu.lowering;
8 const gpu_registration = backends_gpu.registration;
9
10 const GpuDialect = gpu.GpuDialect;
11 const SpirvDialect = spirv.SpirvDialect;
12 const NvptxDialect = nvptx.NvptxDialect;
13 const package_extension = gpu_registration.package_extension;
14 const dialect_registrations = gpu_registration.dialect_registrations;
15 const pass_registrations = lowering.pass_registrations;
16 const pipeline_registrations = lowering.pipeline_registrations;
17 const gpu_to_spirv_pipeline_name = lowering.gpu_to_spirv_pipeline_name;
18 const gpu_to_nvptx_pipeline_name = lowering.gpu_to_nvptx_pipeline_name;
19 const target_lowering_pipeline_count = lowering.target_lowering_pipeline_count;
20 const addTargetLoweringPipeline = lowering.addTargetLoweringPipeline;
21 const addGpuToNvptxPipeline = lowering.addGpuToNvptxPipeline;
22
23 test {
24 _ = @import("features.zig");
25 _ = @import("subgroup.zig");
26 _ = @import("cpu/test.zig");
27 _ = @import("metal/test.zig");
28 _ = @import("nvptx/test.zig");
29 _ = @import("spirv/test.zig");
30 _ = @import("webgpu/test.zig");
31 }
32
33 fn expectSingleTargetLoweringPass(
34 pm: *const choir.passes.PassManager,
35 pass_name: []const u8,
36 ) !void {
37 const testing = std.testing;
38
39 try testing.expectEqual(@as(usize, 1), pm.root.pipeline.items.len);
40 switch (pm.root.pipeline.items[0]) {
41 .pass => |pass| try testing.expectEqualStrings(pass_name, pass.name),
42 .nested => try testing.expect(false),
43 }
44 }
45
46 test "gpu package registers accelerator dialects and target lowering entries" {
47 const testing = std.testing;
48 const allocator = testing.allocator;
49
50 var ctx = try choir.ir.Context.init(allocator, choir.ir.Context.Limits.testing);
51 defer ctx.deinit(allocator);
52
53 try ctx.requireRegistered();
54
55 try testing.expectError(error.UnknownDialect, ctx.getOrLoadDialect("gpu"));
56
57 var extension_registry = choir.extensions.ExtensionRegistry.init(allocator);
58 defer extension_registry.deinit();
59 try extension_registry.registerPackage(&ctx, package_extension);
60
61 inline for (dialect_registrations) |registration| {
62 _ = try ctx.getOrLoadDialect(registration.name);
63 }
64
65 try testing.expect(ctx.lookupOperation(GpuDialect.ThreadIdxOp.operation_name) != null);
66 try testing.expect(ctx.lookupType(gpu.type_names.tma_desc) != null);
67 try testing.expect(ctx.lookupOperation(SpirvDialect.ModuleOp.operation_name) != null);
68 try testing.expect(ctx.lookupOperation(NvptxDialect.ThreadIdxOp.operation_name) != null);
69 try testing.expect(!ctx.isBackendDialect("gpu"));
70 try testing.expect(ctx.isBackendDialect("spirv"));
71 try testing.expect(ctx.isBackendDialect("nvptx"));
72
73 var pass_registry = choir.passes.PassRegistry.init(allocator);
74 defer pass_registry.deinit();
75 try extension_registry.registerPassEntriesTo(&pass_registry);
76
77 inline for (pass_registrations) |registration| {
78 try testing.expect(pass_registry.lookupPass(registration.name) != null);
79 }
80 inline for (pipeline_registrations) |registration| {
81 try testing.expect(pass_registry.lookupPipeline(registration.name) != null);
82 }
83 }
84
85 test "target lowering pipelines materialize from textual Choir registry" {
86 const testing = std.testing;
87 const allocator = testing.allocator;
88
89 var ctx = try choir.ir.Context.init(allocator, choir.ir.Context.Limits.testing);
90 defer ctx.deinit(allocator);
91
92 var extension_registry = choir.extensions.ExtensionRegistry.init(allocator);
93 defer extension_registry.deinit();
94 try extension_registry.registerPackage(&ctx, package_extension);
95
96 var pass_registry = choir.passes.PassRegistry.init(allocator);
97 defer pass_registry.deinit();
98 try extension_registry.registerPassEntriesTo(&pass_registry);
99
100 var pm = choir.passes.PassManager.init(allocator);
101 defer pm.deinit();
102 try choir.passes.parsePassPipeline(
103 &pass_registry,
104 gpu_to_spirv_pipeline_name ++ "," ++ gpu_to_nvptx_pipeline_name,
105 &pm,
106 );
107 try testing.expectEqual(@as(usize, target_lowering_pipeline_count), pm.root.pipeline.items.len);
108
109 const text = try choir.passes.formatPassManagerPipelineAlloc(allocator, &pm);
110 defer allocator.free(text);
111 const expected = spirv.gpu_to_spirv_pass_name ++ "," ++ nvptx.gpu_to_nvptx_pass_name;
112 try testing.expectEqualStrings(expected, text);
113 }
114
115 test "target lowering pipelines preload produced accelerator dialects" {
116 const testing = std.testing;
117 const allocator = testing.allocator;
118
119 var ctx = try choir.ir.Context.init(allocator, choir.ir.Context.Limits.testing);
120 defer ctx.deinit(allocator);
121
122 try ctx.requireRegistered();
123
124 try choir.dialects.registerChoirDialect(&ctx);
125 var extension_registry = choir.extensions.ExtensionRegistry.init(allocator);
126 defer extension_registry.deinit();
127 try extension_registry.registerPackage(&ctx, package_extension);
128
129 _ = try ctx.getOrLoadDialect("builtin");
130 const location = choir.ir.Location.getUnknown();
131 const builtin_module = try choir.dialects.BuiltinDialect.ModuleOp.create(&ctx, location);
132
133 try testing.expect(!ctx.isDialectLoaded("spirv"));
134 var spirv_pm = choir.passes.PassManager.init(allocator);
135 defer spirv_pm.deinit();
136 try addTargetLoweringPipeline(&spirv_pm, .spirv);
137 try expectSingleTargetLoweringPass(&spirv_pm, spirv.gpu_to_spirv_pass_name);
138 try testing.expectEqual(choir.passes.PassResult.success, spirv_pm.run(builtin_module.op, &ctx));
139 try testing.expect(ctx.isDialectLoaded("spirv"));
140
141 try testing.expect(!ctx.isDialectLoaded("nvptx"));
142 var nvptx_pm = choir.passes.PassManager.init(allocator);
143 defer nvptx_pm.deinit();
144 try addGpuToNvptxPipeline(&nvptx_pm);
145 try expectSingleTargetLoweringPass(&nvptx_pm, nvptx.gpu_to_nvptx_pass_name);
146 try testing.expectEqual(choir.passes.PassResult.success, nvptx_pm.run(builtin_module.op, &ctx));
147 try testing.expect(ctx.isDialectLoaded("nvptx"));
148 }
149
150 test "gpu backends declaration coverage" {
151 std.testing.refAllDecls(backends_gpu);
152 }
153
154 const stages = @import("fixture/stages.zig");
155 const calls = @import("calls.zig");
156
157 /// `case`'s source with its one `needle` replaced. The caller owns it.
158 fn stageVariant(case: stages.Case, needle: []const u8, replacement: []const u8) ![]u8 {
159 try std.testing.expectEqual(@as(usize, 1), std.mem.count(u8, case.source, needle));
160 return std.mem.replaceOwned(u8, std.testing.allocator, case.source, needle, replacement);
161 }
162
163 /// Runs SPIR-V, MSL and the CPU twin over `source` and expects each to refuse it as
164 /// `UnsupportedStageMemory`, the twin naming `operation`.
165 fn expectStageMemoryRefused(case: stages.Case, source: []const u8, operation: []const u8) !void {
166 const allocator = std.testing.allocator;
167 {
168 var session: stages.Session = undefined;
169 try session.init(allocator, .testing, source, true);
170 defer session.deinit(allocator);
171 try std.testing.expectError(error.UnsupportedStageMemory, stages.emitSpirv(allocator, session.module));
172 }
173 var session: stages.Session = undefined;
174 try session.init(allocator, .testing, source, false);
175 defer session.deinit(allocator);
176 try std.testing.expectError(error.UnsupportedStageMemory, stages.emitMsl(allocator, session.module, case));
177 var refusal: backends_gpu.cpu.StageRefusal = .{};
178 try std.testing.expectError(error.UnsupportedStageMemory, backends_gpu.cpu.lowerStagesToHost(allocator, session.module, &refusal));
179 try std.testing.expectEqualStrings(operation, refusal.operation);
180 }
181
182 fn expectHelperSignatureRefused(case: stages.Case, helper: []const u8) !void {
183 const allocator = std.testing.allocator;
184 {
185 var session: stages.Session = undefined;
186 try session.init(allocator, .testing, case.source, true);
187 defer session.deinit(allocator);
188 try std.testing.expectError(error.UnsupportedHelperSignature, stages.emitSpirv(allocator, session.module));
189 }
190 var session: stages.Session = undefined;
191 try session.init(allocator, .testing, case.source, false);
192 defer session.deinit(allocator);
193 try std.testing.expectError(error.UnsupportedHelperSignature, stages.emitMsl(allocator, session.module, case));
194 var refusal: backends_gpu.cpu.StageRefusal = .{};
195 try std.testing.expectError(error.UnsupportedHelperSignature, backends_gpu.cpu.lowerStagesToHost(allocator, session.module, &refusal));
196 try std.testing.expectEqualStrings(choir.dialects.func.FuncDialect.FuncOp.operation_name, refusal.operation);
197 try std.testing.expectEqualStrings(helper, refusal.helper);
198 }
199
200 test "every stage backend names an unsupported helper memref parameter" {
201 try expectHelperSignatureRefused(.{
202 .name = "helper_param",
203 .source = @embedFile("fixture/helper_param.txt"),
204 .vertex = "helper_vertex",
205 .fragment = "helper_fragment",
206 }, "read_material");
207 }
208
209 test "every stage backend names an unsupported helper memref result" {
210 try expectHelperSignatureRefused(.{
211 .name = "helper_result",
212 .source = @embedFile("fixture/helper_result.txt"),
213 .vertex = "helper_vertex",
214 .fragment = "helper_fragment",
215 }, "borrow");
216 }
217
218 test "every stage backend refuses a helper that reads a global by one name" {
219 const Memref = choir.dialects.MemrefDialect;
220 const case = stages.calls_case;
221 const source = try stageVariant(
222 case,
223 "%5 = memref.alloca() : !memref<1,arith.f32,local>",
224 "%5 = memref.get_global() {sym_name = \"table\"} : !memref<1,arith.f32,device>",
225 );
226 defer std.testing.allocator.free(source);
227 try expectStageMemoryRefused(case, source, Memref.GetGlobalOp.operation_name);
228 }
229
230 test "every stage backend refuses a stage that allocates shared memory by one name" {
231 const Memref = choir.dialects.MemrefDialect;
232 const case = stages.alloca_case;
233 const source = try stageVariant(
234 case,
235 "%16 = memref.alloca() : !memref<4,arith.f32,local>",
236 "%16 = memref.alloc() : !memref<4,arith.f32,shared>",
237 );
238 defer std.testing.allocator.free(source);
239 try expectStageMemoryRefused(case, source, Memref.AllocOp.operation_name);
240 }
241
242 test "a kernel keeps shared memory and atomics through the call plan" {
243 const allocator = std.testing.allocator;
244 var session: stages.Session = undefined;
245 try session.init(allocator, .testing, @embedFile("metal/fixture/shared.txt"), false);
246 defer session.deinit(allocator);
247 var plan = try calls.Plan.init(allocator, session.module, null);
248 defer plan.deinit();
249 try std.testing.expectEqual(@as(usize, 0), plan.helpers.items.len);
250 }
251
252 test "a stage module's push extent is the one the SPIR-V emitter declares" {
253 const allocator = std.testing.allocator;
254 const all = stages.cases ++ [_]stages.Case{ stages.alloca_case, stages.contraction_case, stages.gradient_case, stages.calls_case };
255 var reads: u32 = 0;
256 for (all) |case| {
257 var session: stages.Session = undefined;
258 try session.init(allocator, .testing, case.source, true);
259 defer session.deinit(allocator);
260 var codegen = spirv.emitter.SpirvCodegen.init(allocator);
261 defer codegen.deinit();
262 const words = try codegen.emitModuleWords(session.module);
263 allocator.free(words);
264 const extent = gpu.stage.pushExtent(session.module).?;
265 try std.testing.expectEqual(codegen.push_extent, extent);
266 if (std.mem.eql(u8, case.name, "pushed")) try std.testing.expectEqual(@as(u32, 16), extent);
267 if (std.mem.eql(u8, case.name, "contraction")) try std.testing.expectEqual(@as(u32, 20), extent);
268 if (extent != 0) reads += 1;
269 }
270 try std.testing.expectEqual(@as(u32, 2), reads);
271 }