tiny.choir.backends.gpu.spirv.emitter.stage

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

1//! Vertex and fragment emission. A stage function's interface slots become2//! Input and Output variables with Location decorations, made the first time3//! an op names them, so emission stays one pass over the function. A sampled4//! texture becomes one UniformConstant combined image sampler per group and5//! binding, shared by every stage function in the module.6//!7//! Block reads name a member by byte offset, so each block is declared as an8//! array of 32-bit words and a read indexes the words its member covers. The9//! push-constant block is `{ uint words[n] }` with stride 4. A uniform block10//! is `{ uvec4 rows[m] }` with stride 16, which std140 requires of arrays; an11//! aligned member never crosses a row, so each word is one row and one12//! component. A scan before emission sizes each array to the furthest member13//! the module reads, which the variable's type needs before its first read.1415const std = @import("std");16const choir = @import("../../../../root.zig");17const gpu_target = @import("../../../../dialects/gpu/root.zig");18const gpu = @import("gpu.zig");19const scalar = @import("scalar.zig");20const spec = @import("spec.zig");21const SpirvOp = @import("ops.zig").SpirvOp;2223const ir = choir.ir;24const GpuDialect = gpu_target.GpuDialect;25const Stage = gpu_target.Stage;26const ScalarKind = scalar.Kind;27const max_locations = gpu_target.stage.max_locations;2829pub const Direction = enum { input, output };3031pub const Slot = struct {32    var_id: u32,33    kind: ScalarKind,34    width: u32,35};3637/// The interface of the stage function being emitted.38pub const Interface = struct {39    stage: Stage,40    inputs: [max_locations]?Slot = @splat(null),41    outputs: [max_locations]?Slot = @splat(null),42};4344/// A descriptor's group and binding.45pub const ResourceKey = struct {46    group: u32,47    binding: u32,48};4950pub const Lod = enum { implicit, explicit };5152/// A uniform block: the bytes its reads reach, and its variable once made.53pub const UniformBlock = struct {54    extent: u32,55    var_id: ?u32 = null,56};5758const stage_dialect = gpu_target.stage;59const push_stride = 4;60const uniform_stride = 16;6162pub fn executionModel(stage: Stage) u32 {63    return switch (stage) {64        .vertex => spec.ExecutionModel.Vertex,65        .fragment => spec.ExecutionModel.Fragment,66    };67}6869/// Integer varyings do not interpolate: a vertex writes them and a fragment70/// reads them Flat. Vulkan forbids the decoration on vertex inputs.71fn isFlat(stage: Stage, direction: Direction, kind: ScalarKind) bool {72    if (scalar.isFloat(kind)) return false;73    return switch (stage) {74        .vertex => direction == .output,75        .fragment => direction == .input,76    };77}7879fn interfaceKind(value_type: ir.Type) !ScalarKind {80    const kind = scalar.kindFromType(value_type) orelse return error.UnsupportedType;81    return switch (kind) {82        .f32, .i32, .u32 => kind,83        else => error.UnsupportedType,84    };85}8687fn slotType(self: anytype, kind: ScalarKind, width: u32) !u32 {88    std.debug.assert(width >= 1);89    std.debug.assert(width <= gpu_target.stage.max_components);90    if (width == 1) return self.getScalarType(kind);91    return self.getVectorType(kind, width);92}9394fn interfaceSlot(95    self: anytype,96    direction: Direction,97    location: u32,98    kind: ScalarKind,99    width: u32,100) !Slot {101    const interface = if (self.stage_interface) |*current|102        current103    else104        return error.UnsupportedOperation;105    if (location >= max_locations) return error.UnsupportedOperation;106    const slots = switch (direction) {107        .input => &interface.inputs,108        .output => &interface.outputs,109    };110    if (slots[location]) |slot| {111        if (slot.kind != kind or slot.width != width) return error.UnsupportedType;112        return slot;113    }114    const storage_class = switch (direction) {115        .input => spec.StorageClass.Input,116        .output => spec.StorageClass.Output,117    };118    const ptr_type = try self.getPointerType(storage_class, try slotType(self, kind, width));119    const var_id = self.builder.newId();120    try self.builder.emit(&self.builder.globals, SpirvOp.Variable, &.{121        ptr_type,122        var_id,123        storage_class,124    });125    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{126        var_id,127        spec.Decoration.Location,128        location,129    });130    if (isFlat(interface.stage, direction, kind)) {131        try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{132            var_id,133            spec.Decoration.Flat,134        });135    }136    try self.addInterfaceVar(var_id);137    const slot: Slot = .{ .var_id = var_id, .kind = kind, .width = width };138    slots[location] = slot;139    return slot;140}141142/// Binds each result of `op` to one component of `composite_id`, or the143/// single result to the whole value.144fn bindComponents(self: anytype, op: *ir.Operation, composite_id: u32, kind: ScalarKind) !void {145    const count = op.getNumResults();146    if (count == 1) return self.bindValue(op.getResult(0).?, composite_id);147    const component_type = try self.getScalarType(kind);148    for (0..count) |index| {149        const result_id = self.builder.newId();150        try self.builder.emit(&self.builder.functions, SpirvOp.CompositeExtract, &.{151            component_type,152            result_id,153            composite_id,154            @as(u32, @intCast(index)),155        });156        try self.bindValue(op.getResult(index).?, result_id);157    }158}159160/// Builds the value that `operands` spell, one component each, after checking161/// that every component has `kind`.162fn composeComponents(self: anytype, operands: []const *ir.Value, kind: ScalarKind) !u32 {163    var ids: [gpu_target.stage.max_components]u32 = undefined;164    for (operands, 0..) |operand, index| {165        if (try interfaceKind(operand.type) != kind) return error.UnsupportedType;166        ids[index] = try self.getValue(operand);167    }168    if (operands.len == 1) return ids[0];169    const vector_type = try self.getVectorType(kind, @intCast(operands.len));170    const result_id = self.builder.newId();171    var words: [2 + gpu_target.stage.max_components]u32 = undefined;172    words[0] = vector_type;173    words[1] = result_id;174    @memcpy(words[2..][0..operands.len], ids[0..operands.len]);175    const construct = words[0 .. 2 + operands.len];176    try self.builder.emit(&self.builder.functions, SpirvOp.CompositeConstruct, construct);177    return result_id;178}179180pub fn emitStageInput(self: anytype, op: *ir.Operation) !void {181    const input = GpuDialect.StageInputOp{ .op = op };182    const location = input.getLocation() orelse return error.MissingAttribute;183    const width: u32 = @intCast(op.getNumResults());184    const kind = try interfaceKind(op.getResult(0).?.type);185    for (0..width) |index| {186        if (try interfaceKind(op.getResult(index).?.type) != kind) return error.UnsupportedType;187    }188    const slot = try interfaceSlot(self, .input, location, kind, width);189    const load_id = self.builder.newId();190    try self.builder.emit(&self.builder.functions, SpirvOp.Load, &.{191        try slotType(self, kind, width),192        load_id,193        slot.var_id,194    });195    try bindComponents(self, op, load_id, kind);196}197198pub fn emitStageOutput(self: anytype, op: *ir.Operation) !void {199    const output = GpuDialect.StageOutputOp{ .op = op };200    const location = output.getLocation() orelse return error.MissingAttribute;201    const operands = op.getOperandValues();202    const kind = try interfaceKind(operands[0].type);203    const slot = try interfaceSlot(self, .output, location, kind, @intCast(operands.len));204    const value_id = try composeComponents(self, operands, kind);205    try self.builder.emit(&self.builder.functions, SpirvOp.Store, &.{ slot.var_id, value_id });206}207208pub fn emitPosition(self: anytype, op: *ir.Operation) !void {209    std.debug.assert(op.getNumOperands() == 4);210    const value_id = try composeComponents(self, op.getOperandValues(), .f32);211    const var_id = try gpu.getBuiltinVar(self, .position);212    try self.builder.emit(&self.builder.functions, SpirvOp.Store, &.{ var_id, value_id });213}214215pub fn emitBuiltinLoad(self: anytype, op: *ir.Operation, builtin: gpu.BuiltinKind) !void {216    const var_id = try gpu.getBuiltinVar(self, builtin);217    const load_id = self.builder.newId();218    try self.builder.emit(&self.builder.functions, SpirvOp.Load, &.{219        try gpu.builtinType(self, builtin),220        load_id,221        var_id,222    });223    const kind = scalar.kindFromType(op.getResult(0).?.type) orelse return error.UnsupportedType;224    try bindComponents(self, op, load_id, kind);225}226227fn sampledImageType(self: anytype) !u32 {228    if (self.sampled_image_type) |id| return id;229    const f32_type = try self.getScalarType(.f32);230    const image_type = self.builder.newId();231    try self.builder.emit(&self.builder.types, SpirvOp.TypeImage, &.{232        image_type,233        f32_type,234        spec.Dim.@"2D",235        0,236        0,237        0,238        1,239        spec.ImageFormat.Unknown,240    });241    const id = self.builder.newId();242    try self.builder.emit(&self.builder.types, SpirvOp.TypeSampledImage, &.{ id, image_type });243    self.sampled_image_type = id;244    return id;245}246247/// Binds a sampled texture to the variable for its group and binding. The248/// map gains a key only once its variable exists, so a failed emit leaves no249/// entry naming an id that was never written.250pub fn emitSampledTexture(self: anytype, op: *ir.Operation) !void {251    const texture = GpuDialect.SampledTextureOp{ .op = op };252    const key: ResourceKey = .{253        .group = texture.getGroup() orelse return error.MissingAttribute,254        .binding = texture.getBinding() orelse return error.MissingAttribute,255    };256    if (self.texture_vars.get(key)) |var_id| return self.bindValue(texture.getResult(), var_id);257    const storage_class = spec.StorageClass.UniformConstant;258    const ptr_type = try self.getPointerType(storage_class, try sampledImageType(self));259    const var_id = self.builder.newId();260    try self.builder.emit(&self.builder.globals, SpirvOp.Variable, &.{261        ptr_type,262        var_id,263        storage_class,264    });265    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{266        var_id,267        spec.Decoration.DescriptorSet,268        key.group,269    });270    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{271        var_id,272        spec.Decoration.Binding,273        key.binding,274    });275    try self.texture_vars.put(self.allocator, key, var_id);276    try self.bindValue(texture.getResult(), var_id);277}278279fn floatOperand(self: anytype, value: *ir.Value) !u32 {280    if (scalar.kindFromType(value.type) != .f32) return error.UnsupportedType;281    return self.getValue(value);282}283284pub fn emitSample(self: anytype, op: *ir.Operation, lod: Lod) !void {285    const operands = op.getOperandValues();286    const sampled_image_type = try sampledImageType(self);287    const sampled_image = self.builder.newId();288    try self.builder.emit(&self.builder.functions, SpirvOp.Load, &.{289        sampled_image_type,290        sampled_image,291        try self.getValue(operands[0]),292    });293    const coord = self.builder.newId();294    try self.builder.emit(&self.builder.functions, SpirvOp.CompositeConstruct, &.{295        try self.getVectorType(.f32, 2),296        coord,297        try floatOperand(self, operands[1]),298        try floatOperand(self, operands[2]),299    });300    const texel_type = try self.getVectorType(.f32, 4);301    const texel = self.builder.newId();302    const functions = &self.builder.functions;303    switch (lod) {304        .implicit => try self.builder.emit(functions, SpirvOp.ImageSampleImplicitLod, &.{305            texel_type,306            texel,307            sampled_image,308            coord,309        }),310        .explicit => try self.builder.emit(functions, SpirvOp.ImageSampleExplicitLod, &.{311            texel_type,312            texel,313            sampled_image,314            coord,315            spec.ImageOperands.Lod,316            try floatOperand(self, operands[3]),317        }),318    }319    try bindComponents(self, op, texel, .f32);320}321322/// Emits a fine derivative. A plain OpDPdx lets the device pick the flavour: lavapipe, the RTX 4090323/// and Metal take fine, and RADV coarse, so only the explicit op draws one frame on all of them.324pub fn emitDerivative(self: anytype, op: *ir.Operation, opcode: u16) !void {325    try self.requireCapability(spec.Capability.DerivativeControl);326    const value_id = try floatOperand(self, op.getOperandValues()[0]);327    const result_id = self.builder.newId();328    try self.builder.emit(&self.builder.functions, opcode, &.{329        try self.getScalarType(.f32),330        result_id,331        value_id,332    });333    try self.bindValue(op.getResult(0).?, result_id);334}335336/// Records how far into each block the module reads, refusing a member the337/// block cannot place.338pub fn scanBlocks(self: anytype, module: *ir.Operation) !void {339    self.push_extent = 0;340    self.push_var = null;341    self.uniform_blocks.clearRetainingCapacity();342    const Scan = struct {343        codegen: @TypeOf(self),344        refused: bool = false,345346        fn visit(scan: *@This(), op: *ir.Operation) !ir.Operation.WalkResult {347            if (std.mem.eql(u8, op.name.name, GpuDialect.PushConstantOp.operation_name)) {348                const member = (GpuDialect.PushConstantOp{ .op = op }).member() orelse349                    return scan.refuse();350                if (!member.placed(stage_dialect.max_push_constant_bytes)) return scan.refuse();351                scan.codegen.push_extent = @max(scan.codegen.push_extent, member.end());352            } else if (std.mem.eql(u8, op.name.name, GpuDialect.UniformOp.operation_name)) {353                const uniform = GpuDialect.UniformOp{ .op = op };354                const member = uniform.member() orelse return scan.refuse();355                if (!member.placed(stage_dialect.max_uniform_bytes)) return scan.refuse();356                const key: ResourceKey = .{357                    .group = uniform.getGroup() orelse return scan.refuse(),358                    .binding = uniform.getBinding() orelse return scan.refuse(),359                };360                const entry = try scan.codegen.uniform_blocks.getOrPut(scan.codegen.allocator, key);361                if (!entry.found_existing) entry.value_ptr.* = .{ .extent = 0 };362                entry.value_ptr.extent = @max(entry.value_ptr.extent, member.end());363            }364            return .advance;365        }366367        fn refuse(scan: *@This()) ir.Operation.WalkResult {368            scan.refused = true;369            return .interrupt;370        }371    };372    var scan: Scan = .{ .codegen = self };373    _ = module.walk(.{ .order = .pre_order }, &scan, Scan.visit) catch |err| switch (err) {374        error.OutOfMemory => return error.OutOfMemory,375        else => unreachable,376    };377    if (scan.refused) return error.UnsupportedOperation;378}379380/// Declares a Block struct whose one member is an array of `length`381/// elements of `element_type` at `stride` bytes, and a variable of it in382/// `storage_class`. The array type is not shared with the cache, since its383/// ArrayStride belongs only in explicitly laid out storage.384fn blockVariable(385    self: anytype,386    storage_class: u32,387    element_type: u32,388    length: u32,389    stride: u32,390) !u32 {391    std.debug.assert(length >= 1);392    const length_id = try self.getIntConstant(try self.getScalarType(.u32), .u32, length);393    const array_type = self.builder.newId();394    try self.builder.emit(&self.builder.types, SpirvOp.TypeArray, &.{395        array_type,396        element_type,397        length_id,398    });399    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{400        array_type,401        spec.Decoration.ArrayStride,402        stride,403    });404    const struct_type = self.builder.newId();405    try self.builder.emit(&self.builder.types, SpirvOp.TypeStruct, &.{ struct_type, array_type });406    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{407        struct_type,408        spec.Decoration.Block,409    });410    try self.builder.emit(&self.builder.annotations, SpirvOp.MemberDecorate, &.{411        struct_type,412        0,413        spec.Decoration.Offset,414        0,415    });416    const var_id = self.builder.newId();417    try self.builder.emit(&self.builder.globals, SpirvOp.Variable, &.{418        try self.getPointerType(storage_class, struct_type),419        var_id,420        storage_class,421    });422    return var_id;423}424425fn pushVariable(self: anytype) !u32 {426    if (self.push_var) |id| return id;427    std.debug.assert(self.push_extent > 0);428    std.debug.assert(self.push_extent % push_stride == 0);429    const u32_type = try self.getScalarType(.u32);430    const id = try blockVariable(431        self,432        spec.StorageClass.PushConstant,433        u32_type,434        self.push_extent / push_stride,435        push_stride,436    );437    self.push_var = id;438    return id;439}440441fn uniformVariable(self: anytype, key: ResourceKey) !u32 {442    const block = self.uniform_blocks.getPtr(key) orelse return error.UnsupportedOperation;443    if (block.var_id) |id| return id;444    std.debug.assert(block.extent > 0);445    const row_type = try self.getVectorType(.u32, 4);446    const rows = std.math.divCeil(u32, block.extent, uniform_stride) catch unreachable;447    const id = try blockVariable(self, spec.StorageClass.Uniform, row_type, rows, uniform_stride);448    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{449        id,450        spec.Decoration.DescriptorSet,451        key.group,452    });453    try self.builder.emit(&self.builder.annotations, SpirvOp.Decorate, &.{454        id,455        spec.Decoration.Binding,456        key.binding,457    });458    block.var_id = id;459    return id;460}461462/// The kind every result of a block read shares.463fn blockReadKind(op: *ir.Operation) !ScalarKind {464    const kind = try interfaceKind(op.getResult(0).?.type);465    for (0..op.getNumResults()) |index| {466        if (try interfaceKind(op.getResult(index).?.type) != kind) return error.UnsupportedType;467    }468    return kind;469}470471/// Loads the word that `indices` reach from `var_id` and binds it, as472/// `kind`, to `result`.473fn loadBlockWord(474    self: anytype,475    result: *ir.Value,476    kind: ScalarKind,477    storage_class: u32,478    var_id: u32,479    indices: []const u32,480) !void {481    const u32_type = try self.getScalarType(.u32);482    var words: [3 + 3]u32 = undefined;483    std.debug.assert(indices.len <= 3);484    const pointer = self.builder.newId();485    words[0] = try self.getPointerType(storage_class, u32_type);486    words[1] = pointer;487    words[2] = var_id;488    for (indices, 0..) |index, slot| {489        words[3 + slot] = try self.getIntConstant(u32_type, .u32, index);490    }491    try self.builder.emit(&self.builder.functions, SpirvOp.AccessChain, words[0 .. 3 + indices.len]);492    const bits = self.builder.newId();493    try self.builder.emit(&self.builder.functions, SpirvOp.Load, &.{ u32_type, bits, pointer });494    if (kind == .u32) return self.bindValue(result, bits);495    const value = self.builder.newId();496    try self.builder.emit(&self.builder.functions, SpirvOp.Bitcast, &.{497        try self.getScalarType(kind),498        value,499        bits,500    });501    try self.bindValue(result, value);502}503504pub fn emitPushConstant(self: anytype, op: *ir.Operation) !void {505    const member = (GpuDialect.PushConstantOp{ .op = op }).member() orelse506        return error.MissingAttribute;507    std.debug.assert(member.end() <= self.push_extent);508    const kind = try blockReadKind(op);509    const var_id = try pushVariable(self);510    const first = member.offset / push_stride;511    for (0..member.width) |component| {512        const word: u32 = first + @as(u32, @intCast(component));513        try loadBlockWord(self, op.getResult(component).?, kind, spec.StorageClass.PushConstant, var_id, &.{514            0,515            word,516        });517    }518}519520pub fn emitUniform(self: anytype, op: *ir.Operation) !void {521    const uniform = GpuDialect.UniformOp{ .op = op };522    const member = uniform.member() orelse return error.MissingAttribute;523    const key: ResourceKey = .{524        .group = uniform.getGroup() orelse return error.MissingAttribute,525        .binding = uniform.getBinding() orelse return error.MissingAttribute,526    };527    const kind = try blockReadKind(op);528    const var_id = try uniformVariable(self, key);529    const row = member.offset / uniform_stride;530    const first = member.offset % uniform_stride / 4;531    std.debug.assert(first + member.width <= 4);532    for (0..member.width) |component| {533        const column: u32 = first + @as(u32, @intCast(component));534        try loadBlockWord(self, op.getResult(component).?, kind, spec.StorageClass.Uniform, var_id, &.{535            0,536            row,537            column,538        });539    }540}541542test "integer varyings are flat only where Vulkan allows the decoration" {543    try std.testing.expect(isFlat(.vertex, .output, .u32));544    try std.testing.expect(isFlat(.fragment, .input, .i32));545    try std.testing.expect(!isFlat(.vertex, .input, .u32));546    try std.testing.expect(!isFlat(.fragment, .output, .u32));547    try std.testing.expect(!isFlat(.fragment, .input, .f32));548}