tiny.choir.backends.gpu.spirv.emitter.validation

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

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_target = @import("../../../../dialects/gpu/root.zig");8const catalog = @import("catalog.zig");910const GpuDialect = gpu_target.GpuDialect;11const Stage = gpu_target.Stage;12const FuncDialect = dialects.func.FuncDialect;13const MemrefDialect = dialects.memref.MemrefDialect;14const SpirvDialect = spirv_target.SpirvDialect;15const BuiltinDialect = dialects.builtin.BuiltinDialect;1617pub fn isFunctionOp(op: *ir.Operation) bool {18    return std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name) or19        std.mem.eql(u8, op.name.name, GpuDialect.FuncOp.operation_name);20}2122pub fn validateModule(module: *ir.Operation) !void {23    if (!std.mem.eql(u8, module.name.name, BuiltinDialect.ModuleOp.operation_name)) {24        return error.InvalidModule;25    }2627    const region = module.getRegion(0) orelse return error.InvalidModule;28    const block = region.getEntryBlock() orelse return error.InvalidModule;2930    var saw_entry = false;31    var op_iter = block.operations.head;32    while (op_iter) |op_ptr| {33        const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));34        if (isFunctionOp(op)) {35            const stage = gpu_target.stage.stageOf(op);36            if (op.getAttr("kernel") != null) {37                if (stage != null) return error.InvalidModule;38                saw_entry = true;39                try validateEntryOp(op, null);40            } else if (stage != null) {41                saw_entry = true;42                try validateEntryOp(op, stage);43            } else if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) {44                return error.UnsupportedOperation;45            }46        } else {47            return error.UnsupportedOperation;48        }49        op_iter = op.next_op;50    }5152    if (!saw_entry) return error.InvalidModule;53}5455/// The ops a stage function takes beyond the gpu stage vocabulary: scalar56/// arithmetic and structured control flow. Buffers, shared memory and the57/// compute builtins that lower into SPIR-V ops stay in kernels.58fn stageAdmits(op_name: []const u8) bool {59    const prefixes = [_][]const u8{ "arith.", "scf.", "func.", "gpu." };60    for (prefixes) |prefix| {61        if (std.mem.startsWith(u8, op_name, prefix)) return true;62    }63    if (std.mem.eql(u8, op_name, MemrefDialect.AllocaOp.operation_name) or64        std.mem.eql(u8, op_name, MemrefDialect.LoadOp.operation_name) or65        std.mem.eql(u8, op_name, MemrefDialect.StoreOp.operation_name)) return true;66    const lowered = [_][]const u8{67        SpirvDialect.ConstantOp.operation_name,68        SpirvDialect.IAddOp.operation_name,69        SpirvDialect.FAddOp.operation_name,70        SpirvDialect.ISubOp.operation_name,71        SpirvDialect.FSubOp.operation_name,72        SpirvDialect.IMulOp.operation_name,73        SpirvDialect.FMulOp.operation_name,74        SpirvDialect.UDivOp.operation_name,75        SpirvDialect.SDivOp.operation_name,76        SpirvDialect.FDivOp.operation_name,77    };78    for (lowered) |name| {79        if (std.mem.eql(u8, op_name, name)) return true;80    }81    return false;82}8384fn validateEntryOp(func_op: *ir.Operation, stage: ?Stage) !void {85    if (func_op.getResultTypes().len != 0) return error.UnsupportedFunctionSignature;8687    const region = func_op.getRegion(0) orelse return error.InvalidModule;88    if (!region.hasOneBlock()) return error.UnsupportedControlFlow;89    const entry = region.getEntryBlock() orelse return error.InvalidModule;9091    var op_iter = entry.operations.head;92    while (op_iter) |op_ptr| {93        const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));94        try validateOperationTree(op, stage, false);95        op_iter = op.next_op;96    }97}9899pub fn validateHelperOp(func_op: *ir.Operation) !void {100    if (func_op.getResultTypes().len > 1) return error.UnsupportedFunctionSignature;101    const region = func_op.getRegion(0) orelse return error.InvalidModule;102    if (!region.hasOneBlock()) return error.UnsupportedControlFlow;103    var ops = region.getEntryBlock().?.getOperations();104    while (ops.next()) |op| try validateOperationTree(op, null, true);105}106107fn validateOperationTree(op: *ir.Operation, stage: ?Stage, helper: bool) !void {108    if (!catalog.supports(op.name.name)) return error.UnsupportedOperation;109    if (!gpu_target.stage.admits(stage, op.name.name)) return error.UnsupportedOperation;110    if (stage != null and !stageAdmits(op.name.name)) return error.UnsupportedOperation;111    if (helper and !stageAdmits(op.name.name)) return error.UnsupportedOperation;112    if (helper and std.mem.startsWith(u8, op.name.name, "gpu.")) return error.UnsupportedOperation;113114    var region_index: usize = 0;115    while (region_index < op.getNumRegions()) : (region_index += 1) {116        const region = op.getRegion(region_index) orelse continue;117        var block_iter = region.getBlocks();118        while (block_iter.next()) |block| {119            var child_iter = block.operations.head;120            while (child_iter) |child_ptr| {121                const child: *ir.Operation = @ptrCast(@alignCast(child_ptr));122                try validateOperationTree(child, stage, helper);123                child_iter = child.next_op;124            }125        }126    }127}128129test "spirv validation rejects module without kernel" {130    var ctx = try ir.Context.init(std.testing.allocator, ir.Context.Limits.testing);131    defer ctx.deinit(std.testing.allocator);132133    const loc = ir.Location.getUnknown();134    const module = try BuiltinDialect.ModuleOp.create(&ctx, loc);135136    try std.testing.expectError(error.InvalidModule, validateModule(module.op));137}138139test "spirv kernel catalog excludes module and function boundaries" {140    try std.testing.expect(!catalog.supports(BuiltinDialect.ModuleOp.operation_name));141    try std.testing.expect(!catalog.supports(FuncDialect.FuncOp.operation_name));142    try std.testing.expect(!catalog.supports(GpuDialect.FuncOp.operation_name));143    try std.testing.expect(!catalog.supports(SpirvDialect.ModuleOp.operation_name));144}145146test "spirv validation rejects a noncatalog boundary inside a kernel" {147    var ctx = try ir.Context.init(std.testing.allocator, ir.Context.Limits.testing);148    defer ctx.deinit(std.testing.allocator);149150    const loc = ir.Location.getUnknown();151    const module = try BuiltinDialect.ModuleOp.create(&ctx, loc);152    const func = try FuncDialect.FuncOp.createKernel(&ctx, loc, "kernel", &.{});153    try module.getBodyBlock().addOperation(func.op);154155    const unsupported = try BuiltinDialect.ModuleOp.create(&ctx, loc);156    try func.getEntryBlock().addOperation(unsupported.op);157158    try std.testing.expectError(error.UnsupportedOperation, validateModule(module.op));159}