lib/choir/src/backends/gpu/webgpu/wgsl.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const abi = @import("choir_abi");
3 const choir_pkg = @import("../../../root.zig");
4
5 const gpu = @import("../../../dialects/gpu/root.zig");
6
7 const ir = choir_pkg.ir;
8 const dialects = choir_pkg.dialects;
9
10 const Allocator = std.mem.Allocator;
11 const ArithDialect = dialects.ArithDialect;
12 const CmpPredicate = dialects.arith.CmpPredicate;
13 const BuiltinDialect = dialects.BuiltinDialect;
14 const FuncDialect = dialects.FuncDialect;
15 const GpuDialect = gpu.GpuDialect;
16 const MemrefDialect = dialects.MemrefDialect;
17 const ScfDialect = dialects.ScfDialect;
18
19 const EmitError = abi.Error || std.Io.Writer.Error;
20
21 const ScalarKind = dialects.arith.ScalarKind;
22 const scalar_kinds = dialects.arith.ScalarSet.init(&.{
23 .bool,
24 .index,
25 .i8,
26 .i16,
27 .i32,
28 .u32,
29 .i64,
30 .f16,
31 .f32,
32 .f64,
33 });
34
35 const YieldTarget = struct {
36 names: []const []const u8,
37 };
38
39 pub fn emitWgsl(
40 result_allocator: Allocator,
41 entry_name: []const u8,
42 module: *ir.Operation,
43 ) abi.Error![]u8 {
44 var emitter = Emitter.init(result_allocator, entry_name, module);
45 defer emitter.deinit();
46 return emitter.emit() catch |err| switch (err) {
47 error.WriteFailed => error.OutOfMemory,
48 else => |other| other,
49 };
50 }
51
52 const Emitter = struct {
53 allocator: Allocator,
54 entry_name: []const u8,
55 module: *ir.Operation,
56 declarations: std.Io.Writer.Allocating,
57 body: std.Io.Writer.Allocating,
58 values: std.AutoHashMapUnmanaged(*const ir.Value, []const u8) = .{},
59 names: std.ArrayListUnmanaged([]u8) = .empty,
60 next_value: u32 = 0,
61 next_loop: u32 = 0,
62 next_shared: u32 = 0,
63 indent: u32 = 1,
64 while_depth: u32 = 0,
65 uses_umulhi: bool = false,
66
67 fn init(allocator: Allocator, entry_name: []const u8, module: *ir.Operation) Emitter {
68 return .{
69 .allocator = allocator,
70 .entry_name = entry_name,
71 .module = module,
72 .declarations = std.Io.Writer.Allocating.init(allocator),
73 .body = std.Io.Writer.Allocating.init(allocator),
74 };
75 }
76
77 fn deinit(self: *Emitter) void {
78 self.values.deinit(self.allocator);
79 for (self.names.items) |name| self.allocator.free(name);
80 self.names.deinit(self.allocator);
81 self.declarations.deinit();
82 self.body.deinit();
83 }
84
85 fn emit(self: *Emitter) EmitError![]u8 {
86 const func = try self.findKernelFunction();
87 try self.emitParameterBindings(func);
88 try self.emitBlock(func.getEntryBlock(), null);
89
90 var out = std.Io.Writer.Allocating.init(self.allocator);
91 errdefer out.deinit();
92 try writeHeader(&out.writer);
93 if (self.uses_umulhi) try writeUmulhiHelper(&out.writer);
94 try out.writer.writeAll(self.declarations.written());
95 if (self.declarations.written().len != 0) try out.writer.writeByte('\n');
96 try self.emitFunctionHeader(&out.writer);
97 try out.writer.writeAll(self.body.written());
98 try out.writer.writeAll("}\n");
99 return out.toOwnedSlice() catch return error.OutOfMemory;
100 }
101
102 fn collectBlockWrites(
103 self: *Emitter,
104 block: *ir.Block,
105 written: *std.AutoHashMapUnmanaged(*const ir.Value, void),
106 ) EmitError!void {
107 var ops = block.getOperations();
108 while (ops.next()) |op| {
109 const name = op.name.name;
110 if (std.mem.eql(u8, name, MemrefDialect.StoreOp.operation_name)) {
111 const store = MemrefDialect.StoreOp{ .op = op };
112 written.put(self.allocator, store.getMemref(), {}) catch return error.OutOfMemory;
113 } else if (std.mem.eql(u8, name, MemrefDialect.AtomicRmwOp.operation_name)) {
114 const atomic = MemrefDialect.AtomicRmwOp{ .op = op };
115 written.put(self.allocator, atomic.getMemref(), {}) catch return error.OutOfMemory;
116 } else if (std.mem.eql(u8, name, MemrefDialect.AtomicCasOp.operation_name)) {
117 const atomic = MemrefDialect.AtomicCasOp{ .op = op };
118 written.put(self.allocator, atomic.getMemref(), {}) catch return error.OutOfMemory;
119 }
120 for (op.regions.items) |*region| {
121 var current = region.blocks.head;
122 while (current) |inner| : (current = inner.next) {
123 try self.collectBlockWrites(inner, written);
124 }
125 }
126 }
127 }
128
129 fn findKernelFunction(self: *Emitter) abi.Error!FuncDialect.FuncOp {
130 if (!std.mem.eql(u8, self.module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
131 return error.InvalidArtifact;
132 }
133 const block = self.module.getRegion(0).?.getEntryBlock() orelse return error.InvalidArtifact;
134 var ops = block.getOperations();
135 while (ops.next()) |op| {
136 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;
137 const func = FuncDialect.FuncOp{ .op = op };
138 if (!func.isKernel()) continue;
139 const name = func.getName() orelse return error.InvalidArtifact;
140 if (std.mem.eql(u8, name, self.entry_name)) return func;
141 }
142 return error.InvalidArtifact;
143 }
144
145 fn emitParameterBindings(self: *Emitter, func: FuncDialect.FuncOp) EmitError!void {
146 var written: std.AutoHashMapUnmanaged(*const ir.Value, void) = .empty;
147 defer written.deinit(self.allocator);
148 try self.collectBlockWrites(func.getEntryBlock(), &written);
149
150 const args = func.getArguments();
151 for (args, 0..) |arg, index| {
152 if (memrefInfo(arg.type)) |info| {
153 if (info.addr_space == .shared) return error.UnsupportedOperation;
154 const name = std.fmt.allocPrint(self.allocator, "arg{d}", .{index}) catch return error.OutOfMemory;
155 const bound = try self.rememberName(name);
156 const access = if (written.contains(arg)) "read_write" else "read";
157 try self.declarations.writer.print(
158 "@group(0) @binding({d}) var<storage, {s}> {s}: array<{s}>;\n",
159 .{ index, access, bound, try wgslStorageScalarType(info.element) },
160 );
161 try self.bind(arg, bound);
162 continue;
163 }
164 const kind = try scalarKind(arg.type);
165 const storage_type = try wgslStorageScalarType(kind);
166 const storage_name = std.fmt.allocPrint(self.allocator, "arg{d}_scalar", .{index}) catch return error.OutOfMemory;
167 defer self.allocator.free(storage_name);
168 try self.declarations.writer.print(
169 "@group(0) @binding({d}) var<storage, read> {s}: array<{s}>;\n",
170 .{ index, storage_name, storage_type },
171 );
172 const expr = std.fmt.allocPrint(self.allocator, "{s}[0]", .{storage_name}) catch return error.OutOfMemory;
173 try self.rememberAndBind(arg, expr);
174 }
175 }
176
177 fn emitFunctionHeader(self: *Emitter, writer: *std.Io.Writer) EmitError!void {
178 try writer.print("@compute @workgroup_size(choir_workgroup_size_x, choir_workgroup_size_y, choir_workgroup_size_z)\nfn {s}(\n", .{self.entry_name});
179 try writer.writeAll(" @builtin(global_invocation_id) choir_global_id: vec3<u32>,\n");
180 try writer.writeAll(" @builtin(local_invocation_id) choir_local_id: vec3<u32>,\n");
181 try writer.writeAll(" @builtin(workgroup_id) choir_workgroup_id: vec3<u32>,\n");
182 try writer.writeAll(" @builtin(num_workgroups) choir_num_workgroups: vec3<u32>,\n");
183 try writer.writeAll(" @builtin(local_invocation_index) choir_local_index: u32\n");
184 try writer.writeAll(") {\n");
185 }
186
187 fn emitBlock(self: *Emitter, block: *ir.Block, yield_target: ?YieldTarget) EmitError!void {
188 var ops = block.getOperations();
189 while (ops.next()) |op| {
190 try self.emitOperation(op, yield_target);
191 }
192 }
193
194 fn emitOperation(self: *Emitter, op: *ir.Operation, yield_target: ?YieldTarget) EmitError!void {
195 const name = op.name.name;
196 if (std.mem.eql(u8, name, FuncDialect.ReturnOp.operation_name)) {
197 try self.line("return;", .{});
198 } else if (std.mem.eql(u8, name, ScfDialect.YieldOp.operation_name)) {
199 try self.emitYield(ScfDialect.YieldOp{ .op = op }, yield_target);
200 } else if (std.mem.eql(u8, name, ScfDialect.IfOp.operation_name)) {
201 try self.emitIf(ScfDialect.IfOp{ .op = op });
202 } else if (std.mem.eql(u8, name, ScfDialect.ForOp.operation_name)) {
203 try self.emitFor(ScfDialect.ForOp{ .op = op });
204 } else if (std.mem.eql(u8, name, ScfDialect.WhileOp.operation_name)) {
205 try self.emitWhile(ScfDialect.WhileOp{ .op = op });
206 } else if (std.mem.eql(u8, name, GpuDialect.GlobalIdxOp.operation_name)) {
207 const wrapped = GpuDialect.GlobalIdxOp{ .op = op };
208 try self.emitGpuVectorRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_global_id");
209 } else if (std.mem.eql(u8, name, GpuDialect.ThreadIdxOp.operation_name)) {
210 const wrapped = GpuDialect.ThreadIdxOp{ .op = op };
211 try self.emitGpuVectorRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_local_id");
212 } else if (std.mem.eql(u8, name, GpuDialect.BlockIdxOp.operation_name)) {
213 const wrapped = GpuDialect.BlockIdxOp{ .op = op };
214 try self.emitGpuVectorRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_workgroup_id");
215 } else if (std.mem.eql(u8, name, GpuDialect.BlockDimOp.operation_name)) {
216 const wrapped = GpuDialect.BlockDimOp{ .op = op };
217 try self.emitBlockDimRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact);
218 } else if (std.mem.eql(u8, name, GpuDialect.GridDimOp.operation_name)) {
219 const wrapped = GpuDialect.GridDimOp{ .op = op };
220 try self.emitGpuVectorRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_num_workgroups");
221 } else if (std.mem.eql(u8, name, GpuDialect.LaneIdOp.operation_name)) {
222 return error.UnsupportedOperation;
223 } else if (std.mem.eql(u8, name, GpuDialect.WarpIdOp.operation_name)) {
224 return error.UnsupportedOperation;
225 } else if (std.mem.eql(u8, name, GpuDialect.BarrierOp.operation_name)) {
226 try self.emitBarrier(GpuDialect.BarrierOp{ .op = op });
227 } else if (std.mem.eql(u8, name, GpuDialect.WarpReduceOp.operation_name)) {
228 return error.UnsupportedOperation;
229 } else if (std.mem.eql(u8, name, GpuDialect.WarpScanOp.operation_name)) {
230 return error.UnsupportedOperation;
231 } else if (std.mem.eql(u8, name, MemrefDialect.AllocOp.operation_name)) {
232 try self.emitAlloc(MemrefDialect.AllocOp{ .op = op });
233 } else if (std.mem.eql(u8, name, MemrefDialect.LoadOp.operation_name)) {
234 try self.emitLoad(MemrefDialect.LoadOp{ .op = op });
235 } else if (std.mem.eql(u8, name, MemrefDialect.StoreOp.operation_name)) {
236 try self.emitStore(MemrefDialect.StoreOp{ .op = op });
237 } else if (std.mem.eql(u8, name, MemrefDialect.AtomicRmwOp.operation_name)) {
238 return error.UnsupportedOperation;
239 } else if (std.mem.eql(u8, name, MemrefDialect.AtomicCasOp.operation_name)) {
240 return error.UnsupportedOperation;
241 } else if (std.mem.eql(u8, name, ArithDialect.ConstantOp.operation_name)) {
242 try self.emitConstant(ArithDialect.ConstantOp{ .op = op });
243 } else if (std.mem.eql(u8, name, ArithDialect.AddOp.operation_name)) {
244 try self.emitBinary(op, "+");
245 } else if (std.mem.eql(u8, name, ArithDialect.SubOp.operation_name)) {
246 try self.emitBinary(op, "-");
247 } else if (std.mem.eql(u8, name, ArithDialect.MulOp.operation_name)) {
248 try self.emitBinary(op, "*");
249 } else if (std.mem.eql(u8, name, ArithDialect.UmulhiOp.operation_name)) {
250 try self.emitUmulhi(op);
251 } else if (std.mem.eql(u8, name, ArithDialect.DivOp.operation_name)) {
252 try self.emitBinary(op, "/");
253 } else if (std.mem.eql(u8, name, ArithDialect.MaxOp.operation_name)) {
254 try self.emitCall2(op, "max");
255 } else if (std.mem.eql(u8, name, ArithDialect.MinOp.operation_name)) {
256 try self.emitCall2(op, "min");
257 } else if (std.mem.eql(u8, name, ArithDialect.AndOp.operation_name)) {
258 try self.emitBitwiseOrLogical(op, "&", "&&");
259 } else if (std.mem.eql(u8, name, ArithDialect.OrOp.operation_name)) {
260 try self.emitBitwiseOrLogical(op, "|", "||");
261 } else if (std.mem.eql(u8, name, ArithDialect.XorOp.operation_name)) {
262 try self.emitBitwiseOrLogical(op, "^", "!=");
263 } else if (std.mem.eql(u8, name, ArithDialect.ShlOp.operation_name)) {
264 try self.emitShift(op, "<<");
265 } else if (std.mem.eql(u8, name, ArithDialect.ShrOp.operation_name)) {
266 try self.emitShift(op, ">>");
267 } else if (std.mem.eql(u8, name, ArithDialect.UshrOp.operation_name)) {
268 try self.emitUnsignedShiftRight(op);
269 } else if (std.mem.eql(u8, name, ArithDialect.NegOp.operation_name)) {
270 try self.emitUnary(op, "-");
271 } else if (std.mem.eql(u8, name, ArithDialect.NotOp.operation_name)) {
272 try self.emitNot(op);
273 } else if (std.mem.eql(u8, name, ArithDialect.AbsOp.operation_name)) {
274 try self.emitCall1(op, "abs");
275 } else if (std.mem.eql(u8, name, ArithDialect.SqrtOp.operation_name)) {
276 try self.emitCall1(op, "sqrt");
277 } else if (std.mem.eql(u8, name, ArithDialect.ExpOp.operation_name)) {
278 try self.emitCall1(op, "exp");
279 } else if (std.mem.eql(u8, name, ArithDialect.LogOp.operation_name)) {
280 try self.emitCall1(op, "log");
281 } else if (std.mem.eql(u8, name, ArithDialect.TanhOp.operation_name)) {
282 try self.emitCall1(op, "tanh");
283 } else if (std.mem.eql(u8, name, ArithDialect.SinOp.operation_name)) {
284 try self.emitCall1(op, "sin");
285 } else if (std.mem.eql(u8, name, ArithDialect.CosOp.operation_name)) {
286 try self.emitCall1(op, "cos");
287 } else if (std.mem.eql(u8, name, ArithDialect.TanOp.operation_name)) {
288 try self.emitCall1(op, "tan");
289 } else if (std.mem.eql(u8, name, ArithDialect.FloorOp.operation_name)) {
290 try self.emitCall1(op, "floor");
291 } else if (std.mem.eql(u8, name, ArithDialect.RoundOp.operation_name)) {
292 try self.emitCall1(op, "round");
293 } else if (std.mem.eql(u8, name, ArithDialect.TruncOp.operation_name)) {
294 try self.emitCall1(op, "trunc");
295 } else if (std.mem.eql(u8, name, ArithDialect.PowOp.operation_name)) {
296 try self.emitCall2(op, "pow");
297 } else if (std.mem.eql(u8, name, ArithDialect.Atan2Op.operation_name)) {
298 try self.emitCall2(op, "atan2");
299 } else if (std.mem.eql(u8, name, ArithDialect.FmaOp.operation_name)) {
300 try self.emitFma(ArithDialect.FmaOp{ .op = op });
301 } else if (std.mem.eql(u8, name, ArithDialect.CmpOp.operation_name)) {
302 try self.emitCompare(ArithDialect.CmpOp{ .op = op });
303 } else if (std.mem.eql(u8, name, ArithDialect.SelectOp.operation_name)) {
304 try self.emitSelect(ArithDialect.SelectOp{ .op = op });
305 } else if (std.mem.eql(u8, name, ArithDialect.CastOp.operation_name)) {
306 try self.emitCast(ArithDialect.CastOp{ .op = op });
307 } else if (std.mem.eql(u8, name, ArithDialect.BitcastOp.operation_name)) {
308 try self.emitBitcast(ArithDialect.BitcastOp{ .op = op });
309 } else {
310 return error.UnsupportedOperation;
311 }
312 }
313
314 fn emitGpuVectorRegister(self: *Emitter, op: *ir.Operation, dim: gpu.Dimension, builtin_name: []const u8) EmitError!void {
315 const result = op.getResult(0) orelse return error.InvalidArtifact;
316 const out = try self.freshValueName();
317 try self.line("let {s}: i32 = i32({s}.{s});", .{ out, builtin_name, dimName(dim) });
318 try self.bind(result, out);
319 }
320
321 fn emitBlockDimRegister(self: *Emitter, op: *ir.Operation, dim: gpu.Dimension) EmitError!void {
322 const result = op.getResult(0) orelse return error.InvalidArtifact;
323 const out = try self.freshValueName();
324 try self.line("let {s}: i32 = i32({s});", .{ out, workgroupSizeOverride(dim) });
325 try self.bind(result, out);
326 }
327
328 fn emitBarrier(self: *Emitter, op: GpuDialect.BarrierOp) EmitError!void {
329 if (self.while_depth != 0) return error.UnsupportedOperation;
330 const scope = op.getScope() orelse return error.InvalidArtifact;
331 switch (scope) {
332 .block => try self.line("workgroupBarrier();", .{}),
333 else => return error.UnsupportedOperation,
334 }
335 }
336
337 fn emitAlloc(self: *Emitter, op: MemrefDialect.AllocOp) EmitError!void {
338 if (op.getDynamicSize() != null) return error.UnsupportedOperation;
339 const result = op.getResult();
340 const info = memrefInfo(result.type) orelse return error.InvalidArtifact;
341 if (info.addr_space != .shared) return error.UnsupportedOperation;
342 const size = info.size orelse return error.InvalidArtifact;
343 if (size == 0) return error.InvalidArtifact;
344 const out = try self.freshSharedName();
345 try self.declarations.writer.print("var<workgroup> {s}: array<{s}, {d}>;\n", .{ out, try wgslStorageScalarType(info.element), size });
346 try self.bind(result, out);
347 }
348
349 fn emitLoad(self: *Emitter, op: MemrefDialect.LoadOp) EmitError!void {
350 const result = op.getResult();
351 const kind = try scalarKind(result.type);
352 const out = try self.freshValueName();
353 try self.line("let {s}: {s} = {s}[{s}];", .{
354 out,
355 try wgslScalarType(kind),
356 try self.require(op.getMemref()),
357 try self.require(op.getIndex()),
358 });
359 try self.bind(result, out);
360 }
361
362 fn emitStore(self: *Emitter, op: MemrefDialect.StoreOp) EmitError!void {
363 try self.line("{s}[{s}] = {s};", .{
364 try self.require(op.getMemref()),
365 try self.require(op.getIndex()),
366 try self.require(op.getValue()),
367 });
368 }
369
370 fn emitConstant(self: *Emitter, op: ArithDialect.ConstantOp) EmitError!void {
371 const result = op.getResult();
372 const kind = try scalarKind(result.type);
373 const out = try self.freshValueName();
374 switch (kind) {
375 .bool => {
376 const bool_attr = op.op.getAttrAs(ir.Attribute.BoolAttr, "value") orelse return error.InvalidArtifact;
377 try self.line("let {s}: bool = {s};", .{ out, if (bool_attr.getValue()) "true" else "false" });
378 },
379 .index => {
380 const value = op.getIntValue() orelse return error.InvalidArtifact;
381 if (value < 0) return error.UnsupportedOperation;
382 try self.line("let {s}: i32 = i32({d});", .{ out, value });
383 },
384 .u32 => {
385 const value = op.getIntValue() orelse return error.InvalidArtifact;
386 if (value < 0) return error.UnsupportedOperation;
387 try self.line("let {s}: u32 = {d}u;", .{ out, value });
388 },
389 .i32 => {
390 const value = op.getIntValue() orelse return error.InvalidArtifact;
391 try self.line("let {s}: i32 = i32({d});", .{ out, value });
392 },
393 .f32 => {
394 const value = op.getFloatValue() orelse return error.InvalidArtifact;
395 try self.line("let {s}: f32 = bitcast<f32>(0x{X:0>8}u);", .{ out, floatBits(value) });
396 },
397 .i8, .i16, .i64, .u8, .u16, .u64, .f16, .bf16, .f64 => return error.UnsupportedOperation,
398 }
399 try self.bind(result, out);
400 }
401
402 fn emitBinary(self: *Emitter, op: *ir.Operation, operator: []const u8) EmitError!void {
403 const result = op.getResult(0) orelse return error.InvalidArtifact;
404 const kind = try scalarKind(result.type);
405 const out = try self.freshValueName();
406 try self.line("let {s}: {s} = {s} {s} {s};", .{
407 out,
408 try wgslScalarType(kind),
409 try self.require(op.operands.items[0].value),
410 operator,
411 try self.require(op.operands.items[1].value),
412 });
413 try self.bind(result, out);
414 }
415
416 fn emitUmulhi(self: *Emitter, op: *ir.Operation) EmitError!void {
417 const result = op.getResult(0) orelse return error.InvalidArtifact;
418 const kind = try scalarKind(result.type);
419 const out = try self.freshValueName();
420 self.uses_umulhi = true;
421 switch (kind) {
422 .u32 => try self.line("let {s}: u32 = choir_umulhi_u32({s}, {s});", .{
423 out,
424 try self.require(op.operands.items[0].value),
425 try self.require(op.operands.items[1].value),
426 }),
427 .index, .i32 => try self.line("let {s}: {s} = bitcast<{s}>(choir_umulhi_u32(bitcast<u32>({s}), bitcast<u32>({s})));", .{
428 out,
429 try wgslScalarType(kind),
430 try wgslScalarType(kind),
431 try self.require(op.operands.items[0].value),
432 try self.require(op.operands.items[1].value),
433 }),
434 else => return error.UnsupportedOperation,
435 }
436 try self.bind(result, out);
437 }
438
439 fn emitBitwiseOrLogical(self: *Emitter, op: *ir.Operation, bitwise_operator: []const u8, bool_operator: []const u8) EmitError!void {
440 const result = op.getResult(0) orelse return error.InvalidArtifact;
441 const kind = try scalarKind(result.type);
442 const operator = if (kind == .bool) bool_operator else bitwise_operator;
443 const out = try self.freshValueName();
444 try self.line("let {s}: {s} = {s} {s} {s};", .{
445 out,
446 try wgslScalarType(kind),
447 try self.require(op.operands.items[0].value),
448 operator,
449 try self.require(op.operands.items[1].value),
450 });
451 try self.bind(result, out);
452 }
453
454 fn emitShift(self: *Emitter, op: *ir.Operation, operator: []const u8) EmitError!void {
455 const result = op.getResult(0) orelse return error.InvalidArtifact;
456 const kind = try scalarKind(result.type);
457 const out = try self.freshValueName();
458 try self.line("let {s}: {s} = {s} {s} u32({s});", .{
459 out,
460 try wgslScalarType(kind),
461 try self.require(op.operands.items[0].value),
462 operator,
463 try self.require(op.operands.items[1].value),
464 });
465 try self.bind(result, out);
466 }
467
468 fn emitUnsignedShiftRight(self: *Emitter, op: *ir.Operation) EmitError!void {
469 const result = op.getResult(0) orelse return error.InvalidArtifact;
470 const kind = try scalarKind(result.type);
471 const ty = try wgslScalarType(kind);
472 const out = try self.freshValueName();
473 if (kind == .u32) {
474 try self.line("let {s}: u32 = {s} >> u32({s});", .{
475 out,
476 try self.require(op.operands.items[0].value),
477 try self.require(op.operands.items[1].value),
478 });
479 } else {
480 _ = try wgslUnsignedScalarType(kind);
481 try self.line("let {s}: {s} = bitcast<{s}>(bitcast<u32>({s}) >> u32({s}));", .{
482 out,
483 ty,
484 ty,
485 try self.require(op.operands.items[0].value),
486 try self.require(op.operands.items[1].value),
487 });
488 }
489 try self.bind(result, out);
490 }
491
492 fn emitUnary(self: *Emitter, op: *ir.Operation, operator: []const u8) EmitError!void {
493 const result = op.getResult(0) orelse return error.InvalidArtifact;
494 const kind = try scalarKind(result.type);
495 if (kind == .u32) return error.UnsupportedOperation;
496 const out = try self.freshValueName();
497 try self.line("let {s}: {s} = {s}{s};", .{
498 out,
499 try wgslScalarType(kind),
500 operator,
501 try self.require(op.operands.items[0].value),
502 });
503 try self.bind(result, out);
504 }
505
506 fn emitNot(self: *Emitter, op: *ir.Operation) EmitError!void {
507 const result = op.getResult(0) orelse return error.InvalidArtifact;
508 const kind = try scalarKind(result.type);
509 const operator = switch (kind) {
510 .bool => "!",
511 .index, .i32, .u32 => "~",
512 else => return error.UnsupportedOperation,
513 };
514 const out = try self.freshValueName();
515 try self.line("let {s}: {s} = {s}{s};", .{
516 out,
517 try wgslScalarType(kind),
518 operator,
519 try self.require(op.operands.items[0].value),
520 });
521 try self.bind(result, out);
522 }
523
524 fn emitCall1(self: *Emitter, op: *ir.Operation, function_name: []const u8) EmitError!void {
525 const result = op.getResult(0) orelse return error.InvalidArtifact;
526 const kind = try scalarKind(result.type);
527 const out = try self.freshValueName();
528 try self.line("let {s}: {s} = {s}({s});", .{
529 out,
530 try wgslScalarType(kind),
531 function_name,
532 try self.require(op.operands.items[0].value),
533 });
534 try self.bind(result, out);
535 }
536
537 fn emitCall2(self: *Emitter, op: *ir.Operation, function_name: []const u8) EmitError!void {
538 const result = op.getResult(0) orelse return error.InvalidArtifact;
539 const kind = try scalarKind(result.type);
540 const out = try self.freshValueName();
541 try self.line("let {s}: {s} = {s}({s}, {s});", .{
542 out,
543 try wgslScalarType(kind),
544 function_name,
545 try self.require(op.operands.items[0].value),
546 try self.require(op.operands.items[1].value),
547 });
548 try self.bind(result, out);
549 }
550
551 fn emitFma(self: *Emitter, op: ArithDialect.FmaOp) EmitError!void {
552 const result = op.getResult();
553 const out = try self.freshValueName();
554 try self.line("let {s}: {s} = fma({s}, {s}, {s});", .{
555 out,
556 try wgslScalarType(try scalarKind(result.type)),
557 try self.require(op.getA()),
558 try self.require(op.getB()),
559 try self.require(op.getC()),
560 });
561 try self.bind(result, out);
562 }
563
564 fn emitCompare(self: *Emitter, op: ArithDialect.CmpOp) EmitError!void {
565 const out = try self.freshValueName();
566 try self.line("let {s}: bool = {s} {s} {s};", .{
567 out,
568 try self.require(op.op.operands.items[0].value),
569 comparisonOperator(op.getPredicate() orelse return error.InvalidArtifact),
570 try self.require(op.op.operands.items[1].value),
571 });
572 try self.bind(op.getResult(), out);
573 }
574
575 fn emitSelect(self: *Emitter, op: ArithDialect.SelectOp) EmitError!void {
576 const result = op.getResult();
577 const out = try self.freshValueName();
578 try self.line("let {s}: {s} = select({s}, {s}, {s});", .{
579 out,
580 try wgslScalarType(try scalarKind(result.type)),
581 try self.require(op.getFalseValue()),
582 try self.require(op.getTrueValue()),
583 try self.require(op.getCondition()),
584 });
585 try self.bind(result, out);
586 }
587
588 fn emitCast(self: *Emitter, op: ArithDialect.CastOp) EmitError!void {
589 const result = op.getResult();
590 const out = try self.freshValueName();
591 const ty = try wgslScalarType(try scalarKind(result.type));
592 try self.line("let {s}: {s} = {s}({s});", .{
593 out,
594 ty,
595 ty,
596 try self.require(op.getInput()),
597 });
598 try self.bind(result, out);
599 }
600
601 fn emitBitcast(self: *Emitter, op: ArithDialect.BitcastOp) EmitError!void {
602 const result = op.getResult();
603 const out = try self.freshValueName();
604 const ty = try wgslScalarType(try scalarKind(result.type));
605 try self.line("let {s}: {s} = bitcast<{s}>({s});", .{
606 out,
607 ty,
608 ty,
609 try self.require(op.getInput()),
610 });
611 try self.bind(result, out);
612 }
613
614 fn emitIf(self: *Emitter, op: ScfDialect.IfOp) EmitError!void {
615 const result_count = op.getNumResults();
616 if (result_count == 0) {
617 try self.line("if ({s}) {{", .{try self.require(op.getCondition())});
618 self.indent += 1;
619 try self.emitBlock(op.getThenBlock(), null);
620 self.indent -= 1;
621 if (op.getElseBlock()) |else_block| {
622 try self.line("}} else {{", .{});
623 self.indent += 1;
624 try self.emitBlock(else_block, null);
625 self.indent -= 1;
626 }
627 try self.line("}}", .{});
628 return;
629 }
630
631 const else_block = op.getElseBlock() orelse return error.InvalidArtifact;
632 const result_names = self.allocator.alloc([]const u8, result_count) catch return error.OutOfMemory;
633 defer self.allocator.free(result_names);
634 for (result_names, 0..) |*name, index| {
635 const result = op.op.getResult(index) orelse return error.InvalidArtifact;
636 name.* = try self.freshLoopName();
637 try self.line("var {s}: {s};", .{ name.*, try wgslScalarType(try scalarKind(result.type)) });
638 }
639
640 try self.line("if ({s}) {{", .{try self.require(op.getCondition())});
641 self.indent += 1;
642 try self.emitBlock(op.getThenBlock(), .{ .names = result_names });
643 self.indent -= 1;
644 try self.line("}} else {{", .{});
645 self.indent += 1;
646 try self.emitBlock(else_block, .{ .names = result_names });
647 self.indent -= 1;
648 try self.line("}}", .{});
649
650 for (result_names, 0..) |name, index| {
651 try self.bind(op.op.getResult(index).?, name);
652 }
653 }
654
655 fn emitWhile(self: *Emitter, op: ScfDialect.WhileOp) EmitError!void {
656 const before = op.getBeforeBlock();
657 const after = op.getAfterBlock();
658 const carry_count = op.op.operands.items.len;
659 if (op.op.results.items.len != carry_count) return error.InvalidArtifact;
660 if (before.arguments.items.len != carry_count) return error.InvalidArtifact;
661 if (after.arguments.items.len != carry_count) return error.InvalidArtifact;
662
663 const carry_names = self.allocator.alloc([]const u8, carry_count) catch return error.OutOfMemory;
664 defer self.allocator.free(carry_names);
665 const exit_names = self.allocator.alloc([]const u8, carry_count) catch return error.OutOfMemory;
666 defer self.allocator.free(exit_names);
667
668 for (0..carry_count) |index| {
669 const operand = op.op.operands.items[index].value;
670 const carry_type = try wgslScalarType(try scalarKind(operand.type));
671 carry_names[index] = try self.freshLoopName();
672 try self.line("var {s}: {s} = {s};", .{ carry_names[index], carry_type, try self.require(operand) });
673 exit_names[index] = try self.freshLoopName();
674 try self.line("var {s}: {s};", .{ exit_names[index], carry_type });
675 try self.bind(before.arguments.items[index], carry_names[index]);
676 }
677
678 self.while_depth += 1;
679 defer self.while_depth -= 1;
680 try self.line("loop {{", .{});
681 self.indent += 1;
682
683 var before_ops = before.getOperations();
684 var saw_condition = false;
685 while (before_ops.next()) |before_op| {
686 if (std.mem.eql(u8, before_op.name.name, ScfDialect.ConditionOp.operation_name)) {
687 const condition = ScfDialect.ConditionOp{ .op = before_op };
688 const args = condition.getArgs();
689 if (args.len != carry_count) return error.InvalidArtifact;
690 try self.line("if (!({s})) {{", .{try self.require(condition.getCondition())});
691 self.indent += 1;
692 for (args, exit_names) |arg, exit_name| {
693 try self.line("{s} = {s};", .{ exit_name, try self.require(arg) });
694 }
695 try self.line("break;", .{});
696 self.indent -= 1;
697 try self.line("}}", .{});
698 for (args, 0..) |arg, index| {
699 try self.bind(after.arguments.items[index], try self.require(arg));
700 }
701 saw_condition = true;
702 break;
703 }
704 try self.emitOperation(before_op, null);
705 }
706 if (!saw_condition) return error.InvalidArtifact;
707
708 var after_ops = after.getOperations();
709 var saw_yield = false;
710 while (after_ops.next()) |after_op| {
711 if (std.mem.eql(u8, after_op.name.name, ScfDialect.YieldOp.operation_name)) {
712 const yield = ScfDialect.YieldOp{ .op = after_op };
713 const yielded = yield.getOperands();
714 if (yielded.len != carry_count) return error.InvalidArtifact;
715 for (yielded, carry_names) |yield_value, carry_name| {
716 try self.line("{s} = {s};", .{ carry_name, try self.require(yield_value) });
717 }
718 saw_yield = true;
719 break;
720 }
721 try self.emitOperation(after_op, null);
722 }
723 if (!saw_yield) return error.InvalidArtifact;
724
725 self.indent -= 1;
726 try self.line("}}", .{});
727
728 for (op.op.results.items, exit_names) |*result, exit_name| {
729 try self.bind(result, exit_name);
730 }
731 }
732
733 fn emitFor(self: *Emitter, op: ScfDialect.ForOp) EmitError!void {
734 const init_args = op.getInitArgs();
735 if (op.op.results.items.len != init_args.len) return error.InvalidArtifact;
736
737 const accumulator_names = self.allocator.alloc([]const u8, init_args.len) catch return error.OutOfMemory;
738 defer self.allocator.free(accumulator_names);
739
740 const iter_args = op.getIterArgs();
741 for (init_args, 0..) |initial, index| {
742 const name = try self.freshLoopName();
743 accumulator_names[index] = name;
744 try self.line("var {s}: {s} = {s};", .{
745 name,
746 try wgslScalarType(try scalarKind(initial.type)),
747 try self.require(initial),
748 });
749 try self.bind(iter_args[index], name);
750 }
751
752 const iv_name = try self.freshLoopName();
753 try self.bind(op.getInductionVar(), iv_name);
754 const lower = try self.require(op.getLowerBound());
755 const upper = try self.require(op.getUpperBound());
756 const step = try self.require(op.getStep());
757 try self.line("for (var {s}: i32 = {s}; {s} < {s}; {s} = {s} + {s}) {{", .{ iv_name, lower, iv_name, upper, iv_name, iv_name, step });
758 self.indent += 1;
759 try self.emitBlock(op.getBodyBlock(), .{ .names = accumulator_names });
760 self.indent -= 1;
761 try self.line("}}", .{});
762
763 for (op.op.results.items, 0..) |*result, index| {
764 try self.bind(result, accumulator_names[index]);
765 }
766 }
767
768 fn emitYield(self: *Emitter, op: ScfDialect.YieldOp, yield_target: ?YieldTarget) EmitError!void {
769 const operands = op.getOperands();
770 const target_names = if (yield_target) |target_binding| target_binding.names else {
771 if (operands.len != 0) return error.UnsupportedOperation;
772 return;
773 };
774 if (operands.len != target_names.len) return error.InvalidArtifact;
775 for (operands, target_names) |operand, target_name| {
776 try self.line("{s} = {s};", .{ target_name, try self.require(operand) });
777 }
778 }
779
780 fn freshValueName(self: *Emitter) abi.Error![]const u8 {
781 const index = self.next_value;
782 self.next_value += 1;
783 const name = std.fmt.allocPrint(self.allocator, "v{d}", .{index}) catch return error.OutOfMemory;
784 return try self.rememberName(name);
785 }
786
787 fn freshLoopName(self: *Emitter) abi.Error![]const u8 {
788 const index = self.next_loop;
789 self.next_loop += 1;
790 const name = std.fmt.allocPrint(self.allocator, "l{d}", .{index}) catch return error.OutOfMemory;
791 return try self.rememberName(name);
792 }
793
794 fn freshSharedName(self: *Emitter) abi.Error![]const u8 {
795 const index = self.next_shared;
796 self.next_shared += 1;
797 const name = std.fmt.allocPrint(self.allocator, "shared{d}", .{index}) catch return error.OutOfMemory;
798 return try self.rememberName(name);
799 }
800
801 fn rememberAndBind(self: *Emitter, value: *const ir.Value, name: []u8) abi.Error!void {
802 const owned = try self.rememberName(name);
803 try self.bind(value, owned);
804 }
805
806 fn rememberName(self: *Emitter, name: []u8) abi.Error![]const u8 {
807 self.names.append(self.allocator, name) catch {
808 self.allocator.free(name);
809 return error.OutOfMemory;
810 };
811 return name;
812 }
813
814 fn bind(self: *Emitter, value: *const ir.Value, name: []const u8) abi.Error!void {
815 self.values.put(self.allocator, value, name) catch return error.OutOfMemory;
816 }
817
818 fn require(self: *Emitter, value: *const ir.Value) abi.Error![]const u8 {
819 return self.values.get(value) orelse error.InvalidArtifact;
820 }
821
822 fn line(self: *Emitter, comptime fmt: []const u8, args: anytype) std.Io.Writer.Error!void {
823 for (0..self.indent) |_| try self.body.writer.writeAll(" ");
824 try self.body.writer.print(fmt, args);
825 try self.body.writer.writeByte('\n');
826 }
827 };
828
829 fn writeHeader(writer: *std.Io.Writer) std.Io.Writer.Error!void {
830 try writer.writeAll("override choir_workgroup_size_x: u32 = 1u;\n");
831 try writer.writeAll("override choir_workgroup_size_y: u32 = 1u;\n");
832 try writer.writeAll("override choir_workgroup_size_z: u32 = 1u;\n\n");
833 }
834
835 fn writeUmulhiHelper(writer: *std.Io.Writer) std.Io.Writer.Error!void {
836 try writer.writeAll(
837 \\fn choir_umulhi_u32(lhs: u32, rhs: u32) -> u32 {
838 \\ let lhs_lo: u32 = lhs & 0xffffu;
839 \\ let lhs_hi: u32 = lhs >> 16u;
840 \\ let rhs_lo: u32 = rhs & 0xffffu;
841 \\ let rhs_hi: u32 = rhs >> 16u;
842 \\ let low: u32 = lhs_lo * rhs_lo;
843 \\ let mid0: u32 = lhs_lo * rhs_hi;
844 \\ let mid1: u32 = lhs_hi * rhs_lo;
845 \\ let high: u32 = lhs_hi * rhs_hi;
846 \\ let carry: u32 = ((low >> 16u) + (mid0 & 0xffffu) + (mid1 & 0xffffu)) >> 16u;
847 \\ return high + (mid0 >> 16u) + (mid1 >> 16u) + carry;
848 \\}
849 \\
850 \\
851 );
852 }
853
854 fn scalarKind(typ: ir.Type) abi.Error!ScalarKind {
855 return scalar_kinds.kindFromType(typ) orelse error.UnsupportedOperation;
856 }
857
858 const MemrefInfo = struct {
859 size: ?u64,
860 element: ScalarKind,
861 addr_space: dialects.AddressSpace,
862 };
863
864 fn memrefInfo(typ: ir.Type) ?MemrefInfo {
865 const name = typ.getDialectTypeName() orelse return null;
866 if (!std.mem.eql(u8, name, MemrefDialect.name)) return null;
867 const params = MemrefDialect.parseMemrefParams(typ.getDialectParamKey() orelse return null) orelse return null;
868 return .{
869 .size = params.size,
870 .element = scalar_kinds.kindFromTypeName(params.element_type_name) orelse return null,
871 .addr_space = params.addr_space,
872 };
873 }
874
875 fn wgslScalarType(kind: ScalarKind) abi.Error![]const u8 {
876 return switch (kind) {
877 .bool => "bool",
878 .index => "i32",
879 .i32 => "i32",
880 .u32 => "u32",
881 .f32 => "f32",
882 .i8, .i16, .i64, .u8, .u16, .u64, .f16, .bf16, .f64 => error.UnsupportedOperation,
883 };
884 }
885
886 fn wgslStorageScalarType(kind: ScalarKind) abi.Error![]const u8 {
887 return switch (kind) {
888 .index => "i32",
889 .i32 => "i32",
890 .u32 => "u32",
891 .f32 => "f32",
892 .bool, .i8, .i16, .i64, .u8, .u16, .u64, .f16, .bf16, .f64 => error.UnsupportedOperation,
893 };
894 }
895
896 fn wgslUnsignedScalarType(kind: ScalarKind) abi.Error![]const u8 {
897 return switch (kind) {
898 .index, .i32, .u32 => "u32",
899 else => error.UnsupportedOperation,
900 };
901 }
902
903 fn dimName(dim: gpu.Dimension) []const u8 {
904 return switch (dim) {
905 .x => "x",
906 .y => "y",
907 .z => "z",
908 };
909 }
910
911 fn workgroupSizeOverride(dim: gpu.Dimension) []const u8 {
912 return switch (dim) {
913 .x => "choir_workgroup_size_x",
914 .y => "choir_workgroup_size_y",
915 .z => "choir_workgroup_size_z",
916 };
917 }
918
919 fn comparisonOperator(pred: CmpPredicate) []const u8 {
920 return switch (pred) {
921 .eq => "==",
922 .ne => "!=",
923 .lt, .slt, .ult => "<",
924 .le, .sle, .ule => "<=",
925 .gt, .sgt, .ugt => ">",
926 .ge, .sge, .uge => ">=",
927 };
928 }
929
930 fn floatBits(value: f64) u32 {
931 const narrowed: f32 = @floatCast(value);
932 return @bitCast(narrowed);
933 }
934
935 test "webgpu scalar support mask follows Choir scalar spellings" {
936 inline for (std.meta.tags(ScalarKind)) |kind| {
937 const expected: ?ScalarKind = if (scalar_kinds.contains(kind)) kind else null;
938 try std.testing.expectEqual(
939 expected,
940 scalar_kinds.kindFromTypeName(dialects.arith.scalarTypeName(kind)),
941 );
942 }
943 try std.testing.expectEqual(@as(?ScalarKind, null), scalar_kinds.kindFromTypeName("arith.unknown"));
944 }