tiny.choir.backends.gpu.metal.msl

1const std = @import("std");2const abi = @import("choir_abi");3const choir_pkg = @import("../../../root.zig");45const gpu = @import("../../../dialects/gpu/root.zig");6const calls = @import("../calls.zig");78const ir = choir_pkg.ir;9const dialects = choir_pkg.dialects;1011const Allocator = std.mem.Allocator;12const ArithDialect = dialects.ArithDialect;13const AtomicRmwKind = dialects.AtomicRmwKind;14const CmpPredicate = dialects.arith.CmpPredicate;15const BuiltinDialect = dialects.BuiltinDialect;16const FuncDialect = dialects.FuncDialect;17const GpuDialect = gpu.GpuDialect;18const MemrefDialect = dialects.MemrefDialect;19const ScfDialect = dialects.ScfDialect;20const Stage = gpu.Stage;21const max_locations = gpu.stage.max_locations;2223const EmitError = RenderError || std.Io.Writer.Error;2425/// The render path's failures: the ABI's, and the stage memory `calls.Plan` refuses by name.26pub const RenderError = abi.Error || gpu.stage.MemoryError || calls.SignatureError;2728const ScalarKind = dialects.arith.ScalarKind;29const 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});4142const YieldTarget = struct {43    names: []const []const u8,44};4546pub const max_texture_bindings = abi.metal.max_texture_bindings;47pub const max_uniform_bindings = abi.metal.max_uniform_bindings;48pub const uniform_buffer_base = abi.metal.uniform_buffer_base;49pub const push_buffer = abi.metal.push_buffer;5051pub const RenderEntries = struct {52    vertex: []const u8,53    fragment: []const u8,54};5556/// Emits one source holding `entries.vertex` as a vertex function and57/// `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 binds61/// 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 gpu65/// contract and Vulkan do, and Metal's puts y = +1 there. So the vertex66/// function negates the y it writes to `[[position]]`, which is exact, and the67/// 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-constant70/// words are `constant uint*` at `[[buffer(push_buffer)]]`, which a host sets71/// with `setVertexBytes` and `setFragmentBytes`. Uniform binding `b` of group72/// 0 is `constant uint4*` at `[[buffer(uniform_buffer_base + b)]]`, rows of73/// 16 bytes as std140 lays them. A read reinterprets its word with `as_type`.74pub 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}101102fn 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}113114const Direction = enum { input, output };115116const Slot = struct {117    kind: ScalarKind,118    width: u32,119};120121/// What a stage function reads and writes, gathered while its body is122/// emitted and written into its structs and signature afterwards.123const 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,133134    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    }151152    fn count(slots: []const ?Slot) usize {153        var total: usize = 0;154        for (slots) |slot| total += @intFromBool(slot != null);155        return total;156    }157158    fn hasInput(self: *const StageInterface) bool {159        return count(&self.inputs) > 0 or self.frag_coord;160    }161};162163/// The struct member prefix of a slot: attributes feed the vertex stage,164/// locations join the stages, and colors leave the fragment stage.165fn 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}171172fn 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}181182const components = [_][]const u8{ "x", "y", "z", "w" };183184fn 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}191192/// The kind every result of a block read shares.193fn 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}200201pub 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}214215const 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;",229230    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    }238239    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    }245246    fn emit(self: *Emitter) EmitError![]u8 {247        const func = try self.findKernelFunction();248        try self.emitParameterBindings(func);249        try self.emitBlock(func.getEntryBlock(), null);250251        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    }260261    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    }276277    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    }293294    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    }311312    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    }331332    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    }365366    /// Emits `op` when it belongs to the stage vocabulary and reports whether367    /// 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    }411412    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    }427428    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    }438439    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}) catch445            return error.OutOfMemory;446        try self.rememberAndBind(op.getResult(), name);447    }448449    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            else459                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    }466467    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            else483                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    }492493    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() orelse501            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    }518519    /// Binds each result of `op` to one component of `vector`, or its single520    /// 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 }) catch527                return error.OutOfMemory;528            try self.rememberAndBind(op.getResult(index).?, name);529        }530    }531532    /// 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    }549550    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    }565566    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    }579580    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    }603604    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    }610611    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) and617                !std.mem.eql(u8, name, MemrefDialect.AllocaOp.operation_name) and618                !std.mem.eql(u8, name, MemrefDialect.LoadOp.operation_name) and619                !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    }749750    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    }756757    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    }763764    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    }770771    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    }778779    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    }792793    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    }806807    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    }815816    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    }827828    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    }839840    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    }859860    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    }873874    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    }882883    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;887888        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    }906907    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;912913        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    }943944    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    }979980    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    }992993    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    }10041005    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    }10231024    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    }10411042    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    }10531054    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    }10661067    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    }10901091    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    }11031104    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    }11141115    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    }11271128    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    }11401141    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    }11531154    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    }11791180    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    }11941195    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;12021203        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);12071208        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        }12171218        try self.line("while (true) {{", .{});1219        self.indent += 1;12201221        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;12451246        try self.emitBlock(after, .{ .names = carry_names });1247        self.indent -= 1;1248        try self.line("}}", .{});12491250        for (op.op.results.items, exit_names) |*result, exit_name| {1251            try self.bind(result, exit_name);1252        }1253    }12541255    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;12581259        const accumulator_names = self.allocator.alloc([]const u8, init_args.len) catch return error.OutOfMemory;1260        defer self.allocator.free(accumulator_names);12611262        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        }12731274        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("}}", .{});12841285        for (op.op.results.items, 0..) |*result, index| {1286            try self.bind(result, accumulator_names[index]);1287        }1288    }12891290    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    }13011302    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    }13081309    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    }13151316    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    }13221323    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    }13271328    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    }13351336    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    }13391340    fn require(self: *Emitter, value: *const ir.Value) abi.Error![]const u8 {1341        return self.values.get(value) orelse error.InvalidArtifact;1342    }13431344    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};13501351fn writeHeader(writer: *std.Io.Writer) std.Io.Writer.Error!void {1352    try writer.writeAll("#include <metal_stdlib>\nusing namespace metal;\n\n");1353}13541355fn 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            else1378                try writer.print("[[user(locn{d})]];\n", .{location}),1379            .output => try writer.print("[[color({d})]];\n", .{location}),1380        },1381    }1382}13831384/// Writes a signature's parameters one per line, with a comma after each but1385/// the last.1386const ParameterList = struct {1387    writer: *std.Io.Writer,1388    total: usize,1389    index: usize = 0,13901391    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};13991400fn 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}14111412fn 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}14231424fn scalarKind(typ: ir.Type) abi.Error!ScalarKind {1425    return scalar_kinds.kindFromType(typ) orelse error.UnsupportedOperation;1426}14271428fn memrefElementKind(typ: ir.Type) ?ScalarKind {1429    return if (memrefInfo(typ)) |info| info.element else null;1430}14311432const MemrefInfo = struct {1433    size: ?u64,1434    element: ScalarKind,1435    addr_space: dialects.AddressSpace,1436};14371438fn 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}14481449fn 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}14631464fn 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}14751476fn 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}14901491fn 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}14981499fn 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}15061507fn 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}15181519fn 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}15291530fn 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}15411542fn dimName(dim: gpu.Dimension) []const u8 {1543    return switch (dim) {1544        .x => "x",1545        .y => "y",1546        .z => "z",1547    };1548}15491550fn 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}15601561fn floatBits(value: f64) u32 {1562    const narrowed: f32 = @floatCast(value);1563    return @bitCast(narrowed);1564}15651566test "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}

1pub const msl = @import("msl.zig");