tiny.choir.backends.gpu.cpu.stage

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

1//! Lowers the vertex and fragment functions of a gpu-dialect module to host functions, so a CPU2//! rasterizer runs the same stage code a GPU does. Each stage function `name` becomes a host3//! function `name(inputs, outputs, context)` over the words `choir_abi.stage` lays out:4//! - an interface read loads its words, and a write stores them;5//! - a builtin loads its word;6//! - a sample stores its coordinates to the sample words, calls `tiny_gpu_sample`, and loads the7//!   texel it left there, because the sampler a texture is read through is chosen when the8//!   texture is bound, after the code exists;9//! - a push-constant or uniform read loads its member's words from the context. A uniform read10//!   outside group 0, at a binding past `uniform_bindings`, or beyond `uniform_window_bytes` is11//!   refused, as is a member the std140 rules of `gpu.stage.Member` do not place.12//!13//! A vertex function also stores, on entry, which smooth and flat words it writes.14//!15//! A fragment function that takes a derivative (`dpdx`, `dpdy`, `fwidth`) reads its quad, as16//! `choir_abi.stage.quad_suffix` states. It recomputes the derivative's operand at the partner lane17//! along each axis it differentiates: each op the operand depends on is lowered once more per18//! axis over the partner's input block, when a stage input or `frag_coord` reaches it, and shared19//! otherwise, however many derivatives read it. A fine derivative is then the odd lane's value20//! minus the even lane's, chosen by two selects on the lane's parity, so it keeps the sign of zero21//! a device's subtraction gives.22//!23//! The recomputation must not observe or change memory, so a derivative's operand may not depend24//! on a memref op or on a call that passes or returns anything but scalars. A call on scalars is25//! made again with the partner's operands, into the same helper. A derivative inside a region is26//! refused, which keeps the quad's four lanes in step, as is a derivative whose operand depends on27//! another derivative.2829const std = @import("std");30const abi = @import("choir_abi");31const choir = @import("../../../root.zig");32const gpu = @import("../../../dialects/gpu/root.zig");33const calls = @import("../calls.zig");3435const ir = choir.ir;36const dialects = choir.dialects;37const ArithDialect = dialects.ArithDialect;38const BuiltinDialect = dialects.BuiltinDialect;39const FuncDialect = dialects.FuncDialect;40const GpuDialect = gpu.GpuDialect;41const MemrefDialect = dialects.MemrefDialect;42const Stage = gpu.Stage;43const layout = abi.stage;4445pub const Error = abi.Error || gpu.stage.MemoryError || calls.SignatureError;4647/// Where a lowering stopped, for a caller that reports why.48pub const Refusal = struct {49    operation: []const u8 = "",50    helper: []const u8 = "",51};5253/// Lowers every stage function and its reachable ordinary callees into a host module the caller54/// owns and erases. Kernels and unreachable functions are left out. A refused op names itself.55pub fn lowerStagesToHost(56    allocator: std.mem.Allocator,57    module: *ir.Operation,58    refusal: ?*Refusal,59) Error!*ir.Operation {60    const ctx = module.getContext();61    const host_module = BuiltinDialect.ModuleOp.create(ctx, module.location) catch |err| return loweringError(err);62    errdefer host_module.op.erase();63    const body = ir.inspection.moduleBodyBlock(module) orelse return error.InvalidArtifact;64    var plan_refused: []const u8 = "";65    var call_plan = calls.Plan.init(allocator, module, &plan_refused) catch |err| switch (err) {66        error.OutOfMemory => return error.OutOfMemory,67        error.UnsupportedStageMemory => {68            if (refusal) |why| why.operation = plan_refused;69            return error.UnsupportedStageMemory;70        },71        error.UnsupportedHelperSignature => {72            if (refusal) |why| {73                why.operation = FuncDialect.FuncOp.operation_name;74                why.helper = plan_refused;75            }76            return error.UnsupportedHelperSignature;77        },78        error.UnsupportedOperation => {79            if (refusal) |why| why.operation = FuncDialect.CallOp.operation_name;80            return error.UnsupportedOperation;81        },82    };83    defer call_plan.deinit();8485    var lowerer = Lowerer{86        .allocator = allocator,87        .ctx = ctx,88        .host_body = host_module.getBodyBlock(),89        .refusal = refusal,90        .f32_type = ArithDialect.getScalarType(ctx, .f32) catch |err| return loweringError(err),91        .io_type = undefined,92        .context_type = undefined,93    };94    lowerer.io_type = MemrefDialect.getMemrefTypeDynamic(ctx, lowerer.f32_type, .host) catch |err| return loweringError(err);95    lowerer.context_type = lowerer.io_type;9697    for (call_plan.helpers.items) |helper| try lowerer.lowerHelperFunction(FuncDialect.FuncOp{ .op = helper });9899    var stages: u32 = 0;100    var ops = body.getOperations();101    while (ops.next()) |op| {102        if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;103        const stage = gpu.stage.stageOf(op) orelse continue;104        try lowerer.lowerFunction(FuncDialect.FuncOp{ .op = op }, stage);105        stages += 1;106    }107    if (stages == 0) return error.InvalidArtifact;108    ir.verifyOperation(host_module.op, ir.verify.default_options) catch return error.CompilationFailed;109    return host_module.op;110}111112const Lowerer = struct {113    allocator: std.mem.Allocator,114    ctx: *ir.Context,115    host_body: *ir.Block,116    refusal: ?*Refusal,117    f32_type: ir.Type,118    io_type: ir.Type,119    context_type: ir.Type,120    sample_declared: bool = false,121    /// The current function's buffers and context.122    inputs: *ir.Value = undefined,123    outputs: *ir.Value = undefined,124    context: *ir.Value = undefined,125    stage: Stage = .vertex,126    helper_body: bool = false,127    loc: ir.Location = .unknown,128    /// For a quad function, the word of `inputs` its input reads start at: the lane's own block,129    /// or a partner's while an op is lowered again over it.130    base: ?*ir.Value = null,131    /// The axis whose partner the lowering reads, while it lowers an op again.132    partner: ?Axis = null,133    /// Per axis of a quad function: the partner block's first word, whether the lane is odd along134    /// the axis, and each needed value as the partner computes it.135    partners: [2]Partner = undefined,136    /// Ops lowered again for a partner, per axis, across every function.137    partner_ops: [2]u32 = .{ 0, 0 },138139    fn lowerFunction(self: *Lowerer, source: FuncDialect.FuncOp, stage: Stage) Error!void {140        if (source.getArguments().len != 0 or source.getNumResults() != 0) return self.refuse(source.op);141        const name = source.getName() orelse return error.InvalidArtifact;142        const body = source.getEntryBlock();143        var needs = Needs.init(self.allocator);144        defer needs.deinit();145        if (stage == .fragment) try needs.collect(body);146        if (needs.axes != 0) return self.lowerQuadFunction(source, name, &needs);147148        const host = FuncDialect.FuncOp.create(149            self.ctx,150            source.op.location,151            name,152            &.{ self.io_type, self.io_type, self.context_type },153            &.{},154        ) catch |err| return loweringError(err);155        self.host_body.addOperation(host.op) catch |err| return loweringError(err);156        const arguments = host.getArguments();157        self.inputs = arguments[0];158        self.outputs = arguments[1];159        self.context = arguments[2];160        self.stage = stage;161        self.helper_body = false;162        self.loc = source.op.location;163        self.base = null;164165        const entry = host.getEntryBlock();166        if (stage == .vertex) try self.storeMasks(entry, body);167        var mapping = ir.Mapping.init(self.allocator);168        defer mapping.deinit();169        try self.cloneBlock(body, entry, &mapping, true);170    }171172    /// Lowers a fragment function that takes derivatives to `name ++ quad_suffix` over its quad.173    /// After each op it needs for an axis, it lowers the op again over that axis's partner block174    /// when the op varies across the quad, and maps it to its own values otherwise.175    fn lowerQuadFunction(self: *Lowerer, source: FuncDialect.FuncOp, name: []const u8, needs: *const Needs) Error!void {176        var name_buffer: [max_symbol_bytes]u8 = undefined;177        const symbol = std.fmt.bufPrint(&name_buffer, "{s}{s}", .{ name, layout.quad_suffix }) catch return self.refuse(source.op);178        const index_type = ArithDialect.getIndexType(self.ctx) catch |err| return loweringError(err);179        const host = FuncDialect.FuncOp.create(180            self.ctx,181            source.op.location,182            symbol,183            &.{ self.io_type, index_type, self.io_type, self.context_type },184            &.{},185        ) catch |err| return loweringError(err);186        self.host_body.addOperation(host.op) catch |err| return loweringError(err);187        const arguments = host.getArguments();188        self.inputs = arguments[0];189        self.outputs = arguments[2];190        self.context = arguments[3];191        self.stage = .fragment;192        self.helper_body = false;193        self.loc = source.op.location;194195        const entry = host.getEntryBlock();196        const lane = arguments[1];197        const block_words = try self.indexConstant(entry, layout.io_words);198        self.base = try self.binary(entry, ArithDialect.MulOp, lane, block_words);199        var mappings: [2]ir.Mapping = .{ ir.Mapping.init(self.allocator), ir.Mapping.init(self.allocator) };200        defer for (&mappings) |*partner_mapping| partner_mapping.deinit();201        for (&self.partners, &mappings, [_]Axis{ .x, .y }) |*partner, *partner_mapping, axis| {202            partner.mapping = partner_mapping;203            if (needs.axes & Needs.bit(axis) == 0) continue;204            const bit = try self.indexConstant(entry, Needs.bit(axis));205            const zero = try self.indexConstant(entry, 0);206            const along = try self.binary(entry, ArithDialect.AndOp, lane, bit);207            const odd = ArithDialect.CmpOp.create(self.ctx, self.loc, .ne, along, zero) catch |err| return loweringError(err);208            try add(entry, odd.op);209            partner.odd = odd.getResult();210            const partner_lane = try self.binary(entry, ArithDialect.XorOp, lane, bit);211            partner.base = try self.binary(entry, ArithDialect.MulOp, partner_lane, block_words);212        }213214        var mapping = ir.Mapping.init(self.allocator);215        defer mapping.deinit();216        var ops = source.getEntryBlock().getOperations();217        while (ops.next()) |op| {218            try self.lowerOperation(op, entry, &mapping, true);219            const axes = needs.of(op);220            for ([_]Axis{ .x, .y }) |axis| {221                if (axes & Needs.bit(axis) != 0) try self.lowerForPartner(op, entry, &mapping, axis);222            }223        }224    }225226    /// Maps `op`'s results in the partner mapping of `axis`: lowered again over the partner's227    /// block when a stage input or `frag_coord` reaches it, or its own lowering otherwise.228    fn lowerForPartner(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping, axis: Axis) Error!void {229        const partner = &self.partners[@backingInt(axis)];230        if (isDerivative(op.name.name) or touchesMemory(op)) return self.refuse(op);231        if (!self.varies(op, mapping, partner.mapping)) {232            for (0..op.getNumResults()) |index| {233                const result = op.getResult(index) orelse return error.InvalidArtifact;234                try mapValue(partner.mapping, result, mapping.lookupValue(result) orelse return error.InvalidArtifact);235            }236            return;237        }238        const own_base = self.base;239        self.base = partner.base;240        self.partner = axis;241        defer {242            self.base = own_base;243            self.partner = null;244        }245        try self.lowerOperation(op, dest, partner.mapping, false);246        self.partner_ops[@backingInt(axis)] += 1;247    }248249    /// Whether `op` computes a different value at a partner lane: it reads a stage input or250    /// `frag_coord`, holds regions, or takes an operand the partner computes differently.251    fn varies(self: *Lowerer, op: *ir.Operation, mapping: *ir.Mapping, partner_mapping: *ir.Mapping) bool {252        _ = self;253        const name = op.name.name;254        if (std.mem.eql(u8, name, GpuDialect.StageInputOp.operation_name)) return true;255        if (std.mem.eql(u8, name, GpuDialect.FragCoordOp.operation_name)) return true;256        if (op.regions.items.len != 0) return true;257        for (op.operand_values) |operand| {258            if (partner_mapping.lookupValue(operand) != mapping.lookupValue(operand)) return true;259        }260        return false;261    }262263    fn lowerDerivative(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {264        const value = op.operand_values[0];265        const name = op.name.name;266        const result = if (std.mem.eql(u8, name, GpuDialect.DpdxOp.operation_name))267            try self.difference(dest, mapping, value, .x)268        else if (std.mem.eql(u8, name, GpuDialect.DpdyOp.operation_name))269            try self.difference(dest, mapping, value, .y)270        else blk: {271            const along_x = try self.unary(dest, ArithDialect.AbsOp, try self.difference(dest, mapping, value, .x));272            const along_y = try self.unary(dest, ArithDialect.AbsOp, try self.difference(dest, mapping, value, .y));273            break :blk try self.binary(dest, ArithDialect.AddOp, along_x, along_y);274        };275        try mapValue(mapping, op.getResult(0) orelse return error.InvalidArtifact, result);276    }277278    /// The fine derivative of `value` along `axis`: the odd lane's value minus the even lane's.279    fn difference(self: *Lowerer, dest: *ir.Block, mapping: *ir.Mapping, value: *ir.Value, axis: Axis) Error!*ir.Value {280        const partner = &self.partners[@backingInt(axis)];281        const own = mapping.lookupValue(value) orelse return error.InvalidArtifact;282        const theirs = partner.mapping.lookupValue(value) orelse return error.InvalidArtifact;283        const odd = try self.select(dest, partner.odd, own, theirs);284        const even = try self.select(dest, partner.odd, theirs, own);285        return self.binary(dest, ArithDialect.SubOp, odd, even);286    }287288    fn select(self: *Lowerer, dest: *ir.Block, condition: *ir.Value, if_true: *ir.Value, if_false: *ir.Value) Error!*ir.Value {289        const op = ArithDialect.SelectOp.create(self.ctx, self.loc, condition, if_true, if_false) catch |err| return loweringError(err);290        try add(dest, op.op);291        return op.getResult();292    }293294    fn binary(self: *Lowerer, dest: *ir.Block, comptime Op: type, lhs: *ir.Value, rhs: *ir.Value) Error!*ir.Value {295        const op = Op.create(self.ctx, self.loc, lhs, rhs) catch |err| return loweringError(err);296        try add(dest, op.op);297        return op.getResult();298    }299300    fn unary(self: *Lowerer, dest: *ir.Block, comptime Op: type, operand: *ir.Value) Error!*ir.Value {301        const op = Op.create(self.ctx, self.loc, operand) catch |err| return loweringError(err);302        try add(dest, op.op);303        return op.getResult();304    }305306    fn lowerHelperFunction(self: *Lowerer, source: FuncDialect.FuncOp) Error!void {307        const name = source.getName() orelse return error.InvalidArtifact;308        const host = FuncDialect.FuncOp.create(309            self.ctx,310            source.op.location,311            name,312            source.getInputTypes() orelse return error.InvalidArtifact,313            source.getResultTypes(),314        ) catch |err| return loweringError(err);315        self.host_body.addOperation(host.op) catch |err| return loweringError(err);316        self.helper_body = true;317        self.base = null;318        self.loc = source.op.location;319        var mapping = ir.Mapping.init(self.allocator);320        defer mapping.deinit();321        for (source.getArguments(), host.getArguments()) |arg, mapped| try mapValue(&mapping, arg, mapped);322        try self.cloneBlock(source.getEntryBlock(), host.getEntryBlock(), &mapping, true);323    }324325    /// Stores the masks of smooth and flat words the function's interface writes reach.326    fn storeMasks(self: *Lowerer, dest: *ir.Block, source: *ir.Block) Error!void {327        var masks = Masks{};328        try masks.collect(source);329        try self.storeBits(dest, @truncate(masks.smooth), layout.smooth_mask_low);330        try self.storeBits(dest, @truncate(masks.smooth >> 32), layout.smooth_mask_high);331        try self.storeBits(dest, @truncate(masks.flat), layout.flat_mask_low);332        try self.storeBits(dest, @truncate(masks.flat >> 32), layout.flat_mask_high);333    }334335    fn storeBits(self: *Lowerer, dest: *ir.Block, bits: u32, word: u32) Error!void {336        const u32_type = ArithDialect.getScalarType(self.ctx, .u32) catch |err| return loweringError(err);337        const constant = ArithDialect.ConstantOp.createInt(self.ctx, self.loc, u32_type, bits) catch |err| return loweringError(err);338        try add(dest, constant.op);339        try self.storeWord(dest, constant.getResult(), self.outputs, word);340    }341342    fn cloneBlock(self: *Lowerer, source: *ir.Block, dest: *ir.Block, mapping: *ir.Mapping, function_body: bool) Error!void {343        var ops = source.getOperations();344        while (ops.next()) |op| try self.lowerOperation(op, dest, mapping, function_body);345    }346347    fn lowerOperation(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping, function_body: bool) Error!void {348        const name = op.name.name;349        if (self.partner != null and touchesMemory(op)) return self.refuse(op);350        if (isDerivative(name)) {351            if (!function_body or self.partner != null or self.base == null) return self.refuse(op);352            try self.lowerDerivative(op, dest, mapping);353        } else if (std.mem.eql(u8, name, FuncDialect.ReturnOp.operation_name)) {354            if (!function_body or (!self.helper_body and op.operand_values.len != 0)) return self.refuse(op);355            const operands = self.allocator.alloc(*ir.Value, op.operand_values.len) catch return error.OutOfMemory;356            defer self.allocator.free(operands);357            for (op.operand_values, operands) |operand, *mapped| mapped.* = mapping.lookupOrDefaultValue(operand);358            const ret = FuncDialect.ReturnOp.create(self.ctx, op.location, operands) catch |err| return loweringError(err);359            try add(dest, ret.op);360        } else if (std.mem.eql(u8, name, GpuDialect.StageInputOp.operation_name)) {361            try self.lowerInput(op, dest, mapping);362        } else if (std.mem.eql(u8, name, GpuDialect.StageOutputOp.operation_name)) {363            try self.lowerOutput(op, dest, mapping);364        } else if (std.mem.eql(u8, name, GpuDialect.PositionOp.operation_name)) {365            for (op.operand_values, 0..) |operand, component| {366                try self.storeWord(dest, mapping.lookupOrDefaultValue(operand), self.outputs, layout.position + @as(u32, @intCast(component)));367            }368        } else if (std.mem.eql(u8, name, GpuDialect.FragCoordOp.operation_name)) {369            for (0..4) |component| {370                const result = op.getResult(component) orelse return error.InvalidArtifact;371                const value = try self.loadInput(dest, layout.frag_coord + @as(u32, @intCast(component)), result.type);372                try mapValue(mapping, result, value);373            }374        } else if (std.mem.eql(u8, name, GpuDialect.VertexIndexOp.operation_name)) {375            try self.lowerBuiltin(op, dest, mapping, layout.vertex_index);376        } else if (std.mem.eql(u8, name, GpuDialect.InstanceIndexOp.operation_name)) {377            try self.lowerBuiltin(op, dest, mapping, layout.instance_index);378        } else if (std.mem.eql(u8, name, GpuDialect.FrontFacingOp.operation_name)) {379            try self.lowerFrontFacing(op, dest, mapping);380        } else if (std.mem.eql(u8, name, GpuDialect.SampledTextureOp.operation_name)) {381            try self.lowerTexture(op, dest, mapping);382        } else if (std.mem.eql(u8, name, MemrefDialect.AllocaOp.operation_name)) {383            const source_alloca = MemrefDialect.AllocaOp{ .op = op };384            const params = MemrefDialect.parseMemrefParams(385                source_alloca.getResult().type.getDialectParamKey() orelse return self.refuse(op),386            ) orelse return self.refuse(op);387            if (params.addr_space != .local or source_alloca.getDynamicSize() != null) return self.refuse(op);388            const size = params.size orelse return self.refuse(op);389            if (size == 0) return self.refuse(op);390            const elem = self.ctx.getDialectTypeFromName(params.element_type_name) catch return self.refuse(op);391            const host_type = MemrefDialect.getMemrefType1DWithAttrs(self.ctx, size, elem, .host, .{392                .alignment = params.alignment,393                .exclusive = params.exclusive,394                .indexing = params.indexing,395            }) catch |err| return loweringError(err);396            const host_alloca = MemrefDialect.AllocaOp.createStatic(self.ctx, op.location, host_type) catch |err| return loweringError(err);397            try add(dest, host_alloca.op);398            try mapValue(mapping, source_alloca.getResult(), host_alloca.getResult());399        } else if (std.mem.eql(u8, name, GpuDialect.PushConstantOp.operation_name)) {400            const member = (GpuDialect.PushConstantOp{ .op = op }).member() orelse return error.InvalidArtifact;401            if (!member.placed(layout.push_words * 4)) return self.refuse(op);402            try self.lowerBlockRead(op, dest, mapping, layout.push + member.offset / 4);403        } else if (std.mem.eql(u8, name, GpuDialect.UniformOp.operation_name)) {404            const uniform = GpuDialect.UniformOp{ .op = op };405            const member = uniform.member() orelse return error.InvalidArtifact;406            const group = uniform.getGroup() orelse return error.InvalidArtifact;407            const binding = uniform.getBinding() orelse return error.InvalidArtifact;408            if (group != 0 or binding >= layout.uniform_bindings) return self.refuse(op);409            if (!member.placed(layout.uniform_window_bytes)) return self.refuse(op);410            try self.lowerBlockRead(op, dest, mapping, layout.uniformWord(binding, member.offset));411        } else if (std.mem.eql(u8, name, GpuDialect.SampleOp.operation_name) or412            std.mem.eql(u8, name, GpuDialect.SampleLodOp.operation_name))413        {414            try self.lowerSample(op, dest, mapping);415        } else if (std.mem.eql(u8, op.name.getDialectNamespace(), GpuDialect.name) or op.successors.items.len != 0) {416            return self.refuse(op);417        } else {418            try self.cloneOperation(op, dest, mapping);419        }420    }421422    fn cloneOperation(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {423        const has_regions = op.regions.items.len != 0;424        const cloned = op.cloneWithoutRegionsMapped(mapping, .{ .clone_operands = !has_regions }) catch |err| return loweringError(err);425        errdefer cloned.erase();426        try add(dest, cloned);427        for (op.regions.items, 0..) |*source_region, region_index| {428            const dest_region = &cloned.regions.items[region_index];429            var source_block = source_region.blocks.head;430            while (source_block) |block| : (source_block = block.next) {431                const cloned_block = dest_region.addBlock() catch |err| return loweringError(err);432                mapping.mapBlock(block, cloned_block) catch return error.OutOfMemory;433                for (block.arguments.items) |argument| {434                    const cloned_argument = cloned_block.addArgument(argument.type, .unknown) catch |err| return loweringError(err);435                    try mapValue(mapping, argument, cloned_argument);436                }437            }438            source_block = source_region.blocks.head;439            while (source_block) |block| : (source_block = block.next) {440                const cloned_block = mapping.lookupBlock(block) orelse return error.InvalidArtifact;441                try self.cloneBlock(block, cloned_block, mapping, false);442            }443        }444        if (has_regions) {445            var operands: [max_region_operands]*ir.Value = undefined;446            if (op.operand_values.len > operands.len) return self.refuse(op);447            for (op.operand_values, operands[0..op.operand_values.len]) |operand, *mapped| {448                mapped.* = mapping.lookupOrDefaultValue(operand);449            }450            cloned.replaceOperands(operands[0..op.operand_values.len]) catch |err| return loweringError(err);451        }452    }453454    fn lowerInput(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {455        const location = (GpuDialect.StageInputOp{ .op = op }).getLocation() orelse return error.InvalidArtifact;456        for (0..op.getNumResults()) |component| {457            const result = op.getResult(component) orelse return error.InvalidArtifact;458            const region = switch (self.stage) {459                .vertex => layout.smooth,460                .fragment => try self.regionOf(op, result.type),461            };462            const word = layout.slot(region, location, @intCast(component));463            try mapValue(mapping, result, try self.loadInput(dest, word, result.type));464        }465    }466467    fn lowerOutput(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {468        const location = (GpuDialect.StageOutputOp{ .op = op }).getLocation() orelse return error.InvalidArtifact;469        for (op.operand_values, 0..) |operand, component| {470            const region = switch (self.stage) {471                .vertex => try self.regionOf(op, operand.type),472                .fragment => layout.smooth,473            };474            const word = layout.slot(region, location, @intCast(component));475            try self.storeWord(dest, mapping.lookupOrDefaultValue(operand), self.outputs, word);476        }477    }478479    /// Loads each result of a block read from consecutive context words from `first` on. Every480    /// result must share one 4-byte scalar type.481    fn lowerBlockRead(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping, first: u32) Error!void {482        const kind = scalarKind(op.getResult(0).?.type) orelse return self.refuse(op);483        switch (kind) {484            .f32, .u32, .i32 => {},485            else => return self.refuse(op),486        }487        for (0..op.getNumResults()) |component| {488            const result = op.getResult(component) orelse return error.InvalidArtifact;489            if (scalarKind(result.type) != kind) return self.refuse(op);490            const word = first + @as(u32, @intCast(component));491            std.debug.assert(word < layout.context_words);492            try mapValue(mapping, result, try self.loadFrom(dest, self.context, word, result.type));493        }494    }495496    fn lowerBuiltin(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping, word: u32) Error!void {497        const result = op.getResult(0) orelse return error.InvalidArtifact;498        try mapValue(mapping, result, try self.loadInput(dest, word, result.type));499    }500501    fn lowerFrontFacing(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {502        const result = op.getResult(0) orelse return error.InvalidArtifact;503        const u32_type = ArithDialect.getScalarType(self.ctx, .u32) catch |err| return loweringError(err);504        const bits = try self.loadInput(dest, layout.front_facing, u32_type);505        const zero = ArithDialect.ConstantOp.createInt(self.ctx, op.location, u32_type, 0) catch |err| return loweringError(err);506        try add(dest, zero.op);507        const facing = ArithDialect.CmpOp.create(self.ctx, op.location, .ne, bits, zero.getResult()) catch |err| return loweringError(err);508        try add(dest, facing.op);509        try mapValue(mapping, result, facing.getResult());510    }511512    /// A texture becomes the key a sample passes, since the host resolves the binding.513    fn lowerTexture(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {514        const texture = GpuDialect.SampledTextureOp{ .op = op };515        const group = texture.getGroup() orelse return error.InvalidArtifact;516        const binding = texture.getBinding() orelse return error.InvalidArtifact;517        if (group > 0xff or binding > 0xffff) return self.refuse(op);518        const i32_type = ArithDialect.getScalarType(self.ctx, .i32) catch |err| return loweringError(err);519        const key = ArithDialect.ConstantOp.createInt(self.ctx, op.location, i32_type, layout.textureKey(group, binding)) catch |err| return loweringError(err);520        try add(dest, key.op);521        try mapValue(mapping, texture.getResult(), key.getResult());522    }523524    fn lowerSample(self: *Lowerer, op: *ir.Operation, dest: *ir.Block, mapping: *ir.Mapping) Error!void {525        const operands = op.operand_values;526        if (operands.len != 3 and operands.len != 4) return error.InvalidArtifact;527        const key = mapping.lookupValue(operands[0]) orelse return self.refuse(op);528        try self.storeWord(dest, mapping.lookupOrDefaultValue(operands[1]), self.outputs, layout.sample);529        try self.storeWord(dest, mapping.lookupOrDefaultValue(operands[2]), self.outputs, layout.sample + 1);530        const lod = if (operands.len == 4) mapping.lookupOrDefaultValue(operands[3]) else blk: {531            const zero = ArithDialect.ConstantOp.createFloat(self.ctx, op.location, self.f32_type, 0) catch |err| return loweringError(err);532            try add(dest, zero.op);533            break :blk zero.getResult();534        };535        try self.storeWord(dest, lod, self.outputs, layout.sample + 2);536        try self.declareSample(key.type);537        const call = FuncDialect.CallOp.create(self.ctx, op.location, layout.sample_symbol, &.{ self.context, key, self.outputs }, &.{}) catch |err| return loweringError(err);538        try add(dest, call.op);539        for (0..4) |component| {540            const result = op.getResult(component) orelse return error.InvalidArtifact;541            try mapValue(mapping, result, try self.loadWord(dest, self.outputs, layout.sample + @as(u32, @intCast(component)), result.type));542        }543    }544545    fn declareSample(self: *Lowerer, key_type: ir.Type) Error!void {546        if (self.sample_declared) return;547        const declaration = FuncDialect.FuncOp.createDeclaration(548            self.ctx,549            self.loc,550            layout.sample_symbol,551            &.{ self.context_type, key_type, self.io_type },552            &.{},553        ) catch |err| return loweringError(err);554        try add(self.host_body, declaration.op);555        self.sample_declared = true;556    }557558    /// The region an interface value travels in between the stages: smooth for a float, flat for559    /// an integer.560    fn regionOf(self: *Lowerer, op: *ir.Operation, typ: ir.Type) Error!u32 {561        return switch (scalarKind(typ) orelse return self.refuse(op)) {562            .f32 => layout.smooth,563            .u32, .i32 => layout.flat,564            else => self.refuse(op),565        };566    }567568    fn loadWord(self: *Lowerer, dest: *ir.Block, buffer: *ir.Value, word: u32, result_type: ir.Type) Error!*ir.Value {569        std.debug.assert(word < layout.io_words);570        return self.loadFrom(dest, buffer, word, result_type);571    }572573    /// Loads `word` of the input block the lowering reads: the only one, or a lane's of a quad.574    fn loadInput(self: *Lowerer, dest: *ir.Block, word: u32, result_type: ir.Type) Error!*ir.Value {575        std.debug.assert(word < layout.io_words);576        const base = self.base orelse return self.loadFrom(dest, self.inputs, word, result_type);577        const index = try self.binary(dest, ArithDialect.AddOp, base, try self.indexConstant(dest, word));578        return self.loadAt(dest, self.inputs, index, result_type);579    }580581    /// Loads `word` of `buffer` as `result_type`, bitcasting a word that holds an integer.582    fn loadFrom(self: *Lowerer, dest: *ir.Block, buffer: *ir.Value, word: u32, result_type: ir.Type) Error!*ir.Value {583        return self.loadAt(dest, buffer, try self.indexConstant(dest, word), result_type);584    }585586    fn loadAt(self: *Lowerer, dest: *ir.Block, buffer: *ir.Value, index: *ir.Value, result_type: ir.Type) Error!*ir.Value {587        const load = MemrefDialect.LoadOp.create(self.ctx, self.loc, buffer, index, self.f32_type) catch |err| return loweringError(err);588        try add(dest, load.op);589        const loaded = load.op.getResult(0) orelse return error.CompilationFailed;590        if (scalarKind(result_type) == .f32) return loaded;591        const cast = ArithDialect.BitcastOp.create(self.ctx, self.loc, loaded, result_type) catch |err| return loweringError(err);592        try add(dest, cast.op);593        return cast.getResult();594    }595596    fn storeWord(self: *Lowerer, dest: *ir.Block, value: *ir.Value, buffer: *ir.Value, word: u32) Error!void {597        std.debug.assert(word < layout.io_words);598        const stored = switch (scalarKind(value.type) orelse return error.UnsupportedOperation) {599            .f32 => value,600            .u32, .i32 => blk: {601                const cast = ArithDialect.BitcastOp.create(self.ctx, self.loc, value, self.f32_type) catch |err| return loweringError(err);602                try add(dest, cast.op);603                break :blk cast.getResult();604            },605            else => return error.UnsupportedOperation,606        };607        const index = try self.indexConstant(dest, word);608        const store = MemrefDialect.StoreOp.create(self.ctx, self.loc, stored, buffer, index) catch |err| return loweringError(err);609        try add(dest, store.op);610    }611612    fn indexConstant(self: *Lowerer, dest: *ir.Block, value: u32) Error!*ir.Value {613        const index_type = ArithDialect.getIndexType(self.ctx) catch |err| return loweringError(err);614        const constant = ArithDialect.ConstantOp.createInt(self.ctx, self.loc, index_type, value) catch |err| return loweringError(err);615        try add(dest, constant.op);616        return constant.getResult();617    }618619    fn refuse(self: *Lowerer, op: *ir.Operation) Error {620        if (self.refusal) |refusal| refusal.operation = op.name.name;621        return error.UnsupportedOperation;622    }623};624625/// Operands an op with regions may carry, such as the condition of an `scf.if`.626const max_region_operands = 16;627/// Bytes a quad function's symbol may take.628const max_symbol_bytes = 256;629630const Axis = enum(u1) { x, y };631632const Partner = struct {633    base: *ir.Value = undefined,634    odd: *ir.Value = undefined,635    mapping: *ir.Mapping = undefined,636};637638/// Whether `op` may read or write memory whose contents its operands do not carry: a memref op,639/// or a call that passes or returns anything but scalars. `calls.Plan` admits a helper only when it640/// reads no stage state and touches only its own local arrays, which live for one call, so a641/// partner may call one again on its own scalars.642fn touchesMemory(op: *ir.Operation) bool {643    if (std.mem.eql(u8, op.name.getDialectNamespace(), MemrefDialect.name)) return true;644    if (!std.mem.eql(u8, op.name.name, FuncDialect.CallOp.operation_name)) return false;645    for (op.operand_values) |operand| {646        if (scalarKind(operand.type) == null) return true;647    }648    for (0..op.getNumResults()) |index| {649        const result = op.getResult(index) orelse return true;650        if (scalarKind(result.type) == null) return true;651    }652    return false;653}654655fn isDerivative(name: []const u8) bool {656    return std.mem.eql(u8, name, GpuDialect.DpdxOp.operation_name) or657        std.mem.eql(u8, name, GpuDialect.DpdyOp.operation_name) or658        std.mem.eql(u8, name, GpuDialect.FwidthOp.operation_name);659}660661/// The axes along which some derivative of a fragment body reads each op of the body's block:662/// the op computes the derivative's operand, or a value that operand depends on.663const Needs = struct {664    allocator: std.mem.Allocator,665    axes_of: std.AutoHashMapUnmanaged(*ir.Operation, u2) = .empty,666    /// Every axis some derivative takes.667    axes: u2 = 0,668669    fn init(allocator: std.mem.Allocator) Needs {670        return .{ .allocator = allocator };671    }672673    fn deinit(self: *Needs) void {674        self.axes_of.deinit(self.allocator);675    }676677    fn bit(axis: Axis) u2 {678        return @as(u2, 1) << @backingInt(axis);679    }680681    fn of(self: *const Needs, op: *ir.Operation) u2 {682        return self.axes_of.get(op) orelse 0;683    }684685    /// Walks `block` from its last op to its first, so an op's marks are complete before it686    /// passes them to the ops it reads. Only derivatives in `block` itself seed marks; one in a687    /// region is refused when it is lowered.688    fn collect(self: *Needs, block: *ir.Block) Error!void {689        var cursor = block.operations.tail;690        while (cursor) |pointer| {691            const op: *ir.Operation = @ptrCast(@alignCast(pointer));692            cursor = op.prev_op;693            const name = op.name.name;694            const seed: u2 = if (std.mem.eql(u8, name, GpuDialect.DpdxOp.operation_name))695                bit(.x)696            else if (std.mem.eql(u8, name, GpuDialect.DpdyOp.operation_name))697                bit(.y)698            else if (std.mem.eql(u8, name, GpuDialect.FwidthOp.operation_name))699                bit(.x) | bit(.y)700            else701                0;702            self.axes |= seed;703            const axes = self.of(op) | seed;704            if (axes == 0) continue;705            for (op.operand_values) |operand| try self.mark(operand, block, axes);706            if (op.regions.items.len == 0) continue;707            var reads = Reads{ .needs = self, .block = block, .axes = axes };708            _ = op.walk(.{ .order = .pre_order }, &reads, Reads.visit) catch |err| return loweringError(err);709        }710    }711712    /// Marks the op of `block` that defines `value`, if any, with `axes`.713    fn mark(self: *Needs, value: *ir.Value, block: *ir.Block, axes: u2) Error!void {714        const pointer = value.getDefiningOp() orelse return;715        const op: *ir.Operation = @ptrCast(@alignCast(pointer));716        if (op.parent_block != block) return;717        const entry = self.axes_of.getOrPut(self.allocator, op) catch return error.OutOfMemory;718        if (!entry.found_existing) entry.value_ptr.* = 0;719        entry.value_ptr.* |= axes;720    }721722    /// Marks the values of the outer block that ops inside a region read.723    const Reads = struct {724        needs: *Needs,725        block: *ir.Block,726        axes: u2,727728        fn visit(reads: *Reads, op: *ir.Operation) !ir.Operation.WalkResult {729            for (op.operand_values) |operand| try reads.needs.mark(operand, reads.block, reads.axes);730            return .advance;731        }732    };733};734735/// The smooth and flat words a vertex function's outputs reach, one bit per word.736const Masks = struct {737    smooth: u64 = 0,738    flat: u64 = 0,739740    fn collect(self: *Masks, block: *ir.Block) Error!void {741        var ops = block.getOperations();742        while (ops.next()) |op| {743            if (std.mem.eql(u8, op.name.name, GpuDialect.StageOutputOp.operation_name)) {744                const location = (GpuDialect.StageOutputOp{ .op = op }).getLocation() orelse return error.InvalidArtifact;745                for (op.operand_values, 0..) |operand, component| {746                    const bit = @as(u64, 1) << @intCast(location * layout.location_components + component);747                    switch (scalarKind(operand.type) orelse return error.UnsupportedOperation) {748                        .f32 => self.smooth |= bit,749                        .u32, .i32 => self.flat |= bit,750                        else => return error.UnsupportedOperation,751                    }752                }753            }754            for (op.regions.items) |*region| {755                var nested = region.blocks.head;756                while (nested) |inner| : (nested = inner.next) try self.collect(inner);757            }758        }759    }760};761762fn scalarKind(typ: ir.Type) ?dialects.arith.ScalarKind {763    const type_name = typ.getDialectTypeName() orelse return null;764    return dialects.arith.scalarKindFromTypeName(type_name);765}766767fn mapValue(mapping: *ir.Mapping, from: *ir.Value, to: *ir.Value) Error!void {768    mapping.mapValue(from, to) catch return error.OutOfMemory;769}770771fn add(block: *ir.Block, op: *ir.Operation) Error!void {772    block.addOperation(op) catch |err| return loweringError(err);773}774775fn loweringError(err: anyerror) Error {776    return switch (err) {777        error.OutOfMemory => error.OutOfMemory,778        else => error.CompilationFailed,779    };780}