tiny.choir.backends.gpu.spirv.emitter.dialect

1const std = @import("std");2const choir = @import("../../../../root.zig");34const ir = choir.ir;5const dialects = choir.dialects;6const spirv_target = @import("../root.zig");7const gpu = @import("gpu.zig");8const binary = @import("module.zig");9const scalar = @import("scalar.zig");10const spec = @import("spec.zig");11const spirv_ops = @import("ops.zig");1213const FuncDialect = dialects.func.FuncDialect;14const SpirvDialect = spirv_target.SpirvDialect;15const SpirvExecutionModel = spirv_target.ExecutionModel;16const SpirvStorageClass = spirv_target.StorageClass;17const Section = binary.Section;18const SpirvOp = spirv_ops.SpirvOp;1920pub fn functionName(op: *ir.Operation) ?[]const u8 {21    return ir.SymbolTable.getSymbolName(op);22}2324fn storageClassToSpv(storage: SpirvStorageClass) u32 {25    return switch (storage) {26        .function => spec.StorageClass.Function,27        .private => spec.StorageClass.Private,28        .workgroup => spec.StorageClass.Workgroup,29        .uniform => spec.StorageClass.Uniform,30        .storage_buffer => spec.StorageClass.StorageBuffer,31        .input => spec.StorageClass.Input,32        .output => spec.StorageClass.Output,33    };34}3536pub fn emitModuleWords(self: anytype, module: *ir.Operation) ![]u32 {37    if (!std.mem.eql(u8, module.name.name, SpirvDialect.ModuleOp.operation_name)) {38        return error.InvalidModule;39    }4041    const spirv_module = SpirvDialect.ModuleOp{ .op = module };42    const addressing_model = spirv_module.getAddressingModel() orelse return error.MissingAttribute;43    const memory_model = spirv_module.getMemoryModel() orelse return error.MissingAttribute;44    const capability = spirv_module.getCapability() orelse return error.MissingAttribute;4546    try self.requireCapability(switch (capability) {47        .shader => spec.Capability.Shader,48    });49    try self.builder.emitMemoryModel(50        switch (addressing_model) {51            .logical => spec.AddressingModel.Logical,52        },53        switch (memory_model) {54            .glsl450 => spec.MemoryModel.GLSL450,55        },56    );5758    const region = module.getRegion(0) orelse return error.InvalidModule;59    const block = region.getEntryBlock() orelse return error.InvalidModule;6061    var funcs: std.ArrayListUnmanaged(*ir.Operation) = .empty;62    var vars: std.ArrayListUnmanaged(*ir.Operation) = .empty;63    var consts: std.ArrayListUnmanaged(*ir.Operation) = .empty;64    defer funcs.deinit(self.allocator);65    defer vars.deinit(self.allocator);66    defer consts.deinit(self.allocator);6768    var op_iter = block.operations.head;69    while (op_iter) |op_ptr| {70        const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));71        const name = op.name.name;72        if (std.mem.eql(u8, name, SpirvDialect.FuncOp.operation_name)) {73            try funcs.append(self.allocator, op);74        } else if (std.mem.eql(u8, name, SpirvDialect.VariableOp.operation_name)) {75            try vars.append(self.allocator, op);76        } else if (std.mem.eql(u8, name, SpirvDialect.ConstantOp.operation_name)) {77            try consts.append(self.allocator, op);78        } else {79            return error.UnsupportedOperation;80        }81        op_iter = op.next_op;82    }8384    for (consts.items) |const_op| {85        try emitConstant(self, const_op);86    }8788    for (vars.items) |var_op| {89        try emitVariable(self, var_op, true);90    }9192    for (funcs.items) |func_op| {93        try emitFunction(self, func_op);94    }9596    return self.builder.toWords(self.allocator);97}9899fn emitTopLevel(self: anytype, op: *ir.Operation) !void {100    const name = op.name.name;101    if (std.mem.eql(u8, name, SpirvDialect.FuncOp.operation_name)) {102        return emitFunction(self, op);103    }104    if (std.mem.eql(u8, name, SpirvDialect.VariableOp.operation_name)) {105        return emitVariable(self, op, true);106    }107    if (std.mem.eql(u8, name, SpirvDialect.ConstantOp.operation_name)) {108        return emitConstant(self, op);109    }110    return error.UnsupportedOperation;111}112113fn emitFunction(self: anytype, func_op: *ir.Operation) !void {114    const name = functionName(func_op) orelse return error.MissingFunctionName;115116    const region = func_op.getRegion(0) orelse return error.InvalidModule;117    if (!region.hasOneBlock()) return error.UnsupportedControlFlow;118    const entry = region.getEntryBlock() orelse return error.InvalidModule;119120    const result_types = func_op.getResultTypes();121    if (result_types.len > 1) return error.UnsupportedFunctionSignature;122123    const returns_void = result_types.len == 0;124    const return_kind: scalar.Kind = if (returns_void)125        .void126    else127        (scalar.kindFromType(result_types[0]) orelse return error.UnsupportedType);128    const return_type_id = try self.getScalarType(return_kind);129130    var param_type_ids = std.ArrayListUnmanaged(u32).empty;131    defer param_type_ids.deinit(self.allocator);132    for (entry.arguments.items) |arg| {133        const kind = scalar.kindFromType(arg.type) orelse return error.UnsupportedType;134        const type_id = try self.getScalarType(kind);135        try param_type_ids.append(self.allocator, type_id);136    }137138    const func_type_id = try emitFunctionType(self, return_type_id, param_type_ids.items);139    const func_id = self.builder.newId();140    try self.builder.emit(&self.builder.functions, SpirvOp.Function, &.{141        return_type_id,142        func_id,143        0,144        func_type_id,145    });146147    for (entry.arguments.items, 0..) |arg, idx| {148        const param_id = self.builder.newId();149        try self.builder.emit(&self.builder.functions, SpirvOp.FunctionParameter, &.{150            param_type_ids.items[idx],151            param_id,152        });153        try self.bindValue(arg, param_id);154    }155156    const label_id = self.builder.newId();157    try self.emitLabel(label_id);158159    var op_iter = entry.operations.head;160    while (op_iter) |op_ptr| {161        const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));162        try emitOp(self, op, returns_void);163        op_iter = op.next_op;164    }165166    try self.builder.emit(&self.builder.functions, SpirvOp.FunctionEnd, &.{});167    self.current_block = null;168169    const spirv_func = SpirvDialect.FuncOp{ .op = func_op };170    if (spirv_func.isEntryPoint()) {171        const exec_model = spirv_func.getExecutionModel() orelse return error.MissingAttribute;172        try self.builder.emitEntryPoint(173            switch (exec_model) {174                SpirvExecutionModel.vertex => spec.ExecutionModel.Vertex,175                SpirvExecutionModel.fragment => spec.ExecutionModel.Fragment,176                SpirvExecutionModel.gl_compute => spec.ExecutionModel.GLCompute,177            },178            func_id,179            name,180            self.interface_vars.items,181        );182        switch (exec_model) {183            SpirvExecutionModel.vertex => {},184            SpirvExecutionModel.fragment => try self.builder.emitExecutionMode(185                func_id,186                spec.ExecutionMode.OriginUpperLeft,187            ),188            SpirvExecutionModel.gl_compute => {189                try self.builder.emitExecutionModeLocalSize(func_id, 1, 1, 1);190            },191        }192        try self.emitFloatExecutionModes(func_id);193    }194}195196fn emitOp(self: anytype, op: *ir.Operation, returns_void: bool) !void {197    const name = op.name.name;198    if (std.mem.eql(u8, name, SpirvDialect.ConstantOp.operation_name)) {199        return emitConstant(self, op);200    }201    if (std.mem.eql(u8, name, SpirvDialect.VariableOp.operation_name)) {202        return emitVariable(self, op, false);203    }204    if (std.mem.eql(u8, name, SpirvDialect.IAddOp.operation_name)) {205        return emitBinary(self, op, SpirvOp.IAdd, .int_any);206    }207    if (std.mem.eql(u8, name, SpirvDialect.FAddOp.operation_name)) {208        return emitBinary(self, op, SpirvOp.FAdd, .float);209    }210    if (std.mem.eql(u8, name, SpirvDialect.ISubOp.operation_name)) {211        return emitBinary(self, op, SpirvOp.ISub, .int_any);212    }213    if (std.mem.eql(u8, name, SpirvDialect.FSubOp.operation_name)) {214        return emitBinary(self, op, SpirvOp.FSub, .float);215    }216    if (std.mem.eql(u8, name, SpirvDialect.IMulOp.operation_name)) {217        return emitBinary(self, op, SpirvOp.IMul, .int_any);218    }219    if (std.mem.eql(u8, name, SpirvDialect.FMulOp.operation_name)) {220        return emitBinary(self, op, SpirvOp.FMul, .float);221    }222    if (std.mem.eql(u8, name, SpirvDialect.UDivOp.operation_name)) {223        return emitBinary(self, op, SpirvOp.UDiv, .int_unsigned);224    }225    if (std.mem.eql(u8, name, SpirvDialect.SDivOp.operation_name)) {226        return emitBinary(self, op, SpirvOp.SDiv, .int_signed);227    }228    if (std.mem.eql(u8, name, SpirvDialect.FDivOp.operation_name)) {229        return emitBinary(self, op, SpirvOp.FDiv, .float);230    }231    if (std.mem.eql(u8, name, SpirvDialect.LocalInvocationIdOp.operation_name)) {232        return gpu.emitIndex(self, op, .local_invocation_id);233    }234    if (std.mem.eql(u8, name, SpirvDialect.WorkgroupIdOp.operation_name)) {235        return gpu.emitIndex(self, op, .workgroup_id);236    }237    if (std.mem.eql(u8, name, SpirvDialect.WorkgroupSizeOp.operation_name)) {238        return gpu.emitIndex(self, op, .workgroup_size);239    }240    if (std.mem.eql(u8, name, SpirvDialect.NumWorkgroupsOp.operation_name)) {241        return gpu.emitIndex(self, op, .num_workgroups);242    }243    if (std.mem.eql(u8, name, SpirvDialect.GlobalInvocationIdOp.operation_name)) {244        return gpu.emitIndex(self, op, .global_invocation_id);245    }246    if (std.mem.eql(u8, name, SpirvDialect.BarrierOp.operation_name)) {247        return gpu.emitBarrier(self, op);248    }249    if (std.mem.eql(u8, name, SpirvDialect.SyncWarpOp.operation_name)) {250        return gpu.emitSyncWarp(self, op);251    }252    if (std.mem.eql(u8, name, SpirvDialect.ActiveMaskOp.operation_name)) {253        return gpu.emitActiveMask(self, op);254    }255    if (std.mem.eql(u8, name, SpirvDialect.AllSyncOp.operation_name)) {256        return gpu.emitAllAny(self, op, .all);257    }258    if (std.mem.eql(u8, name, SpirvDialect.AnySyncOp.operation_name)) {259        return gpu.emitAllAny(self, op, .any);260    }261    if (std.mem.eql(u8, name, SpirvDialect.BallotSyncOp.operation_name)) {262        return gpu.emitBallot(self, op);263    }264    if (std.mem.eql(u8, name, SpirvDialect.ShflSyncOp.operation_name)) {265        return gpu.emitShuffle(self, op);266    }267    if (std.mem.eql(u8, name, SpirvDialect.WarpReduceOp.operation_name)) {268        return gpu.emitWarpReduce(self, op);269    }270    if (std.mem.eql(u8, name, SpirvDialect.WarpScanOp.operation_name)) {271        return gpu.emitWarpScan(self, op);272    }273    if (std.mem.eql(u8, name, FuncDialect.ReturnOp.operation_name)) {274        return emitReturn(self, op, returns_void);275    }276    return error.UnsupportedOperation;277}278279pub fn emitConstant(self: anytype, op: *ir.Operation) !void {280    const result = op.getResult(0) orelse return error.UnsupportedOperation;281    const type_id = try self.getTypeForValue(result);282283    const constant = SpirvDialect.ConstantOp{ .op = op };284    if (constant.getIntValue()) |int_value| {285        const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;286        const const_id = try self.getIntConstant(type_id, kind, int_value);287        try self.bindValue(result, const_id);288        return;289    }290291    if (constant.getFloatValue()) |float_value| {292        const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;293        const const_id = try self.getFloatConstant(type_id, kind, float_value);294        try self.bindValue(result, const_id);295        return;296    }297298    if (op.getAttrAs(ir.Attribute.BoolAttr, "value")) |bool_attr| {299        const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;300        if (kind != .bool) return error.UnsupportedType;301        const const_id = try self.getBoolConstant(bool_attr.getValue());302        try self.bindValue(result, const_id);303        return;304    }305306    return error.UnsupportedOperation;307}308309pub const SpirvBinaryExpectation = enum {310    int_any,311    int_signed,312    int_unsigned,313    float,314};315316pub fn emitBinary(317    self: anytype,318    op: *ir.Operation,319    opcode: u16,320    expectation: SpirvBinaryExpectation,321) !void {322    if (op.operands.items.len != 2) return error.UnsupportedOperation;323    const lhs = op.operands.items[0].value;324    const rhs = op.operands.items[1].value;325    const result = op.getResult(0) orelse return error.UnsupportedOperation;326327    const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;328    switch (expectation) {329        .float => if (!scalar.isFloat(kind)) return error.UnsupportedType,330        .int_any => if (!scalar.isInt(kind)) return error.UnsupportedType,331        .int_signed => if (!scalar.isSignedInt(kind)) return error.UnsupportedType,332        .int_unsigned => if (!scalar.isUnsignedInt(kind)) return error.UnsupportedType,333    }334335    const lhs_id = try self.getValue(lhs);336    const rhs_id = try self.getValue(rhs);337    const result_type_id = try self.getScalarType(kind);338339    const result_id = self.builder.newId();340    try self.builder.emit(&self.builder.functions, opcode, &.{341        result_type_id,342        result_id,343        lhs_id,344        rhs_id,345    });346    try self.bindValue(result, result_id);347}348349fn emitVariable(350    self: anytype,351    op: *ir.Operation,352    module_scope: bool,353) !void {354    const var_op = SpirvDialect.VariableOp{ .op = op };355    const result = op.getResult(0) orelse return error.UnsupportedOperation;356    const storage = var_op.getStorageClass() orelse return error.MissingAttribute;357    const storage_class = storageClassToSpv(storage);358359    const elem_kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;360    const elem_type = try self.getScalarType(elem_kind);361    const ptr_type = try self.getPointerType(storage_class, elem_type);362363    const var_id = self.builder.newId();364    const initializer = var_op.getInitializer();365366    const section = if (module_scope) &self.builder.globals else &self.builder.functions;367    if (initializer) |init_val| {368        const init_id = try self.getValue(init_val);369        try self.builder.emit(section, SpirvOp.Variable, &.{370            ptr_type,371            var_id,372            storage_class,373            init_id,374        });375    } else {376        try self.builder.emit(section, SpirvOp.Variable, &.{377            ptr_type,378            var_id,379            storage_class,380        });381    }382383    try self.bindValue(result, var_id);384    if (module_scope) {385        try self.addInterfaceVar(var_id);386    }387}388389fn emitReturn(self: anytype, op: *ir.Operation, returns_void: bool) !void {390    if (returns_void) {391        if (op.operands.items.len != 0) return error.UnsupportedFunctionSignature;392        try self.builder.emit(&self.builder.functions, SpirvOp.Return, &.{});393        return;394    }395396    if (op.operands.items.len != 1) return error.UnsupportedFunctionSignature;397    const value_id = try self.getValue(op.operands.items[0].value);398    try self.builder.emit(&self.builder.functions, SpirvOp.ReturnValue, &.{value_id});399}400401/// The function type of a signature, declared once per module as SPIR-V402/// requires of a non-aggregate type.403pub fn emitFunctionType(404    self: anytype,405    return_type: u32,406    param_types: []const u32,407) !u32 {408    const types = self.builder.types.items;409    if (findFunctionType(types, return_type, param_types)) |existing| return existing;410    const id = self.builder.newId();411    var operands = Section.empty;412    defer operands.deinit(self.allocator);413414    try operands.append(self.allocator, id);415    try operands.append(self.allocator, return_type);416    if (param_types.len > 0) {417        try operands.appendSlice(self.allocator, param_types);418    }419420    try self.builder.emit(&self.builder.types, SpirvOp.TypeFunction, operands.items);421    return id;422}423424fn findFunctionType(types: []const u32, return_type: u32, param_types: []const u32) ?u32 {425    var index: usize = 0;426    while (index < types.len) {427        const count = types[index] >> 16;428        std.debug.assert(count > 0);429        const opcode: u16 = @truncate(types[index]);430        if (opcode == SpirvOp.TypeFunction and count == 3 + param_types.len and431            types[index + 2] == return_type and432            std.mem.eql(u32, types[index + 3 ..][0..param_types.len], param_types))433        {434            return types[index + 1];435        }436        index += count;437    }438    return null;439}440441test "spirv dialect writer maps storage classes" {442    try std.testing.expectEqual(@as(u32, spec.StorageClass.Workgroup), storageClassToSpv(.workgroup));443    try std.testing.expectEqual(@as(u32, spec.StorageClass.StorageBuffer), storageClassToSpv(.storage_buffer));444}

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