lib/accy/src/target/cpu/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const accy_root = @import("../../root.zig");
4
5 const ir = choir.ir;
6 const dialects = choir.dialects;
7 const ArithDialect = dialects.ArithDialect;
8 const FuncDialect = dialects.FuncDialect;
9 const GpuDialect = dialects.gpu.GpuDialect;
10 const MemrefDialect = dialects.MemrefDialect;
11 const ScfDialect = dialects.ScfDialect;
12 const lowerKernelToHostLoop = choir.backends.gpu.cpu.lowerKernelToHostLoop;
13
14 fn containsOperationNamed(op: *ir.Operation, name: []const u8) bool {
15 if (std.mem.eql(u8, op.name.name, name)) return true;
16 for (op.regions.items) |*region| {
17 var block_iter = region.getBlocks();
18 while (block_iter.next()) |block| {
19 var op_iter = block.getOperations();
20 while (op_iter.next()) |child| {
21 if (containsOperationNamed(child, name)) return true;
22 }
23 }
24 }
25 return false;
26 }
27
28 fn countOperationNamed(op: *ir.Operation, name: []const u8) usize {
29 var count: usize = if (std.mem.eql(u8, op.name.name, name)) 1 else 0;
30 for (op.regions.items) |*region| {
31 var block_iter = region.getBlocks();
32 while (block_iter.next()) |block| {
33 var op_iter = block.getOperations();
34 while (op_iter.next()) |child| {
35 count += countOperationNamed(child, name);
36 }
37 }
38 }
39 return count;
40 }
41
42 fn containsResultTypeNamed(op: *ir.Operation, name: []const u8, type_name: []const u8) bool {
43 if (std.mem.eql(u8, op.name.name, name)) {
44 for (op.results.items) |result| {
45 if (result.type.getDialectTypeName()) |found| {
46 if (std.mem.eql(u8, found, type_name)) return true;
47 }
48 }
49 }
50 for (op.regions.items) |*region| {
51 var block_iter = region.getBlocks();
52 while (block_iter.next()) |block| {
53 var op_iter = block.getOperations();
54 while (op_iter.next()) |child| {
55 if (containsResultTypeNamed(child, name, type_name)) return true;
56 }
57 }
58 }
59 return false;
60 }
61
62 fn containsStoreValueTypeNamed(op: *ir.Operation, type_name: []const u8) bool {
63 if (std.mem.eql(u8, op.name.name, MemrefDialect.StoreOp.operation_name)) {
64 const store = MemrefDialect.StoreOp{ .op = op };
65 if (store.getValue().type.getDialectTypeName()) |found| {
66 if (std.mem.eql(u8, found, type_name)) return true;
67 }
68 }
69 for (op.regions.items) |*region| {
70 var block_iter = region.getBlocks();
71 while (block_iter.next()) |block| {
72 var op_iter = block.getOperations();
73 while (op_iter.next()) |child| {
74 if (containsStoreValueTypeNamed(child, type_name)) return true;
75 }
76 }
77 }
78 return false;
79 }
80
81 fn firstFunction(module: *ir.Operation) ?FuncDialect.FuncOp {
82 const region = module.getRegion(0) orelse return null;
83 const block = region.getEntryBlock() orelse return null;
84 var iter = block.getOperations();
85 while (iter.next()) |op| {
86 if (std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) return .{ .op = op };
87 }
88 return null;
89 }
90
91 test "cpu lowering wraps global x kernel body in host scf loop" {
92 const allocator = std.testing.allocator;
93 const kernel = accy_root.kernel;
94
95 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_copy_f32", &.{
96 kernel.dynamicBuffer(.f32),
97 kernel.dynamicBuffer(.f32),
98 });
99 errdefer builder.deinit();
100
101 const src = builder.argument(0);
102 const dst = builder.argument(1);
103 const index = try builder.globalId(.x);
104 const value = try builder.load(src, index);
105 try builder.store(value, dst, index);
106 try builder.return_();
107
108 var program = try builder.finish();
109 defer program.deinit();
110
111 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{ .entry_name = "cpu_copy_f32" });
112 defer lowered.erase();
113
114 try ir.verifyOperation(lowered, ir.verify.default_options);
115 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
116 try std.testing.expect(!host_func.isKernel());
117 try std.testing.expectEqual(@as(usize, 9), host_func.getNumArguments());
118 try std.testing.expect(containsOperationNamed(host_func.op, ScfDialect.ForOp.operation_name));
119 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.GlobalIdxOp.operation_name));
120 try std.testing.expect(containsOperationNamed(host_func.op, ArithDialect.RemOp.operation_name));
121 try std.testing.expect(containsOperationNamed(host_func.op, dialects.MemrefDialect.LoadOp.operation_name));
122 try std.testing.expect(containsOperationNamed(host_func.op, dialects.MemrefDialect.StoreOp.operation_name));
123 }
124
125 test "cpu lowering vectorizes straight line global x add kernel" {
126 const allocator = std.testing.allocator;
127 const kernel = accy_root.kernel;
128
129 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_vectorized_add_f32", &.{
130 kernel.dynamicBuffer(.f32),
131 kernel.dynamicBuffer(.f32),
132 kernel.dynamicBuffer(.f32),
133 });
134 errdefer builder.deinit();
135
136 const axis = try builder.axis("i", 4);
137 try builder.bind(axis, .thread_x);
138 const dst = builder.argument(0);
139 const lhs = builder.argument(1);
140 const rhs = builder.argument(2);
141 const index = try builder.globalId(.x);
142 const lhs_value = try builder.load(lhs, index);
143 const rhs_value = try builder.load(rhs, index);
144 const sum = try builder.add(lhs_value, rhs_value);
145 try builder.store(sum, dst, index);
146 try builder.return_();
147
148 var program = try builder.finish();
149 defer program.deinit();
150
151 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{
152 .entry_name = "cpu_vectorized_add_f32",
153 .vector_width = 4,
154 });
155 defer lowered.erase();
156
157 try ir.verifyOperation(lowered, ir.verify.default_options);
158 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
159 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.GlobalIdxOp.operation_name));
160 try std.testing.expect(containsResultTypeNamed(host_func.op, MemrefDialect.LoadOp.operation_name, "arith.vec4xf32"));
161 try std.testing.expect(containsResultTypeNamed(host_func.op, ArithDialect.AddOp.operation_name, "arith.vec4xf32"));
162 try std.testing.expect(containsStoreValueTypeNamed(host_func.op, "arith.vec4xf32"));
163 }
164
165 test "cpu lowering leaves unsupported packed cpu vector width unvectorized" {
166 const allocator = std.testing.allocator;
167 const kernel = accy_root.kernel;
168
169 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_f64_add_width_four", &.{
170 kernel.dynamicBuffer(.f64),
171 kernel.dynamicBuffer(.f64),
172 kernel.dynamicBuffer(.f64),
173 });
174 errdefer builder.deinit();
175
176 const axis = try builder.axis("i", 4);
177 try builder.bind(axis, .thread_x);
178 const dst = builder.argument(0);
179 const lhs = builder.argument(1);
180 const rhs = builder.argument(2);
181 const index = try builder.globalId(.x);
182 const lhs_value = try builder.load(lhs, index);
183 const rhs_value = try builder.load(rhs, index);
184 const sum = try builder.add(lhs_value, rhs_value);
185 try builder.store(sum, dst, index);
186 try builder.return_();
187
188 var program = try builder.finish();
189 defer program.deinit();
190
191 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{
192 .entry_name = "cpu_f64_add_width_four",
193 .vector_width = 4,
194 });
195 defer lowered.erase();
196
197 try ir.verifyOperation(lowered, ir.verify.default_options);
198 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
199 try std.testing.expect(!containsResultTypeNamed(host_func.op, MemrefDialect.LoadOp.operation_name, "arith.vec4xf64"));
200 try std.testing.expect(!containsResultTypeNamed(host_func.op, ArithDialect.AddOp.operation_name, "arith.vec4xf64"));
201 try std.testing.expect(!containsStoreValueTypeNamed(host_func.op, "arith.vec4xf64"));
202 }
203
204 test "cpu lowering vectorizes straight line global x f64 add kernel at supported width" {
205 const allocator = std.testing.allocator;
206 const kernel = accy_root.kernel;
207
208 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_vectorized_add_f64", &.{
209 kernel.dynamicBuffer(.f64),
210 kernel.dynamicBuffer(.f64),
211 kernel.dynamicBuffer(.f64),
212 });
213 errdefer builder.deinit();
214
215 const axis = try builder.axis("i", 2);
216 try builder.bind(axis, .thread_x);
217 const dst = builder.argument(0);
218 const lhs = builder.argument(1);
219 const rhs = builder.argument(2);
220 const index = try builder.globalId(.x);
221 const lhs_value = try builder.load(lhs, index);
222 const rhs_value = try builder.load(rhs, index);
223 const sum = try builder.add(lhs_value, rhs_value);
224 try builder.store(sum, dst, index);
225 try builder.return_();
226
227 var program = try builder.finish();
228 defer program.deinit();
229
230 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{
231 .entry_name = "cpu_vectorized_add_f64",
232 .vector_width = 2,
233 });
234 defer lowered.erase();
235
236 try ir.verifyOperation(lowered, ir.verify.default_options);
237 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
238 try std.testing.expect(containsResultTypeNamed(host_func.op, MemrefDialect.LoadOp.operation_name, "arith.vec2xf64"));
239 try std.testing.expect(containsResultTypeNamed(host_func.op, ArithDialect.AddOp.operation_name, "arith.vec2xf64"));
240 try std.testing.expect(containsStoreValueTypeNamed(host_func.op, "arith.vec2xf64"));
241 }
242
243 test "cpu lowering vectorizes scalar argument splat and scalar tail" {
244 const allocator = std.testing.allocator;
245 const kernel = accy_root.kernel;
246
247 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_scalar_argument_add_f32", &.{
248 kernel.dynamicBuffer(.f32),
249 kernel.scalar(.f32),
250 kernel.dynamicBuffer(.f32),
251 });
252 errdefer builder.deinit();
253
254 const axis = try builder.axis("i", 6);
255 try builder.bind(axis, .thread_x);
256 const dst = builder.argument(0);
257 const bias = builder.argument(1);
258 const src = builder.argument(2);
259 const index = try builder.globalId(.x);
260 const value = try builder.load(src, index);
261 const sum = try builder.add(value, bias);
262 try builder.store(sum, dst, index);
263 try builder.return_();
264
265 var program = try builder.finish();
266 defer program.deinit();
267
268 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{
269 .entry_name = "cpu_scalar_argument_add_f32",
270 .vector_width = 4,
271 });
272 defer lowered.erase();
273
274 try ir.verifyOperation(lowered, ir.verify.default_options);
275 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
276 try std.testing.expectEqual(@as(usize, 2), countOperationNamed(host_func.op, ScfDialect.ForOp.operation_name));
277 try std.testing.expect(containsResultTypeNamed(host_func.op, MemrefDialect.LoadOp.operation_name, "arith.vec4xf32"));
278 try std.testing.expect(containsResultTypeNamed(host_func.op, ArithDialect.SplatOp.operation_name, "arith.vec4xf32"));
279 try std.testing.expect(containsResultTypeNamed(host_func.op, ArithDialect.AddOp.operation_name, "arith.vec4xf32"));
280 try std.testing.expect(containsStoreValueTypeNamed(host_func.op, "arith.vec4xf32"));
281 }
282
283 test "cpu lowering vectorizes expanded pure op coverage" {
284 const allocator = std.testing.allocator;
285 const kernel = accy_root.kernel;
286
287 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_vectorized_max_neg_f32", &.{
288 kernel.dynamicBuffer(.f32),
289 kernel.dynamicBuffer(.f32),
290 kernel.dynamicBuffer(.f32),
291 });
292 errdefer builder.deinit();
293
294 const axis = try builder.axis("i", 4);
295 try builder.bind(axis, .thread_x);
296 const dst = builder.argument(0);
297 const lhs = builder.argument(1);
298 const rhs = builder.argument(2);
299 const index = try builder.globalId(.x);
300 const lhs_value = try builder.load(lhs, index);
301 const rhs_value = try builder.load(rhs, index);
302 const larger = try builder.max(lhs_value, rhs_value);
303 const value = try builder.neg(larger);
304 try builder.store(value, dst, index);
305 try builder.return_();
306
307 var program = try builder.finish();
308 defer program.deinit();
309
310 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{
311 .entry_name = "cpu_vectorized_max_neg_f32",
312 .vector_width = 4,
313 });
314 defer lowered.erase();
315
316 try ir.verifyOperation(lowered, ir.verify.default_options);
317 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
318 try std.testing.expect(containsResultTypeNamed(host_func.op, ArithDialect.MaxOp.operation_name, "arith.vec4xf32"));
319 try std.testing.expect(containsResultTypeNamed(host_func.op, ArithDialect.NegOp.operation_name, "arith.vec4xf32"));
320 try std.testing.expect(containsStoreValueTypeNamed(host_func.op, "arith.vec4xf32"));
321 }
322
323 test "cpu lowering derives global y and z from flat launch shape" {
324 const allocator = std.testing.allocator;
325 const kernel = accy_root.kernel;
326
327 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_global_yz", &.{});
328 errdefer builder.deinit();
329
330 _ = try builder.globalId(.y);
331 _ = try builder.globalId(.z);
332 try builder.return_();
333
334 var program = try builder.finish();
335 defer program.deinit();
336
337 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{ .entry_name = "cpu_global_yz" });
338 defer lowered.erase();
339
340 try ir.verifyOperation(lowered, ir.verify.default_options);
341 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
342 try std.testing.expectEqual(@as(usize, 7), host_func.getNumArguments());
343 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.GlobalIdxOp.operation_name));
344 try std.testing.expect(containsOperationNamed(host_func.op, ArithDialect.DivOp.operation_name));
345 try std.testing.expect(containsOperationNamed(host_func.op, ArithDialect.RemOp.operation_name));
346 try std.testing.expect(containsOperationNamed(host_func.op, ArithDialect.MulOp.operation_name));
347 }
348
349 test "cpu lowering derives thread block and launch dimension ops" {
350 const allocator = std.testing.allocator;
351 const kernel = accy_root.kernel;
352
353 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_thread_block_dims", &.{});
354 errdefer builder.deinit();
355
356 _ = try builder.threadId(.x);
357 _ = try builder.blockId(.x);
358 _ = try builder.blockDim(.x);
359 _ = try builder.gridDim(.x);
360 try builder.return_();
361
362 var program = try builder.finish();
363 defer program.deinit();
364
365 const lowered = try lowerKernelToHostLoop(allocator, program.kernelModule(), .{ .entry_name = "cpu_thread_block_dims" });
366 defer lowered.erase();
367
368 try ir.verifyOperation(lowered, ir.verify.default_options);
369 const host_func = firstFunction(lowered) orelse return error.ExpectedHostFunction;
370 try std.testing.expectEqual(@as(usize, 7), host_func.getNumArguments());
371 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.ThreadIdxOp.operation_name));
372 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.BlockIdxOp.operation_name));
373 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.BlockDimOp.operation_name));
374 try std.testing.expect(!containsOperationNamed(host_func.op, GpuDialect.GridDimOp.operation_name));
375 try std.testing.expect(containsOperationNamed(host_func.op, ArithDialect.DivOp.operation_name));
376 try std.testing.expect(containsOperationNamed(host_func.op, ArithDialect.RemOp.operation_name));
377 }
378
379 test "cpu lowering rejects unsupported gpu lane context ops" {
380 const allocator = std.testing.allocator;
381 const kernel = accy_root.kernel;
382
383 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "cpu_lane_id_reject", &.{});
384 errdefer builder.deinit();
385
386 _ = try builder.laneId();
387 try builder.return_();
388
389 var program = try builder.finish();
390 defer program.deinit();
391
392 try std.testing.expectError(
393 error.UnsupportedOperation,
394 lowerKernelToHostLoop(allocator, program.kernelModule(), .{ .entry_name = "cpu_lane_id_reject" }),
395 );
396 }