lib/choir/src/backends/gpu/metal/msl.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 const calls = @import("../calls.zig");
7
8 const ir = choir_pkg.ir;
9 const dialects = choir_pkg.dialects;
10
11 const Allocator = std.mem.Allocator;
12 const ArithDialect = dialects.ArithDialect;
13 const AtomicRmwKind = dialects.AtomicRmwKind;
14 const CmpPredicate = dialects.arith.CmpPredicate;
15 const BuiltinDialect = dialects.BuiltinDialect;
16 const FuncDialect = dialects.FuncDialect;
17 const GpuDialect = gpu.GpuDialect;
18 const MemrefDialect = dialects.MemrefDialect;
19 const ScfDialect = dialects.ScfDialect;
20 const Stage = gpu.Stage;
21 const max_locations = gpu.stage.max_locations;
22
23 const EmitError = RenderError || std.Io.Writer.Error;
24
25 /// The render path's failures: the ABI's, and the stage memory `calls.Plan` refuses by name.
26 pub const RenderError = abi.Error || gpu.stage.MemoryError || calls.SignatureError;
27
28 const ScalarKind = dialects.arith.ScalarKind;
29 const scalar_kinds = dialects.arith.ScalarSet.init(&.{
30 .bool,
31 .index,
32 .i8,
33 .i16,
34 .i32,
35 .u32,
36 .i64,
37 .f16,
38 .f32,
39 .f64,
40 });
41
42 const YieldTarget = struct {
43 names: []const []const u8,
44 };
45
46 pub const max_texture_bindings = abi.metal.max_texture_bindings;
47 pub const max_uniform_bindings = abi.metal.max_uniform_bindings;
48 pub const uniform_buffer_base = abi.metal.uniform_buffer_base;
49 pub const push_buffer = abi.metal.push_buffer;
50
51 pub const RenderEntries = struct {
52 vertex: []const u8,
53 fragment: []const u8,
54 };
55
56 /// Emits one source holding `entries.vertex` as a vertex function and
57 /// `entries.fragment` as a fragment function.
58 ///
59 /// Interface slot n is `[[attribute(n)]]` for a vertex input, `[[user(locnN)]]`
60 /// between the stages, and `[[color(n)]]` for a fragment output. Metal binds
61 /// resources by one flat index per stage, so a sampled texture takes group 0,
62 /// and its binding is the index of both its texture and its sampler.
63 ///
64 /// A stage's clip space puts y = -1 at the target's top edge, as the gpu
65 /// contract and Vulkan do, and Metal's puts y = +1 there. So the vertex
66 /// function negates the y it writes to `[[position]]`, which is exact, and the
67 /// host's viewport is Metal's own, origin at the top left.
68 ///
69 /// Block reads index 32-bit words, as the SPIR-V blocks do. The push-constant
70 /// words are `constant uint*` at `[[buffer(push_buffer)]]`, which a host sets
71 /// with `setVertexBytes` and `setFragmentBytes`. Uniform binding `b` of group
72 /// 0 is `constant uint4*` at `[[buffer(uniform_buffer_base + b)]]`, rows of
73 /// 16 bytes as std140 lays them. A read reinterprets its word with `as_type`.
74 pub fn emitMslRender(
75 result_allocator: Allocator,
76 module: *ir.Operation,
77 entries: RenderEntries,
78 ) RenderError![]u8 {
79 var out = std.Io.Writer.Allocating.init(result_allocator);
80 errdefer out.deinit();
81 const stages = [_]struct { name: []const u8, stage: Stage }{
82 .{ .name = entries.vertex, .stage = .vertex },
83 .{ .name = entries.fragment, .stage = .fragment },
84 };
85 writeHeader(&out.writer) catch return error.OutOfMemory;
86 writeHelpers(&out.writer, result_allocator, module) catch |err| switch (err) {
87 error.WriteFailed => return error.OutOfMemory,
88 else => |other| return other,
89 };
90 for (stages, 0..) |entry, index| {
91 if (index > 0) out.writer.writeByte('\n') catch return error.OutOfMemory;
92 var emitter = Emitter.init(result_allocator, entry.name, module);
93 defer emitter.deinit();
94 emitter.emitStage(&out.writer, entry.stage) catch |err| switch (err) {
95 error.WriteFailed => return error.OutOfMemory,
96 else => |other| return other,
97 };
98 }
99 return out.toOwnedSlice() catch return error.OutOfMemory;
100 }
101
102 fn writeHelpers(writer: *std.Io.Writer, allocator: Allocator, module: *ir.Operation) EmitError!void {
103 var plan = try calls.Plan.init(allocator, module, null);
104 defer plan.deinit();
105 for (plan.helpers.items) |func| {
106 const name = (FuncDialect.FuncOp{ .op = func }).getName() orelse return error.InvalidArtifact;
107 var emitter = Emitter.init(allocator, name, module);
108 defer emitter.deinit();
109 try emitter.emitHelperSource(writer, FuncDialect.FuncOp{ .op = func });
110 try writer.writeByte('\n');
111 }
112 }
113
114 const Direction = enum { input, output };
115
116 const Slot = struct {
117 kind: ScalarKind,
118 width: u32,
119 };
120
121 /// What a stage function reads and writes, gathered while its body is
122 /// emitted and written into its structs and signature afterwards.
123 const StageInterface = struct {
124 inputs: [max_locations]?Slot = @splat(null),
125 outputs: [max_locations]?Slot = @splat(null),
126 frag_coord: bool = false,
127 vertex_id: bool = false,
128 instance_id: bool = false,
129 front_facing: bool = false,
130 textures: std.StaticBitSet(max_texture_bindings) = .empty,
131 push: bool = false,
132 uniforms: std.StaticBitSet(max_uniform_bindings) = .empty,
133
134 fn record(
135 self: *StageInterface,
136 direction: Direction,
137 location: u32,
138 slot: Slot,
139 ) abi.Error!void {
140 if (location >= max_locations) return error.InvalidArtifact;
141 const slots = switch (direction) {
142 .input => &self.inputs,
143 .output => &self.outputs,
144 };
145 if (slots[location]) |existing| {
146 if (existing.kind != slot.kind) return error.InvalidArtifact;
147 if (existing.width != slot.width) return error.InvalidArtifact;
148 }
149 slots[location] = slot;
150 }
151
152 fn count(slots: []const ?Slot) usize {
153 var total: usize = 0;
154 for (slots) |slot| total += @intFromBool(slot != null);
155 return total;
156 }
157
158 fn hasInput(self: *const StageInterface) bool {
159 return count(&self.inputs) > 0 or self.frag_coord;
160 }
161 };
162
163 /// The struct member prefix of a slot: attributes feed the vertex stage,
164 /// locations join the stages, and colors leave the fragment stage.
165 fn memberPrefix(stage: Stage, direction: Direction) []const u8 {
166 return switch (stage) {
167 .vertex => if (direction == .input) "a" else "l",
168 .fragment => if (direction == .input) "l" else "c",
169 };
170 }
171
172 fn widthSuffix(width: u32) []const u8 {
173 return switch (width) {
174 1 => "",
175 2 => "2",
176 3 => "3",
177 4 => "4",
178 else => unreachable,
179 };
180 }
181
182 const components = [_][]const u8{ "x", "y", "z", "w" };
183
184 fn interfaceKind(typ: ir.Type) abi.Error!ScalarKind {
185 const kind = try scalarKind(typ);
186 return switch (kind) {
187 .f32, .i32, .u32 => kind,
188 else => error.UnsupportedOperation,
189 };
190 }
191
192 /// The kind every result of a block read shares.
193 fn blockReadKind(op: *ir.Operation) abi.Error!ScalarKind {
194 const kind = try interfaceKind(op.getResult(0).?.type);
195 for (0..op.getNumResults()) |index| {
196 if (try interfaceKind(op.getResult(index).?.type) != kind) return error.InvalidArtifact;
197 }
198 return kind;
199 }
200
201 pub fn emitMsl(
202 result_allocator: Allocator,
203 entry_name: []const u8,
204 module: *ir.Operation,
205 ) abi.Error![]u8 {
206 var emitter = Emitter.init(result_allocator, entry_name, module);
207 defer emitter.deinit();
208 return emitter.emit() catch |err| switch (err) {
209 error.WriteFailed => error.OutOfMemory,
210 error.UnsupportedStageMemory, error.UnsupportedHelperSignature => error.UnsupportedOperation,
211 else => |other| other,
212 };
213 }
214
215 const Emitter = struct {
216 allocator: Allocator,
217 entry_name: []const u8,
218 module: *ir.Operation,
219 body: std.Io.Writer.Allocating,
220 values: std.AutoHashMapUnmanaged(*const ir.Value, []const u8) = .{},
221 names: std.ArrayListUnmanaged([]u8) = .empty,
222 next_value: u32 = 0,
223 next_loop: u32 = 0,
224 next_shared: u32 = 0,
225 indent: u32 = 1,
226 stage: ?Stage = null,
227 interface: StageInterface = .{},
228 return_statement: []const u8 = "return;",
229
230 fn init(allocator: Allocator, entry_name: []const u8, module: *ir.Operation) Emitter {
231 return .{
232 .allocator = allocator,
233 .entry_name = entry_name,
234 .module = module,
235 .body = std.Io.Writer.Allocating.init(allocator),
236 };
237 }
238
239 fn deinit(self: *Emitter) void {
240 self.values.deinit(self.allocator);
241 for (self.names.items) |name| self.allocator.free(name);
242 self.names.deinit(self.allocator);
243 self.body.deinit();
244 }
245
246 fn emit(self: *Emitter) EmitError![]u8 {
247 const func = try self.findKernelFunction();
248 try self.emitParameterBindings(func);
249 try self.emitBlock(func.getEntryBlock(), null);
250
251 var out = std.Io.Writer.Allocating.init(self.allocator);
252 errdefer out.deinit();
253 try writeHeader(&out.writer);
254 try writeHelpers(&out.writer, self.allocator, self.module);
255 try self.emitFunctionHeader(&out.writer, func);
256 try out.writer.writeAll(self.body.written());
257 try out.writer.writeAll("}\n");
258 return out.toOwnedSlice() catch return error.OutOfMemory;
259 }
260
261 fn emitStage(self: *Emitter, writer: *std.Io.Writer, stage: Stage) EmitError!void {
262 const func = try self.findStageFunction(stage);
263 if (func.getNumArguments() != 0) return error.InvalidArtifact;
264 self.stage = stage;
265 self.return_statement = "return choir_out;";
266 try self.line("{s}_out choir_out = {{}};", .{self.entry_name});
267 try self.emitBlock(func.getEntryBlock(), null);
268 if (stage == .fragment and StageInterface.count(&self.interface.outputs) == 0) {
269 return error.InvalidArtifact;
270 }
271 try self.writeStageStructs(writer, stage);
272 try self.writeStageSignature(writer, stage);
273 try writer.writeAll(self.body.written());
274 try writer.writeAll("}\n");
275 }
276
277 fn emitHelperSource(self: *Emitter, writer: *std.Io.Writer, func: FuncDialect.FuncOp) EmitError!void {
278 const results = func.getResultTypes();
279 if (results.len > 1) return error.UnsupportedOperation;
280 const return_type = if (results.len == 0) "void" else try mslScalarType(try scalarKind(results[0]));
281 try writer.print("{s} {s}(", .{ return_type, self.entry_name });
282 for (func.getArguments(), 0..) |arg, index| {
283 if (index != 0) try writer.writeAll(", ");
284 const name = try self.freshValueName();
285 try writer.print("{s} {s}", .{ try mslScalarType(try scalarKind(arg.type)), name });
286 try self.bind(arg, name);
287 }
288 try writer.writeAll(") {\n");
289 try self.emitBlock(func.getEntryBlock(), null);
290 try writer.writeAll(self.body.written());
291 try writer.writeAll("}\n");
292 }
293
294 fn findStageFunction(self: *Emitter, stage: Stage) abi.Error!FuncDialect.FuncOp {
295 if (!std.mem.eql(u8, self.module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
296 return error.InvalidArtifact;
297 }
298 const region = self.module.getRegion(0).?;
299 const block = region.getEntryBlock() orelse return error.InvalidArtifact;
300 var ops = block.getOperations();
301 while (ops.next()) |op| {
302 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;
303 const func = FuncDialect.FuncOp{ .op = op };
304 const name = func.getName() orelse continue;
305 if (!std.mem.eql(u8, name, self.entry_name)) continue;
306 if (gpu.stage.stageOf(op) != stage) return error.InvalidArtifact;
307 return func;
308 }
309 return error.InvalidArtifact;
310 }
311
312 fn writeStageStructs(self: *Emitter, writer: *std.Io.Writer, stage: Stage) EmitError!void {
313 const interface = &self.interface;
314 if (interface.hasInput()) {
315 try writer.print("struct {s}_in {{\n", .{self.entry_name});
316 if (interface.frag_coord) try writer.writeAll(" float4 position [[position]];\n");
317 for (interface.inputs, 0..) |maybe_slot, location| {
318 const slot = maybe_slot orelse continue;
319 try writeMember(writer, stage, .input, location, slot);
320 }
321 try writer.writeAll("};\n\n");
322 }
323 try writer.print("struct {s}_out {{\n", .{self.entry_name});
324 if (stage == .vertex) try writer.writeAll(" float4 position [[position]];\n");
325 for (interface.outputs, 0..) |maybe_slot, location| {
326 const slot = maybe_slot orelse continue;
327 try writeMember(writer, stage, .output, location, slot);
328 }
329 try writer.writeAll("};\n\n");
330 }
331
332 fn writeStageSignature(self: *Emitter, writer: *std.Io.Writer, stage: Stage) EmitError!void {
333 const interface = &self.interface;
334 const total = @as(usize, @intFromBool(interface.hasInput())) +
335 @intFromBool(interface.vertex_id) +
336 @intFromBool(interface.instance_id) +
337 @intFromBool(interface.front_facing) +
338 2 * interface.textures.count() +
339 @intFromBool(interface.push) +
340 interface.uniforms.count();
341 const name = self.entry_name;
342 try writer.print("{s} {s}_out {s}(", .{ @tagName(stage), name, name });
343 if (total == 0) return writer.writeAll(") {\n");
344 try writer.writeByte('\n');
345 var list: ParameterList = .{ .writer = writer, .total = total };
346 if (interface.hasInput()) try list.add("{s}_in choir_in [[stage_in]]", .{name});
347 if (interface.vertex_id) try list.add("uint choir_vertex_id [[vertex_id]]", .{});
348 if (interface.instance_id) try list.add("uint choir_instance_id [[instance_id]]", .{});
349 if (interface.front_facing) try list.add("bool choir_front_facing [[front_facing]]", .{});
350 var textures = interface.textures.iterator(.{});
351 while (textures.next()) |binding| {
352 const texture = "texture2d<float> choir_texture{d} [[texture({d})]]";
353 try list.add(texture, .{ binding, binding });
354 try list.add("sampler choir_sampler{d} [[sampler({d})]]", .{ binding, binding });
355 }
356 var uniforms = interface.uniforms.iterator(.{});
357 while (uniforms.next()) |binding| {
358 const uniform = "constant uint4* choir_uniform{d} [[buffer({d})]]";
359 try list.add(uniform, .{ binding, uniform_buffer_base + binding });
360 }
361 if (interface.push) try list.add("constant uint* choir_push [[buffer({d})]]", .{push_buffer});
362 std.debug.assert(list.index == total);
363 try writer.writeAll(") {\n");
364 }
365
366 /// Emits `op` when it belongs to the stage vocabulary and reports whether
367 /// it did.
368 fn emitStageOperation(self: *Emitter, op: *ir.Operation) EmitError!bool {
369 const name = op.name.name;
370 if (std.mem.eql(u8, name, GpuDialect.StageInputOp.operation_name)) {
371 try self.emitStageInput(GpuDialect.StageInputOp{ .op = op });
372 } else if (std.mem.eql(u8, name, GpuDialect.StageOutputOp.operation_name)) {
373 try self.emitStageOutput(GpuDialect.StageOutputOp{ .op = op });
374 } else if (std.mem.eql(u8, name, GpuDialect.PositionOp.operation_name)) {
375 const value = try self.composite(.f32, op.getOperandValues());
376 try self.line("choir_out.position = {s};", .{value});
377 try self.line("choir_out.position.y = -choir_out.position.y;", .{});
378 } else if (std.mem.eql(u8, name, GpuDialect.FragCoordOp.operation_name)) {
379 self.interface.frag_coord = true;
380 try self.bindComponents(op, "choir_in.position");
381 } else if (std.mem.eql(u8, name, GpuDialect.VertexIndexOp.operation_name)) {
382 self.interface.vertex_id = true;
383 try self.bind(op.getResult(0).?, "choir_vertex_id");
384 } else if (std.mem.eql(u8, name, GpuDialect.InstanceIndexOp.operation_name)) {
385 self.interface.instance_id = true;
386 try self.bind(op.getResult(0).?, "choir_instance_id");
387 } else if (std.mem.eql(u8, name, GpuDialect.FrontFacingOp.operation_name)) {
388 self.interface.front_facing = true;
389 try self.bind(op.getResult(0).?, "choir_front_facing");
390 } else if (std.mem.eql(u8, name, GpuDialect.SampledTextureOp.operation_name)) {
391 try self.recordTexture(GpuDialect.SampledTextureOp{ .op = op });
392 } else if (std.mem.eql(u8, name, GpuDialect.SampleOp.operation_name)) {
393 try self.emitSample(op, false);
394 } else if (std.mem.eql(u8, name, GpuDialect.SampleLodOp.operation_name)) {
395 try self.emitSample(op, true);
396 } else if (std.mem.eql(u8, name, GpuDialect.PushConstantOp.operation_name)) {
397 try self.emitPushConstant(GpuDialect.PushConstantOp{ .op = op });
398 } else if (std.mem.eql(u8, name, GpuDialect.UniformOp.operation_name)) {
399 try self.emitUniform(GpuDialect.UniformOp{ .op = op });
400 } else if (std.mem.eql(u8, name, GpuDialect.DpdxOp.operation_name)) {
401 try self.emitCall1(op, "dfdx");
402 } else if (std.mem.eql(u8, name, GpuDialect.DpdyOp.operation_name)) {
403 try self.emitCall1(op, "dfdy");
404 } else if (std.mem.eql(u8, name, GpuDialect.FwidthOp.operation_name)) {
405 try self.emitCall1(op, "fwidth");
406 } else {
407 return false;
408 }
409 return true;
410 }
411
412 fn emitStageInput(self: *Emitter, op: GpuDialect.StageInputOp) EmitError!void {
413 const location = op.getLocation() orelse return error.InvalidArtifact;
414 const width: u32 = @intCast(op.op.getNumResults());
415 const kind = try interfaceKind(op.op.getResult(0).?.type);
416 for (0..width) |index| {
417 const component_kind = try interfaceKind(op.op.getResult(index).?.type);
418 if (component_kind != kind) return error.InvalidArtifact;
419 }
420 try self.interface.record(.input, location, .{ .kind = kind, .width = width });
421 const member = std.fmt.allocPrint(self.allocator, "choir_in.{s}{d}", .{
422 memberPrefix(self.stage.?, .input),
423 location,
424 }) catch return error.OutOfMemory;
425 try self.bindComponents(op.op, try self.rememberName(member));
426 }
427
428 fn emitStageOutput(self: *Emitter, op: GpuDialect.StageOutputOp) EmitError!void {
429 const location = op.getLocation() orelse return error.InvalidArtifact;
430 const values = op.op.getOperandValues();
431 const kind = try interfaceKind(values[0].type);
432 const slot: Slot = .{ .kind = kind, .width = @intCast(values.len) };
433 try self.interface.record(.output, location, slot);
434 const value = try self.composite(kind, values);
435 const member = memberPrefix(self.stage.?, .output);
436 try self.line("choir_out.{s}{d} = {s};", .{ member, location, value });
437 }
438
439 fn recordTexture(self: *Emitter, op: GpuDialect.SampledTextureOp) EmitError!void {
440 const group = op.getGroup() orelse return error.InvalidArtifact;
441 const binding = op.getBinding() orelse return error.InvalidArtifact;
442 if (group != 0 or binding >= max_texture_bindings) return error.CapabilityMismatch;
443 self.interface.textures.set(binding);
444 const name = std.fmt.allocPrint(self.allocator, "choir_texture{d}", .{binding}) catch
445 return error.OutOfMemory;
446 try self.rememberAndBind(op.getResult(), name);
447 }
448
449 fn emitPushConstant(self: *Emitter, op: GpuDialect.PushConstantOp) EmitError!void {
450 const member = op.member() orelse return error.InvalidArtifact;
451 if (!member.placed(gpu.stage.max_push_constant_bytes)) return error.CapabilityMismatch;
452 const kind = try blockReadKind(op.op);
453 self.interface.push = true;
454 for (0..member.width) |component| {
455 const word = member.offset / 4 + component;
456 const name = if (kind == .u32)
457 std.fmt.allocPrint(self.allocator, "choir_push[{d}]", .{word})
458 else
459 std.fmt.allocPrint(self.allocator, "as_type<{s}>(choir_push[{d}])", .{
460 try mslScalarType(kind),
461 word,
462 });
463 try self.rememberAndBind(op.op.getResult(component).?, name catch return error.OutOfMemory);
464 }
465 }
466
467 fn emitUniform(self: *Emitter, op: GpuDialect.UniformOp) EmitError!void {
468 const member = op.member() orelse return error.InvalidArtifact;
469 const group = op.getGroup() orelse return error.InvalidArtifact;
470 const binding = op.getBinding() orelse return error.InvalidArtifact;
471 if (group != 0 or binding >= max_uniform_bindings) return error.CapabilityMismatch;
472 if (!member.placed(gpu.stage.max_uniform_bytes)) return error.CapabilityMismatch;
473 const kind = try blockReadKind(op.op);
474 self.interface.uniforms.set(binding);
475 const row = member.offset / 16;
476 const first = member.offset % 16 / 4;
477 std.debug.assert(first + member.width <= 4);
478 for (0..member.width) |component| {
479 const column = components[first + component];
480 const name = if (kind == .u32)
481 std.fmt.allocPrint(self.allocator, "choir_uniform{d}[{d}].{s}", .{ binding, row, column })
482 else
483 std.fmt.allocPrint(self.allocator, "as_type<{s}>(choir_uniform{d}[{d}].{s})", .{
484 try mslScalarType(kind),
485 binding,
486 row,
487 column,
488 });
489 try self.rememberAndBind(op.op.getResult(component).?, name catch return error.OutOfMemory);
490 }
491 }
492
493 fn emitSample(self: *Emitter, op: *ir.Operation, explicit_lod: bool) EmitError!void {
494 const operands = op.getOperandValues();
495 const texture_ptr = operands[0].getDefiningOp() orelse return error.InvalidArtifact;
496 const texture_op: *ir.Operation = @ptrCast(@alignCast(texture_ptr));
497 if (!std.mem.eql(u8, texture_op.name.name, GpuDialect.SampledTextureOp.operation_name)) {
498 return error.InvalidArtifact;
499 }
500 const binding = (GpuDialect.SampledTextureOp{ .op = texture_op }).getBinding() orelse
501 return error.InvalidArtifact;
502 for (operands[1..]) |coordinate| {
503 if (try scalarKind(coordinate.type) != .f32) return error.UnsupportedOperation;
504 }
505 const out = try self.freshValueName();
506 const u = try self.require(operands[1]);
507 const v = try self.require(operands[2]);
508 const sample = "const float4 {s} = " ++
509 "choir_texture{d}.sample(choir_sampler{d}, float2({s}, {s})";
510 if (explicit_lod) {
511 const lod = try self.require(operands[3]);
512 try self.line(sample ++ ", level({s}));", .{ out, binding, binding, u, v, lod });
513 } else {
514 try self.line(sample ++ ");", .{ out, binding, binding, u, v });
515 }
516 try self.bindComponents(op, out);
517 }
518
519 /// Binds each result of `op` to one component of `vector`, or its single
520 /// result to `vector` whole.
521 fn bindComponents(self: *Emitter, op: *ir.Operation, vector: []const u8) abi.Error!void {
522 const count = op.getNumResults();
523 if (count == 1) return self.bind(op.getResult(0).?, vector);
524 for (0..count) |index| {
525 const component = components[index];
526 const name = std.fmt.allocPrint(self.allocator, "{s}.{s}", .{ vector, component }) catch
527 return error.OutOfMemory;
528 try self.rememberAndBind(op.getResult(index).?, name);
529 }
530 }
531
532 /// The expression that builds `values` into one value of `kind`.
533 fn composite(self: *Emitter, kind: ScalarKind, values: []const *ir.Value) EmitError![]const u8 {
534 for (values) |value| {
535 if (try interfaceKind(value.type) != kind) return error.InvalidArtifact;
536 }
537 if (values.len == 1) return self.require(values[0]);
538 var text = std.Io.Writer.Allocating.init(self.allocator);
539 defer text.deinit();
540 const width: u32 = @intCast(values.len);
541 try text.writer.print("{s}{s}(", .{ try mslScalarType(kind), widthSuffix(width) });
542 for (values, 0..) |value, index| {
543 if (index > 0) try text.writer.writeAll(", ");
544 try text.writer.writeAll(try self.require(value));
545 }
546 try text.writer.writeByte(')');
547 return self.rememberName(text.toOwnedSlice() catch return error.OutOfMemory);
548 }
549
550 fn findKernelFunction(self: *Emitter) abi.Error!FuncDialect.FuncOp {
551 if (!std.mem.eql(u8, self.module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
552 return error.InvalidArtifact;
553 }
554 const block = self.module.getRegion(0).?.getEntryBlock() orelse return error.InvalidArtifact;
555 var ops = block.getOperations();
556 while (ops.next()) |op| {
557 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;
558 const func = FuncDialect.FuncOp{ .op = op };
559 if (!func.isKernel()) continue;
560 const name = func.getName() orelse return error.InvalidArtifact;
561 if (std.mem.eql(u8, name, self.entry_name)) return func;
562 }
563 return error.InvalidArtifact;
564 }
565
566 fn emitParameterBindings(self: *Emitter, func: FuncDialect.FuncOp) abi.Error!void {
567 const args = func.getArguments();
568 for (args, 0..) |arg, index| {
569 if (memrefElementKind(arg.type) != null) {
570 const name = std.fmt.allocPrint(self.allocator, "arg{d}", .{index}) catch return error.OutOfMemory;
571 try self.rememberAndBind(arg, name);
572 continue;
573 }
574 _ = try scalarKind(arg.type);
575 const expr = std.fmt.allocPrint(self.allocator, "(*arg{d})", .{index}) catch return error.OutOfMemory;
576 try self.rememberAndBind(arg, expr);
577 }
578 }
579
580 fn emitFunctionHeader(self: *Emitter, writer: *std.Io.Writer, func: FuncDialect.FuncOp) EmitError!void {
581 try writer.print("kernel void {s}(\n", .{self.entry_name});
582 const args = func.getArguments();
583 const total = args.len + 7;
584 var index: usize = 0;
585 for (args) |arg| {
586 const suffix = if (index + 1 == total) "" else ",";
587 if (memrefElementKind(arg.type)) |kind| {
588 try writer.print(" device {s}* arg{d} [[buffer({d})]]{s}\n", .{ try mslScalarType(kind), index, index, suffix });
589 } else {
590 try writer.print(" constant {s}* arg{d} [[buffer({d})]]{s}\n", .{ try mslScalarType(try scalarKind(arg.type)), index, index, suffix });
591 }
592 index += 1;
593 }
594 try writeBuiltinParameter(writer, &index, total, "choir_thread_position_in_grid", "thread_position_in_grid");
595 try writeBuiltinParameter(writer, &index, total, "choir_thread_position_in_threadgroup", "thread_position_in_threadgroup");
596 try writeBuiltinParameter(writer, &index, total, "choir_threadgroup_position_in_grid", "threadgroup_position_in_grid");
597 try writeBuiltinParameter(writer, &index, total, "choir_threads_per_threadgroup", "threads_per_threadgroup");
598 try writeBuiltinParameter(writer, &index, total, "choir_threads_per_grid", "threads_per_grid");
599 try writeScalarBuiltinParameter(writer, &index, total, "choir_thread_index_in_simdgroup", "thread_index_in_simdgroup");
600 try writeScalarBuiltinParameter(writer, &index, total, "choir_simdgroup_index_in_threadgroup", "simdgroup_index_in_threadgroup");
601 try writer.writeAll(") {\n");
602 }
603
604 fn emitBlock(self: *Emitter, block: *ir.Block, yield_target: ?YieldTarget) EmitError!void {
605 var ops = block.getOperations();
606 while (ops.next()) |op| {
607 try self.emitOperation(op, yield_target);
608 }
609 }
610
611 fn emitOperation(self: *Emitter, op: *ir.Operation, yield_target: ?YieldTarget) EmitError!void {
612 const name = op.name.name;
613 if (!gpu.stage.admits(self.stage, name)) return error.UnsupportedOperation;
614 if (self.stage != null) {
615 const memref_prefix = MemrefDialect.name ++ ".";
616 if (std.mem.startsWith(u8, name, memref_prefix) and
617 !std.mem.eql(u8, name, MemrefDialect.AllocaOp.operation_name) and
618 !std.mem.eql(u8, name, MemrefDialect.LoadOp.operation_name) and
619 !std.mem.eql(u8, name, MemrefDialect.StoreOp.operation_name)) return error.UnsupportedOperation;
620 if (try self.emitStageOperation(op)) return;
621 }
622 if (std.mem.eql(u8, name, FuncDialect.ReturnOp.operation_name)) {
623 if (self.stage != null) {
624 try self.line("{s}", .{self.return_statement});
625 } else if (op.getNumOperands() == 0) {
626 try self.line("return;", .{});
627 } else if (op.getNumOperands() == 1) {
628 try self.line("return {s};", .{try self.require(op.getOperandValues()[0])});
629 } else return error.UnsupportedOperation;
630 } else if (std.mem.eql(u8, name, FuncDialect.CallOp.operation_name)) {
631 try self.emitCall(FuncDialect.CallOp{ .op = op });
632 } else if (std.mem.eql(u8, name, ScfDialect.YieldOp.operation_name)) {
633 try self.emitYield(ScfDialect.YieldOp{ .op = op }, yield_target);
634 } else if (std.mem.eql(u8, name, ScfDialect.IfOp.operation_name)) {
635 try self.emitIf(ScfDialect.IfOp{ .op = op });
636 } else if (std.mem.eql(u8, name, ScfDialect.ForOp.operation_name)) {
637 try self.emitFor(ScfDialect.ForOp{ .op = op });
638 } else if (std.mem.eql(u8, name, ScfDialect.WhileOp.operation_name)) {
639 try self.emitWhile(ScfDialect.WhileOp{ .op = op });
640 } else if (std.mem.eql(u8, name, GpuDialect.GlobalIdxOp.operation_name)) {
641 const wrapped = GpuDialect.GlobalIdxOp{ .op = op };
642 try self.emitGpuRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_thread_position_in_grid");
643 } else if (std.mem.eql(u8, name, GpuDialect.ThreadIdxOp.operation_name)) {
644 const wrapped = GpuDialect.ThreadIdxOp{ .op = op };
645 try self.emitGpuRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_thread_position_in_threadgroup");
646 } else if (std.mem.eql(u8, name, GpuDialect.BlockIdxOp.operation_name)) {
647 const wrapped = GpuDialect.BlockIdxOp{ .op = op };
648 try self.emitGpuRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_threadgroup_position_in_grid");
649 } else if (std.mem.eql(u8, name, GpuDialect.BlockDimOp.operation_name)) {
650 const wrapped = GpuDialect.BlockDimOp{ .op = op };
651 try self.emitGpuRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact, "choir_threads_per_threadgroup");
652 } else if (std.mem.eql(u8, name, GpuDialect.GridDimOp.operation_name)) {
653 const wrapped = GpuDialect.GridDimOp{ .op = op };
654 try self.emitGridDimRegister(op, wrapped.getDimension() orelse return error.InvalidArtifact);
655 } else if (std.mem.eql(u8, name, GpuDialect.LaneIdOp.operation_name)) {
656 try self.emitGpuScalarRegister(op, "choir_thread_index_in_simdgroup");
657 } else if (std.mem.eql(u8, name, GpuDialect.WarpIdOp.operation_name)) {
658 try self.emitGpuScalarRegister(op, "choir_simdgroup_index_in_threadgroup");
659 } else if (std.mem.eql(u8, name, GpuDialect.BarrierOp.operation_name)) {
660 try self.emitBarrier(GpuDialect.BarrierOp{ .op = op });
661 } else if (std.mem.eql(u8, name, GpuDialect.WarpReduceOp.operation_name)) {
662 try self.emitWarpReduce(GpuDialect.WarpReduceOp{ .op = op });
663 } else if (std.mem.eql(u8, name, GpuDialect.WarpScanOp.operation_name)) {
664 try self.emitWarpScan(GpuDialect.WarpScanOp{ .op = op });
665 } else if (std.mem.eql(u8, name, MemrefDialect.AllocOp.operation_name)) {
666 try self.emitAlloc(MemrefDialect.AllocOp{ .op = op });
667 } else if (std.mem.eql(u8, name, MemrefDialect.AllocaOp.operation_name)) {
668 try self.emitAlloca(MemrefDialect.AllocaOp{ .op = op });
669 } else if (std.mem.eql(u8, name, MemrefDialect.LoadOp.operation_name)) {
670 try self.emitLoad(MemrefDialect.LoadOp{ .op = op });
671 } else if (std.mem.eql(u8, name, MemrefDialect.StoreOp.operation_name)) {
672 try self.emitStore(MemrefDialect.StoreOp{ .op = op });
673 } else if (std.mem.eql(u8, name, MemrefDialect.AtomicRmwOp.operation_name)) {
674 try self.emitAtomicRmw(MemrefDialect.AtomicRmwOp{ .op = op });
675 } else if (std.mem.eql(u8, name, MemrefDialect.AtomicCasOp.operation_name)) {
676 try self.emitAtomicCas(MemrefDialect.AtomicCasOp{ .op = op });
677 } else if (std.mem.eql(u8, name, ArithDialect.ConstantOp.operation_name)) {
678 try self.emitConstant(ArithDialect.ConstantOp{ .op = op });
679 } else if (std.mem.eql(u8, name, ArithDialect.AddOp.operation_name)) {
680 try self.emitBinary(op, "+");
681 } else if (std.mem.eql(u8, name, ArithDialect.SubOp.operation_name)) {
682 try self.emitBinary(op, "-");
683 } else if (std.mem.eql(u8, name, ArithDialect.MulOp.operation_name)) {
684 try self.emitBinary(op, "*");
685 } else if (std.mem.eql(u8, name, ArithDialect.UmulhiOp.operation_name)) {
686 try self.emitUmulhi(op);
687 } else if (std.mem.eql(u8, name, ArithDialect.DivOp.operation_name)) {
688 try self.emitBinary(op, "/");
689 } else if (std.mem.eql(u8, name, ArithDialect.MaxOp.operation_name)) {
690 try self.emitCall2(op, "max");
691 } else if (std.mem.eql(u8, name, ArithDialect.MinOp.operation_name)) {
692 try self.emitCall2(op, "min");
693 } else if (std.mem.eql(u8, name, ArithDialect.AndOp.operation_name)) {
694 try self.emitBinary(op, "&");
695 } else if (std.mem.eql(u8, name, ArithDialect.OrOp.operation_name)) {
696 try self.emitBinary(op, "|");
697 } else if (std.mem.eql(u8, name, ArithDialect.XorOp.operation_name)) {
698 try self.emitBinary(op, "^");
699 } else if (std.mem.eql(u8, name, ArithDialect.ShlOp.operation_name)) {
700 try self.emitBinary(op, "<<");
701 } else if (std.mem.eql(u8, name, ArithDialect.ShrOp.operation_name)) {
702 try self.emitBinary(op, ">>");
703 } else if (std.mem.eql(u8, name, ArithDialect.UshrOp.operation_name)) {
704 try self.emitUnsignedShiftRight(op);
705 } else if (std.mem.eql(u8, name, ArithDialect.NegOp.operation_name)) {
706 try self.emitUnary(op, "-");
707 } else if (std.mem.eql(u8, name, ArithDialect.NotOp.operation_name)) {
708 try self.emitNot(op);
709 } else if (std.mem.eql(u8, name, ArithDialect.AbsOp.operation_name)) {
710 try self.emitCall1(op, "abs");
711 } else if (std.mem.eql(u8, name, ArithDialect.SqrtOp.operation_name)) {
712 try self.emitCall1(op, "sqrt");
713 } else if (std.mem.eql(u8, name, ArithDialect.ExpOp.operation_name)) {
714 try self.emitCall1(op, "exp");
715 } else if (std.mem.eql(u8, name, ArithDialect.LogOp.operation_name)) {
716 try self.emitCall1(op, "log");
717 } else if (std.mem.eql(u8, name, ArithDialect.TanhOp.operation_name)) {
718 try self.emitCall1(op, "tanh");
719 } else if (std.mem.eql(u8, name, ArithDialect.SinOp.operation_name)) {
720 try self.emitCall1(op, "sin");
721 } else if (std.mem.eql(u8, name, ArithDialect.CosOp.operation_name)) {
722 try self.emitCall1(op, "cos");
723 } else if (std.mem.eql(u8, name, ArithDialect.TanOp.operation_name)) {
724 try self.emitCall1(op, "tan");
725 } else if (std.mem.eql(u8, name, ArithDialect.FloorOp.operation_name)) {
726 try self.emitCall1(op, "floor");
727 } else if (std.mem.eql(u8, name, ArithDialect.RoundOp.operation_name)) {
728 try self.emitRound(op);
729 } else if (std.mem.eql(u8, name, ArithDialect.TruncOp.operation_name)) {
730 try self.emitCall1(op, "trunc");
731 } else if (std.mem.eql(u8, name, ArithDialect.PowOp.operation_name)) {
732 try self.emitCall2(op, "pow");
733 } else if (std.mem.eql(u8, name, ArithDialect.Atan2Op.operation_name)) {
734 try self.emitCall2(op, "atan2");
735 } else if (std.mem.eql(u8, name, ArithDialect.FmaOp.operation_name)) {
736 try self.emitFma(ArithDialect.FmaOp{ .op = op });
737 } else if (std.mem.eql(u8, name, ArithDialect.CmpOp.operation_name)) {
738 try self.emitCompare(ArithDialect.CmpOp{ .op = op });
739 } else if (std.mem.eql(u8, name, ArithDialect.SelectOp.operation_name)) {
740 try self.emitSelect(ArithDialect.SelectOp{ .op = op });
741 } else if (std.mem.eql(u8, name, ArithDialect.CastOp.operation_name)) {
742 try self.emitCast(ArithDialect.CastOp{ .op = op });
743 } else if (std.mem.eql(u8, name, ArithDialect.BitcastOp.operation_name)) {
744 try self.emitBitcast(ArithDialect.BitcastOp{ .op = op });
745 } else {
746 return error.UnsupportedOperation;
747 }
748 }
749
750 fn emitGpuRegister(self: *Emitter, op: *ir.Operation, dim: gpu.Dimension, builtin_name: []const u8) EmitError!void {
751 const result = op.getResult(0) orelse return error.InvalidArtifact;
752 const out = try self.freshValueName();
753 try self.line("const int {s} = int({s}.{s});", .{ out, builtin_name, dimName(dim) });
754 try self.bind(result, out);
755 }
756
757 fn emitGridDimRegister(self: *Emitter, op: *ir.Operation, dim: gpu.Dimension) EmitError!void {
758 const result = op.getResult(0) orelse return error.InvalidArtifact;
759 const out = try self.freshValueName();
760 try self.line("const int {s} = int(choir_threads_per_grid.{s} / choir_threads_per_threadgroup.{s});", .{ out, dimName(dim), dimName(dim) });
761 try self.bind(result, out);
762 }
763
764 fn emitGpuScalarRegister(self: *Emitter, op: *ir.Operation, builtin_name: []const u8) EmitError!void {
765 const result = op.getResult(0) orelse return error.InvalidArtifact;
766 const out = try self.freshValueName();
767 try self.line("const int {s} = int({s});", .{ out, builtin_name });
768 try self.bind(result, out);
769 }
770
771 fn emitBarrier(self: *Emitter, op: GpuDialect.BarrierOp) EmitError!void {
772 const scope = op.getScope() orelse return error.InvalidArtifact;
773 switch (scope) {
774 .block => try self.line("threadgroup_barrier(mem_flags::mem_threadgroup);", .{}),
775 else => return error.UnsupportedOperation,
776 }
777 }
778
779 fn emitWarpReduce(self: *Emitter, op: GpuDialect.WarpReduceOp) EmitError!void {
780 try self.requireFullWarpMask(op.getMask());
781 const result = op.getResult();
782 const kind = try scalarKind(result.type);
783 const out = try self.freshValueName();
784 try self.line("const {s} {s} = {s}({s});", .{
785 try mslScalarType(kind),
786 out,
787 try mslWarpReduceFunction(op.getOpKind() orelse return error.InvalidArtifact, kind),
788 try self.require(op.getValue()),
789 });
790 try self.bind(result, out);
791 }
792
793 fn emitWarpScan(self: *Emitter, op: GpuDialect.WarpScanOp) EmitError!void {
794 try self.requireFullWarpMask(op.getMask());
795 const result = op.getResult();
796 const kind = try scalarKind(result.type);
797 const out = try self.freshValueName();
798 try self.line("const {s} {s} = {s}({s});", .{
799 try mslScalarType(kind),
800 out,
801 try mslWarpScanFunction(op.getOpKind() orelse return error.InvalidArtifact, op.isInclusive(), kind),
802 try self.require(op.getValue()),
803 });
804 try self.bind(result, out);
805 }
806
807 fn requireFullWarpMask(_: *Emitter, mask: *ir.Value) abi.Error!void {
808 const defining = mask.getDefiningOp() orelse return error.UnsupportedOperation;
809 const op: *ir.Operation = @ptrCast(@alignCast(defining));
810 if (!std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) return error.UnsupportedOperation;
811 const constant = ArithDialect.ConstantOp{ .op = op };
812 const int_value = constant.getIntValue() orelse return error.UnsupportedOperation;
813 if (int_value != -1) return error.UnsupportedOperation;
814 }
815
816 fn emitAlloc(self: *Emitter, op: MemrefDialect.AllocOp) EmitError!void {
817 if (op.getDynamicSize() != null) return error.UnsupportedOperation;
818 const result = op.getResult();
819 const info = memrefInfo(result.type) orelse return error.InvalidArtifact;
820 if (info.addr_space != .shared) return error.UnsupportedOperation;
821 const size = info.size orelse return error.InvalidArtifact;
822 if (size == 0) return error.InvalidArtifact;
823 const out = try self.freshSharedName();
824 try self.line("threadgroup {s} {s}[{d}];", .{ try mslScalarType(info.element), out, size });
825 try self.bind(result, out);
826 }
827
828 fn emitAlloca(self: *Emitter, op: MemrefDialect.AllocaOp) EmitError!void {
829 if (op.getDynamicSize() != null) return error.UnsupportedOperation;
830 const result = op.getResult();
831 const info = memrefInfo(result.type) orelse return error.InvalidArtifact;
832 if (info.addr_space != .local) return error.UnsupportedOperation;
833 const size = info.size orelse return error.UnsupportedOperation;
834 if (size == 0 or size > std.math.maxInt(u32)) return error.UnsupportedOperation;
835 const out = try self.freshValueName();
836 try self.line("thread {s} {s}[{d}];", .{ try mslScalarType(info.element), out, size });
837 try self.bind(result, out);
838 }
839
840 fn emitCall(self: *Emitter, call: FuncDialect.CallOp) EmitError!void {
841 const name = call.getCallee() orelse return error.InvalidArtifact;
842 if (call.getNumResults() > 1) return error.UnsupportedOperation;
843 var arguments = std.Io.Writer.Allocating.init(self.allocator);
844 defer arguments.deinit();
845 for (call.getOperands(), 0..) |arg, index| {
846 if (index != 0) try arguments.writer.writeAll(", ");
847 try arguments.writer.writeAll(try self.require(arg));
848 }
849 if (call.getResult(0)) |result| {
850 const out = try self.freshValueName();
851 try self.line("const {s} {s} = {s}({s});", .{
852 try mslScalarType(try scalarKind(result.type)), out, name, arguments.written(),
853 });
854 try self.bind(result, out);
855 } else {
856 try self.line("{s}({s});", .{ name, arguments.written() });
857 }
858 }
859
860 fn emitLoad(self: *Emitter, op: MemrefDialect.LoadOp) EmitError!void {
861 if (self.stage != null and (memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact).addr_space != .local) return error.UnsupportedOperation;
862 const result = op.getResult();
863 const kind = try scalarKind(result.type);
864 const out = try self.freshValueName();
865 try self.line("const {s} {s} = {s}[{s}];", .{
866 try mslScalarType(kind),
867 out,
868 try self.require(op.getMemref()),
869 try self.require(op.getIndex()),
870 });
871 try self.bind(result, out);
872 }
873
874 fn emitStore(self: *Emitter, op: MemrefDialect.StoreOp) EmitError!void {
875 if (self.stage != null and (memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact).addr_space != .local) return error.UnsupportedOperation;
876 try self.line("{s}[{s}] = {s};", .{
877 try self.require(op.getMemref()),
878 try self.require(op.getIndex()),
879 try self.require(op.getValue()),
880 });
881 }
882
883 fn emitAtomicRmw(self: *Emitter, op: MemrefDialect.AtomicRmwOp) EmitError!void {
884 const result = op.getResult();
885 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
886 if (try scalarKind(result.type) != info.element) return error.InvalidArtifact;
887
888 const atomic_kind = op.getKind() orelse return error.InvalidArtifact;
889 const atomic_function = try mslAtomicFunction(atomic_kind);
890 const atomic_type = try mslAtomicScalarType(atomic_kind, info.element, info.addr_space);
891 const address_space = try mslAtomicAddressSpace(info.addr_space);
892 const value_type = try mslScalarType(info.element);
893 const out = try self.freshValueName();
894 try self.line("const {s} {s} = {s}_explicit(({s} {s}*)(&{s}[{s}]), {s}, memory_order_relaxed);", .{
895 value_type,
896 out,
897 atomic_function,
898 address_space,
899 atomic_type,
900 try self.require(op.getMemref()),
901 try self.require(op.getIndex()),
902 try self.require(op.getValue()),
903 });
904 try self.bind(result, out);
905 }
906
907 fn emitAtomicCas(self: *Emitter, op: MemrefDialect.AtomicCasOp) EmitError!void {
908 const result = op.getResult();
909 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
910 const kind = try scalarKind(result.type);
911 if (kind != info.element) return error.InvalidArtifact;
912
913 const atomic_type = try mslAtomicCasScalarType(info.element);
914 const address_space = try mslAtomicAddressSpace(info.addr_space);
915 const value_type = try mslScalarType(info.element);
916 const old_value = try self.freshValueName();
917 const changed = try self.freshValueName();
918 try self.line("{s} {s} = {s};", .{
919 value_type,
920 old_value,
921 try self.require(op.getExpected()),
922 });
923 try self.line("bool {s};", .{changed});
924 try self.line("do {{", .{});
925 self.indent += 1;
926 try self.line("{s} = atomic_compare_exchange_weak_explicit(({s} {s}*)(&{s}[{s}]), &{s}, {s}, memory_order_relaxed, memory_order_relaxed);", .{
927 changed,
928 address_space,
929 atomic_type,
930 try self.require(op.getMemref()),
931 try self.require(op.getIndex()),
932 old_value,
933 try self.require(op.getDesired()),
934 });
935 self.indent -= 1;
936 try self.line("}} while (!{s} && {s} == {s});", .{
937 changed,
938 old_value,
939 try self.require(op.getExpected()),
940 });
941 try self.bind(result, old_value);
942 }
943
944 fn emitConstant(self: *Emitter, op: ArithDialect.ConstantOp) EmitError!void {
945 const result = op.getResult();
946 const kind = try scalarKind(result.type);
947 const out = try self.freshValueName();
948 switch (kind) {
949 .bool => {
950 const bool_attr = op.op.getAttrAs(ir.Attribute.BoolAttr, "value") orelse return error.InvalidArtifact;
951 try self.line("const bool {s} = {s};", .{ out, if (bool_attr.getValue()) "true" else "false" });
952 },
953 .index => {
954 const value = op.getIntValue() orelse return error.InvalidArtifact;
955 if (value < 0) return error.UnsupportedOperation;
956 try self.line("const int {s} = int({d});", .{ out, value });
957 },
958 .u32 => {
959 const value = op.getIntValue() orelse return error.InvalidArtifact;
960 if (value < 0) return error.UnsupportedOperation;
961 try self.line("const uint {s} = uint({d});", .{ out, value });
962 },
963 .i8, .i16, .i32, .i64 => {
964 const value = op.getIntValue() orelse return error.InvalidArtifact;
965 try self.line("const {s} {s} = {s}({d});", .{ try mslScalarType(kind), out, try mslScalarType(kind), value });
966 },
967 .f16 => {
968 const value = op.getFloatValue() orelse return error.InvalidArtifact;
969 try self.line("const half {s} = half(as_type<float>(0x{X:0>8}u));", .{ out, floatBits(value) });
970 },
971 .f32 => {
972 const value = op.getFloatValue() orelse return error.InvalidArtifact;
973 try self.line("const float {s} = as_type<float>(0x{X:0>8}u);", .{ out, floatBits(value) });
974 },
975 .u8, .u16, .u64, .bf16, .f64 => return error.UnsupportedOperation,
976 }
977 try self.bind(result, out);
978 }
979
980 fn emitBinary(self: *Emitter, op: *ir.Operation, operator: []const u8) EmitError!void {
981 const result = op.getResult(0) orelse return error.InvalidArtifact;
982 const out = try self.freshValueName();
983 try self.line("const {s} {s} = {s} {s} {s};", .{
984 try mslScalarType(try scalarKind(result.type)),
985 out,
986 try self.require(op.operands.items[0].value),
987 operator,
988 try self.require(op.operands.items[1].value),
989 });
990 try self.bind(result, out);
991 }
992
993 fn emitUnary(self: *Emitter, op: *ir.Operation, operator: []const u8) EmitError!void {
994 const result = op.getResult(0) orelse return error.InvalidArtifact;
995 const out = try self.freshValueName();
996 try self.line("const {s} {s} = {s}{s};", .{
997 try mslScalarType(try scalarKind(result.type)),
998 out,
999 operator,
1000 try self.require(op.operands.items[0].value),
1001 });
1002 try self.bind(result, out);
1003 }
1004
1005 fn emitUnsignedShiftRight(self: *Emitter, op: *ir.Operation) EmitError!void {
1006 const result = op.getResult(0) orelse return error.InvalidArtifact;
1007 const kind = try scalarKind(result.type);
1008 const signed_type = try mslScalarType(kind);
1009 const unsigned_type = try mslUnsignedScalarType(kind);
1010 const out = try self.freshValueName();
1011 try self.line("const {s} {s} = as_type<{s}>({s}(as_type<{s}>({s}) >> as_type<{s}>({s})));", .{
1012 signed_type,
1013 out,
1014 signed_type,
1015 unsigned_type,
1016 unsigned_type,
1017 try self.require(op.operands.items[0].value),
1018 unsigned_type,
1019 try self.require(op.operands.items[1].value),
1020 });
1021 try self.bind(result, out);
1022 }
1023
1024 fn emitNot(self: *Emitter, op: *ir.Operation) EmitError!void {
1025 const result = op.getResult(0) orelse return error.InvalidArtifact;
1026 const out = try self.freshValueName();
1027 const kind = try scalarKind(result.type);
1028 const operator = switch (kind) {
1029 .bool => "!",
1030 .index, .i8, .i16, .i32, .u32, .i64 => "~",
1031 else => return error.UnsupportedOperation,
1032 };
1033 try self.line("const {s} {s} = {s}{s};", .{
1034 try mslScalarType(kind),
1035 out,
1036 operator,
1037 try self.require(op.operands.items[0].value),
1038 });
1039 try self.bind(result, out);
1040 }
1041
1042 fn emitCall1(self: *Emitter, op: *ir.Operation, function_name: []const u8) EmitError!void {
1043 const result = op.getResult(0) orelse return error.InvalidArtifact;
1044 const out = try self.freshValueName();
1045 try self.line("const {s} {s} = {s}({s});", .{
1046 try mslScalarType(try scalarKind(result.type)),
1047 out,
1048 function_name,
1049 try self.require(op.operands.items[0].value),
1050 });
1051 try self.bind(result, out);
1052 }
1053
1054 fn emitCall2(self: *Emitter, op: *ir.Operation, function_name: []const u8) EmitError!void {
1055 const result = op.getResult(0) orelse return error.InvalidArtifact;
1056 const out = try self.freshValueName();
1057 try self.line("const {s} {s} = {s}({s}, {s});", .{
1058 try mslScalarType(try scalarKind(result.type)),
1059 out,
1060 function_name,
1061 try self.require(op.operands.items[0].value),
1062 try self.require(op.operands.items[1].value),
1063 });
1064 try self.bind(result, out);
1065 }
1066
1067 fn emitUmulhi(self: *Emitter, op: *ir.Operation) EmitError!void {
1068 const result = op.getResult(0) orelse return error.InvalidArtifact;
1069 const kind = try scalarKind(result.type);
1070 const signed_type = try mslScalarType(kind);
1071 const unsigned_type: []const u8 = switch (kind) {
1072 .index => "uint",
1073 .i32 => "uint",
1074 .u32 => "uint",
1075 .i64 => "ulong",
1076 else => return error.UnsupportedOperation,
1077 };
1078 const out = try self.freshValueName();
1079 try self.line("const {s} {s} = as_type<{s}>(mulhi(as_type<{s}>({s}), as_type<{s}>({s})));", .{
1080 signed_type,
1081 out,
1082 signed_type,
1083 unsigned_type,
1084 try self.require(op.operands.items[0].value),
1085 unsigned_type,
1086 try self.require(op.operands.items[1].value),
1087 });
1088 try self.bind(result, out);
1089 }
1090
1091 fn emitFma(self: *Emitter, op: ArithDialect.FmaOp) EmitError!void {
1092 const result = op.getResult();
1093 const out = try self.freshValueName();
1094 try self.line("const {s} {s} = fma({s}, {s}, {s});", .{
1095 try mslScalarType(try scalarKind(result.type)),
1096 out,
1097 try self.require(op.getA()),
1098 try self.require(op.getB()),
1099 try self.require(op.getC()),
1100 });
1101 try self.bind(result, out);
1102 }
1103
1104 fn emitCompare(self: *Emitter, op: ArithDialect.CmpOp) EmitError!void {
1105 const out = try self.freshValueName();
1106 try self.line("const bool {s} = {s} {s} {s};", .{
1107 out,
1108 try self.require(op.op.operands.items[0].value),
1109 comparisonOperator(op.getPredicate() orelse return error.InvalidArtifact),
1110 try self.require(op.op.operands.items[1].value),
1111 });
1112 try self.bind(op.getResult(), out);
1113 }
1114
1115 fn emitSelect(self: *Emitter, op: ArithDialect.SelectOp) EmitError!void {
1116 const result = op.getResult();
1117 const out = try self.freshValueName();
1118 try self.line("const {s} {s} = {s} ? {s} : {s};", .{
1119 try mslScalarType(try scalarKind(result.type)),
1120 out,
1121 try self.require(op.getCondition()),
1122 try self.require(op.getTrueValue()),
1123 try self.require(op.getFalseValue()),
1124 });
1125 try self.bind(result, out);
1126 }
1127
1128 fn emitCast(self: *Emitter, op: ArithDialect.CastOp) EmitError!void {
1129 const result = op.getResult();
1130 const out = try self.freshValueName();
1131 const ty = try mslScalarType(try scalarKind(result.type));
1132 try self.line("const {s} {s} = {s}({s});", .{
1133 ty,
1134 out,
1135 ty,
1136 try self.require(op.getInput()),
1137 });
1138 try self.bind(result, out);
1139 }
1140
1141 fn emitBitcast(self: *Emitter, op: ArithDialect.BitcastOp) EmitError!void {
1142 const result = op.getResult();
1143 const out = try self.freshValueName();
1144 const ty = try mslScalarType(try scalarKind(result.type));
1145 try self.line("const {s} {s} = as_type<{s}>({s});", .{
1146 ty,
1147 out,
1148 ty,
1149 try self.require(op.getInput()),
1150 });
1151 try self.bind(result, out);
1152 }
1153
1154 fn emitRound(self: *Emitter, op: *ir.Operation) EmitError!void {
1155 const result = op.getResult(0) orelse return error.InvalidArtifact;
1156 const input = try self.require(op.operands.items[0].value);
1157 const kind = try scalarKind(result.type);
1158 const out = try self.freshValueName();
1159 const magnitude = try self.freshValueName();
1160 const sign = try self.freshValueName();
1161 const bits = try self.freshValueName();
1162 switch (kind) {
1163 .f16 => {
1164 try self.line("const half {s} = floor(abs({s}) + half(0.5));", .{ magnitude, input });
1165 try self.line("const ushort {s} = as_type<ushort>({s}) & ushort(0x8000);", .{ sign, input });
1166 try self.line("const ushort {s} = (as_type<ushort>({s}) & ushort(0x7fff)) | {s};", .{ bits, magnitude, sign });
1167 try self.line("const half {s} = as_type<half>({s});", .{ out, bits });
1168 },
1169 .f32 => {
1170 try self.line("const float {s} = floor(abs({s}) + float(0.5));", .{ magnitude, input });
1171 try self.line("const uint {s} = as_type<uint>({s}) & 0x80000000u;", .{ sign, input });
1172 try self.line("const uint {s} = (as_type<uint>({s}) & 0x7fffffffu) | {s};", .{ bits, magnitude, sign });
1173 try self.line("const float {s} = as_type<float>({s});", .{ out, bits });
1174 },
1175 else => return error.UnsupportedOperation,
1176 }
1177 try self.bind(result, out);
1178 }
1179
1180 fn emitIf(self: *Emitter, op: ScfDialect.IfOp) EmitError!void {
1181 if (op.getNumResults() != 0) return error.UnsupportedOperation;
1182 try self.line("if ({s}) {{", .{try self.require(op.getCondition())});
1183 self.indent += 1;
1184 try self.emitBlock(op.getThenBlock(), null);
1185 self.indent -= 1;
1186 if (op.getElseBlock()) |else_block| {
1187 try self.line("}} else {{", .{});
1188 self.indent += 1;
1189 try self.emitBlock(else_block, null);
1190 self.indent -= 1;
1191 }
1192 try self.line("}}", .{});
1193 }
1194
1195 fn emitWhile(self: *Emitter, op: ScfDialect.WhileOp) EmitError!void {
1196 const before = op.getBeforeBlock();
1197 const after = op.getAfterBlock();
1198 const carry_count = op.op.operands.items.len;
1199 if (op.op.results.items.len != carry_count) return error.InvalidArtifact;
1200 if (before.arguments.items.len != carry_count) return error.InvalidArtifact;
1201 if (after.arguments.items.len != carry_count) return error.InvalidArtifact;
1202
1203 const carry_names = self.allocator.alloc([]const u8, carry_count) catch return error.OutOfMemory;
1204 defer self.allocator.free(carry_names);
1205 const exit_names = self.allocator.alloc([]const u8, carry_count) catch return error.OutOfMemory;
1206 defer self.allocator.free(exit_names);
1207
1208 for (0..carry_count) |index| {
1209 const operand = op.op.operands.items[index].value;
1210 const carry_type = try mslScalarType(try scalarKind(operand.type));
1211 carry_names[index] = try self.freshLoopName();
1212 try self.line("{s} {s} = {s};", .{ carry_type, carry_names[index], try self.require(operand) });
1213 exit_names[index] = try self.freshLoopName();
1214 try self.line("{s} {s};", .{ carry_type, exit_names[index] });
1215 try self.bind(before.arguments.items[index], carry_names[index]);
1216 }
1217
1218 try self.line("while (true) {{", .{});
1219 self.indent += 1;
1220
1221 var before_ops = before.getOperations();
1222 var saw_condition = false;
1223 while (before_ops.next()) |before_op| {
1224 if (std.mem.eql(u8, before_op.name.name, ScfDialect.ConditionOp.operation_name)) {
1225 const condition = ScfDialect.ConditionOp{ .op = before_op };
1226 const args = condition.getArgs();
1227 if (args.len != carry_count) return error.InvalidArtifact;
1228 try self.line("if (!({s})) {{", .{try self.require(condition.getCondition())});
1229 self.indent += 1;
1230 for (args, exit_names) |arg, exit_name| {
1231 try self.line("{s} = {s};", .{ exit_name, try self.require(arg) });
1232 }
1233 try self.line("break;", .{});
1234 self.indent -= 1;
1235 try self.line("}}", .{});
1236 for (args, 0..) |arg, index| {
1237 try self.bind(after.arguments.items[index], try self.require(arg));
1238 }
1239 saw_condition = true;
1240 break;
1241 }
1242 try self.emitOperation(before_op, null);
1243 }
1244 if (!saw_condition) return error.InvalidArtifact;
1245
1246 try self.emitBlock(after, .{ .names = carry_names });
1247 self.indent -= 1;
1248 try self.line("}}", .{});
1249
1250 for (op.op.results.items, exit_names) |*result, exit_name| {
1251 try self.bind(result, exit_name);
1252 }
1253 }
1254
1255 fn emitFor(self: *Emitter, op: ScfDialect.ForOp) EmitError!void {
1256 const init_args = op.getInitArgs();
1257 if (op.op.results.items.len != init_args.len) return error.InvalidArtifact;
1258
1259 const accumulator_names = self.allocator.alloc([]const u8, init_args.len) catch return error.OutOfMemory;
1260 defer self.allocator.free(accumulator_names);
1261
1262 const iter_args = op.getIterArgs();
1263 for (init_args, 0..) |initial, index| {
1264 const name = try self.freshLoopName();
1265 accumulator_names[index] = name;
1266 try self.line("{s} {s} = {s};", .{
1267 try mslScalarType(try scalarKind(initial.type)),
1268 name,
1269 try self.require(initial),
1270 });
1271 try self.bind(iter_args[index], name);
1272 }
1273
1274 const iv_name = try self.freshLoopName();
1275 try self.bind(op.getInductionVar(), iv_name);
1276 const lower = try self.require(op.getLowerBound());
1277 const upper = try self.require(op.getUpperBound());
1278 const step = try self.require(op.getStep());
1279 try self.line("for (int {s} = {s}; {s} < {s}; {s} += {s}) {{", .{ iv_name, lower, iv_name, upper, iv_name, step });
1280 self.indent += 1;
1281 try self.emitBlock(op.getBodyBlock(), .{ .names = accumulator_names });
1282 self.indent -= 1;
1283 try self.line("}}", .{});
1284
1285 for (op.op.results.items, 0..) |*result, index| {
1286 try self.bind(result, accumulator_names[index]);
1287 }
1288 }
1289
1290 fn emitYield(self: *Emitter, op: ScfDialect.YieldOp, yield_target: ?YieldTarget) EmitError!void {
1291 const operands = op.getOperands();
1292 const target_names = if (yield_target) |target_binding| target_binding.names else {
1293 if (operands.len != 0) return error.UnsupportedOperation;
1294 return;
1295 };
1296 if (operands.len != target_names.len) return error.InvalidArtifact;
1297 for (operands, target_names) |operand, target_name| {
1298 try self.line("{s} = {s};", .{ target_name, try self.require(operand) });
1299 }
1300 }
1301
1302 fn freshValueName(self: *Emitter) abi.Error![]const u8 {
1303 const index = self.next_value;
1304 self.next_value += 1;
1305 const name = std.fmt.allocPrint(self.allocator, "v{d}", .{index}) catch return error.OutOfMemory;
1306 return try self.rememberName(name);
1307 }
1308
1309 fn freshLoopName(self: *Emitter) abi.Error![]const u8 {
1310 const index = self.next_loop;
1311 self.next_loop += 1;
1312 const name = std.fmt.allocPrint(self.allocator, "l{d}", .{index}) catch return error.OutOfMemory;
1313 return try self.rememberName(name);
1314 }
1315
1316 fn freshSharedName(self: *Emitter) abi.Error![]const u8 {
1317 const index = self.next_shared;
1318 self.next_shared += 1;
1319 const name = std.fmt.allocPrint(self.allocator, "shared{d}", .{index}) catch return error.OutOfMemory;
1320 return try self.rememberName(name);
1321 }
1322
1323 fn rememberAndBind(self: *Emitter, value: *const ir.Value, name: []u8) abi.Error!void {
1324 const owned = try self.rememberName(name);
1325 try self.bind(value, owned);
1326 }
1327
1328 fn rememberName(self: *Emitter, name: []u8) abi.Error![]const u8 {
1329 self.names.append(self.allocator, name) catch {
1330 self.allocator.free(name);
1331 return error.OutOfMemory;
1332 };
1333 return name;
1334 }
1335
1336 fn bind(self: *Emitter, value: *const ir.Value, name: []const u8) abi.Error!void {
1337 self.values.put(self.allocator, value, name) catch return error.OutOfMemory;
1338 }
1339
1340 fn require(self: *Emitter, value: *const ir.Value) abi.Error![]const u8 {
1341 return self.values.get(value) orelse error.InvalidArtifact;
1342 }
1343
1344 fn line(self: *Emitter, comptime fmt: []const u8, args: anytype) std.Io.Writer.Error!void {
1345 for (0..self.indent) |_| try self.body.writer.writeAll(" ");
1346 try self.body.writer.print(fmt, args);
1347 try self.body.writer.writeByte('\n');
1348 }
1349 };
1350
1351 fn writeHeader(writer: *std.Io.Writer) std.Io.Writer.Error!void {
1352 try writer.writeAll("#include <metal_stdlib>\nusing namespace metal;\n\n");
1353 }
1354
1355 fn writeMember(
1356 writer: *std.Io.Writer,
1357 stage: Stage,
1358 direction: Direction,
1359 location: usize,
1360 slot: Slot,
1361 ) EmitError!void {
1362 try writer.print(" {s}{s} {s}{d} ", .{
1363 try mslScalarType(slot.kind),
1364 widthSuffix(slot.width),
1365 memberPrefix(stage, direction),
1366 location,
1367 });
1368 const flat = slot.kind != .f32;
1369 switch (stage) {
1370 .vertex => switch (direction) {
1371 .input => try writer.print("[[attribute({d})]];\n", .{location}),
1372 .output => try writer.print("[[user(locn{d})]];\n", .{location}),
1373 },
1374 .fragment => switch (direction) {
1375 .input => if (flat)
1376 try writer.print("[[user(locn{d}), flat]];\n", .{location})
1377 else
1378 try writer.print("[[user(locn{d})]];\n", .{location}),
1379 .output => try writer.print("[[color({d})]];\n", .{location}),
1380 },
1381 }
1382 }
1383
1384 /// Writes a signature's parameters one per line, with a comma after each but
1385 /// the last.
1386 const ParameterList = struct {
1387 writer: *std.Io.Writer,
1388 total: usize,
1389 index: usize = 0,
1390
1391 fn add(self: *ParameterList, comptime fmt: []const u8, args: anytype) std.Io.Writer.Error!void {
1392 std.debug.assert(self.index < self.total);
1393 try self.writer.writeAll(" ");
1394 try self.writer.print(fmt, args);
1395 try self.writer.writeAll(if (self.index + 1 == self.total) "\n" else ",\n");
1396 self.index += 1;
1397 }
1398 };
1399
1400 fn writeBuiltinParameter(
1401 writer: *std.Io.Writer,
1402 index: *usize,
1403 total: usize,
1404 name: []const u8,
1405 attribute: []const u8,
1406 ) std.Io.Writer.Error!void {
1407 const suffix = if (index.* + 1 == total) "" else ",";
1408 try writer.print(" uint3 {s} [[{s}]]{s}\n", .{ name, attribute, suffix });
1409 index.* += 1;
1410 }
1411
1412 fn writeScalarBuiltinParameter(
1413 writer: *std.Io.Writer,
1414 index: *usize,
1415 total: usize,
1416 name: []const u8,
1417 attribute: []const u8,
1418 ) std.Io.Writer.Error!void {
1419 const suffix = if (index.* + 1 == total) "" else ",";
1420 try writer.print(" uint {s} [[{s}]]{s}\n", .{ name, attribute, suffix });
1421 index.* += 1;
1422 }
1423
1424 fn scalarKind(typ: ir.Type) abi.Error!ScalarKind {
1425 return scalar_kinds.kindFromType(typ) orelse error.UnsupportedOperation;
1426 }
1427
1428 fn memrefElementKind(typ: ir.Type) ?ScalarKind {
1429 return if (memrefInfo(typ)) |info| info.element else null;
1430 }
1431
1432 const MemrefInfo = struct {
1433 size: ?u64,
1434 element: ScalarKind,
1435 addr_space: dialects.AddressSpace,
1436 };
1437
1438 fn memrefInfo(typ: ir.Type) ?MemrefInfo {
1439 const name = typ.getDialectTypeName() orelse return null;
1440 if (!std.mem.eql(u8, name, MemrefDialect.name)) return null;
1441 const params = MemrefDialect.parseMemrefParams(typ.getDialectParamKey() orelse return null) orelse return null;
1442 return .{
1443 .size = params.size,
1444 .element = scalar_kinds.kindFromTypeName(params.element_type_name) orelse return null,
1445 .addr_space = params.addr_space,
1446 };
1447 }
1448
1449 fn mslScalarType(kind: ScalarKind) abi.Error![]const u8 {
1450 return switch (kind) {
1451 .bool => "bool",
1452 .index => "int",
1453 .i8 => "char",
1454 .i16 => "short",
1455 .i32 => "int",
1456 .u32 => "uint",
1457 .i64 => "long",
1458 .f16 => "half",
1459 .f32 => "float",
1460 .u8, .u16, .u64, .bf16, .f64 => error.UnsupportedOperation,
1461 };
1462 }
1463
1464 fn mslUnsignedScalarType(kind: ScalarKind) abi.Error![]const u8 {
1465 return switch (kind) {
1466 .index => "uint",
1467 .i8 => "uchar",
1468 .i16 => "ushort",
1469 .i32 => "uint",
1470 .u32 => "uint",
1471 .i64 => "ulong",
1472 else => error.UnsupportedOperation,
1473 };
1474 }
1475
1476 fn mslAtomicScalarType(kind: AtomicRmwKind, scalar: ScalarKind, addr_space: dialects.AddressSpace) abi.Error![]const u8 {
1477 return switch (scalar) {
1478 .index, .i32 => "atomic_int",
1479 .u32 => "atomic_uint",
1480 .f32 => switch (kind) {
1481 .add => switch (addr_space) {
1482 .host, .device, .unified => "atomic_float",
1483 .shared, .constant, .local => error.UnsupportedOperation,
1484 },
1485 .min, .max, .bit_and, .bit_or, .bit_xor, .exchange => error.UnsupportedOperation,
1486 },
1487 else => error.UnsupportedOperation,
1488 };
1489 }
1490
1491 fn mslAtomicCasScalarType(kind: ScalarKind) abi.Error![]const u8 {
1492 return switch (kind) {
1493 .index, .i32 => "atomic_int",
1494 .u32 => "atomic_uint",
1495 else => error.UnsupportedOperation,
1496 };
1497 }
1498
1499 fn mslAtomicAddressSpace(addr_space: dialects.AddressSpace) abi.Error![]const u8 {
1500 return switch (addr_space) {
1501 .host, .device, .unified => "device",
1502 .shared => "threadgroup",
1503 .constant, .local => error.UnsupportedOperation,
1504 };
1505 }
1506
1507 fn mslAtomicFunction(kind: AtomicRmwKind) abi.Error![]const u8 {
1508 return switch (kind) {
1509 .add => "atomic_fetch_add",
1510 .min => "atomic_fetch_min",
1511 .max => "atomic_fetch_max",
1512 .bit_and => "atomic_fetch_and",
1513 .bit_or => "atomic_fetch_or",
1514 .bit_xor => "atomic_fetch_xor",
1515 .exchange => "atomic_exchange",
1516 };
1517 }
1518
1519 fn mslWarpReduceFunction(kind: gpu.WarpOpKind, scalar: ScalarKind) abi.Error![]const u8 {
1520 switch (scalar) {
1521 .index, .i32, .u32, .f32 => {},
1522 else => return error.UnsupportedOperation,
1523 }
1524 return switch (kind) {
1525 .add => "simd_sum",
1526 else => error.UnsupportedOperation,
1527 };
1528 }
1529
1530 fn mslWarpScanFunction(kind: gpu.WarpOpKind, inclusive: bool, scalar: ScalarKind) abi.Error![]const u8 {
1531 switch (scalar) {
1532 .index, .i32, .u32, .f32 => {},
1533 else => return error.UnsupportedOperation,
1534 }
1535 switch (kind) {
1536 .add => {},
1537 else => return error.UnsupportedOperation,
1538 }
1539 return if (inclusive) "simd_prefix_inclusive_sum" else "simd_prefix_exclusive_sum";
1540 }
1541
1542 fn dimName(dim: gpu.Dimension) []const u8 {
1543 return switch (dim) {
1544 .x => "x",
1545 .y => "y",
1546 .z => "z",
1547 };
1548 }
1549
1550 fn comparisonOperator(pred: CmpPredicate) []const u8 {
1551 return switch (pred) {
1552 .eq => "==",
1553 .ne => "!=",
1554 .lt, .slt, .ult => "<",
1555 .le, .sle, .ule => "<=",
1556 .gt, .sgt, .ugt => ">",
1557 .ge, .sge, .uge => ">=",
1558 };
1559 }
1560
1561 fn floatBits(value: f64) u32 {
1562 const narrowed: f32 = @floatCast(value);
1563 return @bitCast(narrowed);
1564 }
1565
1566 test "metal scalar support mask follows Choir scalar spellings" {
1567 inline for (std.meta.tags(ScalarKind)) |kind| {
1568 const expected: ?ScalarKind = if (scalar_kinds.contains(kind)) kind else null;
1569 try std.testing.expectEqual(
1570 expected,
1571 scalar_kinds.kindFromTypeName(dialects.arith.scalarTypeName(kind)),
1572 );
1573 }
1574 try std.testing.expectEqual(@as(?ScalarKind, null), scalar_kinds.kindFromTypeName("arith.unknown"));
1575 }