lib/accy/src/target/webgpu/test.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const choir = @import("choir");
  3 const accy_root = @import("../../root.zig");
  4 
  5 const emitWgsl = choir.backends.gpu.webgpu.wgsl.emitWgsl;
  6 
  7 test "webgpu choir wgsl emitter lowers simple elementwise kernel" {
  8     const testing = std.testing;
  9     const builder_mod = accy_root.kernel;
 10     const params = [_]builder_mod.Param{
 11         builder_mod.dynamicBuffer(.f32),
 12         builder_mod.dynamicBuffer(.f32),
 13         builder_mod.dynamicBuffer(.f32),
 14         builder_mod.dynamicBuffer(.i32),
 15     };
 16 
 17     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_add", &params);
 18     errdefer builder.deinit();
 19     const out = builder.argument(0);
 20     const lhs = builder.argument(1);
 21     const rhs = builder.argument(2);
 22     const count_buffer = builder.argument(3);
 23     const index = try builder.globalId(.x);
 24     const zero = try builder.constantIndex(0);
 25     const count = try builder.load(count_buffer, zero);
 26     const index_i32 = try builder.cast(index, .i32);
 27     const in_bounds = try builder.compare(.lt, index_i32, count);
 28     var if_op = try builder.if_(in_bounds, &.{});
 29     const previous = builder.insertionBlock();
 30     builder.setInsertionBlock(if_op.thenBlock());
 31     const left = try builder.load(lhs, index);
 32     const right = try builder.load(rhs, index);
 33     const sum = try builder.add(try builder.add(try builder.floor(left), try builder.trunc(right)), try builder.atan2(left, right));
 34     try builder.store(sum, out, index);
 35     try builder.yield_(&.{});
 36     builder.setInsertionBlock(if_op.elseBlock().?);
 37     try builder.yield_(&.{});
 38     builder.setInsertionBlock(previous);
 39     try builder.return_();
 40 
 41     var program = try builder.finish();
 42     defer program.deinit();
 43     const text = try emitWgsl(testing.allocator, "accy_wgsl_test_add", program.kernelModule());
 44     defer testing.allocator.free(text);
 45 
 46     try testing.expect(std.mem.indexOf(u8, text, "@compute @workgroup_size") != null);
 47     try testing.expect(std.mem.indexOf(u8, text, "fn accy_wgsl_test_add") != null);
 48     try testing.expect(std.mem.indexOf(u8, text, "@group(0) @binding(0) var<storage, read_write> arg0: array<f32>;") != null);
 49     try testing.expect(std.mem.indexOf(u8, text, "@group(0) @binding(1) var<storage, read> arg1: array<f32>;") != null);
 50     try testing.expect(std.mem.indexOf(u8, text, "@group(0) @binding(2) var<storage, read> arg2: array<f32>;") != null);
 51     try testing.expect(std.mem.indexOf(u8, text, "@group(0) @binding(3) var<storage, read> arg3: array<i32>;") != null);
 52     try testing.expect(std.mem.indexOf(u8, text, "@builtin(global_invocation_id)") != null);
 53     try testing.expect(std.mem.indexOf(u8, text, "if (") != null);
 54     try testing.expect(std.mem.indexOf(u8, text, "floor(") != null);
 55     try testing.expect(std.mem.indexOf(u8, text, "trunc(") != null);
 56     try testing.expect(std.mem.indexOf(u8, text, "atan2(") != null);
 57     try testing.expect(std.mem.indexOf(u8, text, "arg0[") != null);
 58 }
 59 
 60 test "webgpu choir wgsl emitter lowers loop-carried reduction shape" {
 61     const testing = std.testing;
 62     const builder_mod = accy_root.kernel;
 63     const params = [_]builder_mod.Param{
 64         builder_mod.dynamicBuffer(.f32),
 65         builder_mod.dynamicBuffer(.f32),
 66     };
 67 
 68     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_reduce", &params);
 69     errdefer builder.deinit();
 70     const out = builder.argument(0);
 71     const input = builder.argument(1);
 72     const zero_index = try builder.constantIndex(0);
 73     const upper = try builder.constantIndex(4);
 74     const step = try builder.constantIndex(1);
 75     const initial = try builder.constantFloat(.f32, 0.0);
 76     var loop_scope = try builder.forScope(zero_index, upper, step, &.{initial}, &.{initial.valueType()});
 77     errdefer loop_scope.abort();
 78     const value = try builder.load(input, loop_scope.inductionVar());
 79     const next = try builder.add(loop_scope.iterArg(0).?, value);
 80     try loop_scope.leave(&.{next});
 81     const result = loop_scope.result(0).?;
 82     try builder.store(result, out, zero_index);
 83     try builder.return_();
 84 
 85     var program = try builder.finish();
 86     defer program.deinit();
 87     const text = try emitWgsl(testing.allocator, "accy_wgsl_test_reduce", program.kernelModule());
 88     defer testing.allocator.free(text);
 89 
 90     try testing.expect(std.mem.indexOf(u8, text, "for (var") != null);
 91     try testing.expect(std.mem.indexOf(u8, text, ": i32 =") != null);
 92     try testing.expect(std.mem.indexOf(u8, text, " = ") != null);
 93     try testing.expect(std.mem.indexOf(u8, text, "arg0[") != null);
 94 }
 95 
 96 test "webgpu choir wgsl emitter lowers workgroup memory and barriers" {
 97     const testing = std.testing;
 98     const builder_mod = accy_root.kernel;
 99     const params = [_]builder_mod.Param{
100         builder_mod.dynamicBuffer(.f32),
101     };
102 
103     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_shared", &params);
104     errdefer builder.deinit();
105     const out = builder.argument(0);
106     const shared = try builder.sharedBuffer(.f32, 16);
107     const zero_index = try builder.constantIndex(0);
108     const one = try builder.constantFloat(.f32, 1.0);
109     try builder.store(one, shared, zero_index);
110     try builder.barrier(.block);
111     const value = try builder.load(shared, zero_index);
112     try builder.store(value, out, zero_index);
113     try builder.return_();
114 
115     var program = try builder.finish();
116     defer program.deinit();
117     const text = try emitWgsl(testing.allocator, "accy_wgsl_test_shared", program.kernelModule());
118     defer testing.allocator.free(text);
119 
120     try testing.expect(std.mem.indexOf(u8, text, "var<workgroup> shared0: array<f32, 16>;") != null);
121     try testing.expect(std.mem.indexOf(u8, text, "workgroupBarrier();") != null);
122 }
123 
124 test "webgpu choir wgsl emitter lowers unsigned high multiply" {
125     const testing = std.testing;
126     const builder_mod = accy_root.kernel;
127     const params = [_]builder_mod.Param{
128         builder_mod.dynamicBuffer(.i32),
129         builder_mod.dynamicBuffer(.i32),
130         builder_mod.dynamicBuffer(.i32),
131     };
132 
133     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_umulhi", &params);
134     errdefer builder.deinit();
135     const out = builder.argument(0);
136     const lhs = builder.argument(1);
137     const rhs = builder.argument(2);
138     const index = try builder.globalId(.x);
139     const left = try builder.load(lhs, index);
140     const right = try builder.load(rhs, index);
141     const high = try builder.umulhi(left, right);
142     try builder.store(high, out, index);
143     try builder.return_();
144 
145     var program = try builder.finish();
146     defer program.deinit();
147     const text = try emitWgsl(testing.allocator, "accy_wgsl_test_umulhi", program.kernelModule());
148     defer testing.allocator.free(text);
149 
150     try testing.expect(std.mem.indexOf(u8, text, "fn choir_umulhi_u32") != null);
151     try testing.expect(std.mem.indexOf(u8, text, "bitcast<i32>(choir_umulhi_u32") != null);
152     try testing.expect(std.mem.indexOf(u8, text, "fn accy_wgsl_test_umulhi") != null);
153 }
154 
155 test "webgpu choir wgsl emitter lowers while loops with carried exits" {
156     const testing = std.testing;
157     const builder_mod = accy_root.kernel;
158     const params = [_]builder_mod.Param{
159         builder_mod.dynamicBuffer(.i32),
160         builder_mod.dynamicBuffer(.i32),
161     };
162 
163     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_while", &params);
164     errdefer builder.deinit();
165     const src = builder.argument(0);
166     const dst = builder.argument(1);
167     const index = try builder.globalId(.x);
168     const start = try builder.load(src, index);
169     const zero = try builder.constantInt(.i32, 0);
170     const one = try builder.constantInt(.i32, 1);
171 
172     var scope = try builder.whileScope(&.{ start, zero }, &.{ start.valueType(), zero.valueType() });
173     errdefer scope.abort();
174     const remaining = scope.beforeArg(0).?;
175     const total = scope.beforeArg(1).?;
176     const proceed = try builder.compare(.gt, remaining, zero);
177     try scope.condition(proceed, &.{ remaining, total });
178     const after_remaining = scope.afterArg(0).?;
179     const after_total = scope.afterArg(1).?;
180     const next_total = try builder.add(after_total, after_remaining);
181     const next_remaining = try builder.sub(after_remaining, one);
182     try scope.leave(&.{ next_remaining, next_total });
183 
184     try builder.store(scope.result(1).?, dst, index);
185     try builder.return_();
186 
187     var program = try builder.finish();
188     defer program.deinit();
189     const text = try emitWgsl(testing.allocator, "accy_wgsl_test_while", program.kernelModule());
190     defer testing.allocator.free(text);
191 
192     try testing.expect(std.mem.indexOf(u8, text, "loop {") != null);
193     try testing.expect(std.mem.indexOf(u8, text, "if (!(") != null);
194     try testing.expect(std.mem.indexOf(u8, text, "break;") != null);
195     try testing.expect(std.mem.count(u8, text, "var l") >= 4);
196 }
197 
198 test "webgpu choir wgsl emitter rejects barriers inside while loops" {
199     const testing = std.testing;
200     const builder_mod = accy_root.kernel;
201     const params = [_]builder_mod.Param{
202         builder_mod.dynamicBuffer(.i32),
203     };
204 
205     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_while_barrier", &params);
206     errdefer builder.deinit();
207     const dst = builder.argument(0);
208     const index = try builder.globalId(.x);
209     const start = try builder.cast(index, .i32);
210     const zero = try builder.constantInt(.i32, 0);
211     const one = try builder.constantInt(.i32, 1);
212 
213     var scope = try builder.whileScope(&.{start}, &.{start.valueType()});
214     errdefer scope.abort();
215     const remaining = scope.beforeArg(0).?;
216     const proceed = try builder.compare(.gt, remaining, zero);
217     try scope.condition(proceed, &.{remaining});
218     const after_remaining = scope.afterArg(0).?;
219     try builder.barrier(.block);
220     const next_remaining = try builder.sub(after_remaining, one);
221     try scope.leave(&.{next_remaining});
222 
223     try builder.store(scope.result(0).?, dst, index);
224     try builder.return_();
225 
226     var program = try builder.finish();
227     defer program.deinit();
228     try testing.expectError(
229         error.UnsupportedOperation,
230         emitWgsl(testing.allocator, "accy_wgsl_test_while_barrier", program.kernelModule()),
231     );
232 }
233 
234 test "webgpu choir wgsl emitter lowers value producing if arms" {
235     const testing = std.testing;
236     const builder_mod = accy_root.kernel;
237     const params = [_]builder_mod.Param{
238         builder_mod.dynamicBuffer(.f32),
239         builder_mod.dynamicBuffer(.f32),
240     };
241 
242     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_if_value", &params);
243     errdefer builder.deinit();
244     const dst = builder.argument(0);
245     const src = builder.argument(1);
246     const index = try builder.globalId(.x);
247     const value = try builder.load(src, index);
248     const zero = try builder.constantFloat(.f32, 0.0);
249     const positive = try builder.compare(.gt, value, zero);
250 
251     var if_op = try builder.if_(positive, &.{value.valueType()});
252     const previous = builder.insertionBlock();
253     builder.setInsertionBlock(if_op.thenBlock());
254     const doubled = try builder.add(value, value);
255     try builder.yield_(&.{doubled});
256     builder.setInsertionBlock(if_op.elseBlock().?);
257     try builder.yield_(&.{zero});
258     builder.setInsertionBlock(previous);
259 
260     try builder.store(if_op.result(0).?, dst, index);
261     try builder.return_();
262 
263     var program = try builder.finish();
264     defer program.deinit();
265     const text = try emitWgsl(testing.allocator, "accy_wgsl_test_if_value", program.kernelModule());
266     defer testing.allocator.free(text);
267 
268     try testing.expect(std.mem.indexOf(u8, text, "var l0: f32;") != null);
269     try testing.expect(std.mem.indexOf(u8, text, "} else {") != null);
270 }
271 
272 test "webgpu choir wgsl emitter rejects warp collectives" {
273     const testing = std.testing;
274     const builder_mod = accy_root.kernel;
275     const params = [_]builder_mod.Param{
276         builder_mod.dynamicBuffer(.f32),
277         builder_mod.dynamicBuffer(.f32),
278     };
279 
280     var builder = try builder_mod.Builder.init(testing.allocator, builder_mod.Builder.Limits.testing, "accy_wgsl_test_warp_reject", &params);
281     errdefer builder.deinit();
282     const out = builder.argument(0);
283     const input = builder.argument(1);
284     const lane = try builder.laneId();
285     const value = try builder.load(input, lane);
286     const reduced = try builder.warpReduce(.add, value);
287     try builder.store(reduced, out, lane);
288     try builder.return_();
289 
290     var program = try builder.finish();
291     defer program.deinit();
292     try testing.expectError(
293         error.UnsupportedOperation,
294         emitWgsl(testing.allocator, "accy_wgsl_test_warp_reject", program.kernelModule()),
295     );
296 }