lib/choir/src/backends/gpu/spirv/emitter/validation.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("../../../../root.zig");
3
4 const ir = choir.ir;
5 const dialects = choir.dialects;
6 const spirv_target = @import("../root.zig");
7 const gpu_target = @import("../../../../dialects/gpu/root.zig");
8 const catalog = @import("catalog.zig");
9
10 const GpuDialect = gpu_target.GpuDialect;
11 const Stage = gpu_target.Stage;
12 const FuncDialect = dialects.func.FuncDialect;
13 const MemrefDialect = dialects.memref.MemrefDialect;
14 const SpirvDialect = spirv_target.SpirvDialect;
15 const BuiltinDialect = dialects.builtin.BuiltinDialect;
16
17 pub fn isFunctionOp(op: *ir.Operation) bool {
18 return std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name) or
19 std.mem.eql(u8, op.name.name, GpuDialect.FuncOp.operation_name);
20 }
21
22 pub fn validateModule(module: *ir.Operation) !void {
23 if (!std.mem.eql(u8, module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
24 return error.InvalidModule;
25 }
26
27 const region = module.getRegion(0) orelse return error.InvalidModule;
28 const block = region.getEntryBlock() orelse return error.InvalidModule;
29
30 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 }
51
52 if (!saw_entry) return error.InvalidModule;
53 }
54
55 /// The ops a stage function takes beyond the gpu stage vocabulary: scalar
56 /// arithmetic and structured control flow. Buffers, shared memory and the
57 /// compute builtins that lower into SPIR-V ops stay in kernels.
58 fn 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) or
64 std.mem.eql(u8, op_name, MemrefDialect.LoadOp.operation_name) or
65 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 }
83
84 fn validateEntryOp(func_op: *ir.Operation, stage: ?Stage) !void {
85 if (func_op.getResultTypes().len != 0) return error.UnsupportedFunctionSignature;
86
87 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;
90
91 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 }
98
99 pub 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 }
106
107 fn 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;
113
114 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 }
128
129 test "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);
132
133 const loc = ir.Location.getUnknown();
134 const module = try BuiltinDialect.ModuleOp.create(&ctx, loc);
135
136 try std.testing.expectError(error.InvalidModule, validateModule(module.op));
137 }
138
139 test "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 }
145
146 test "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);
149
150 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);
154
155 const unsupported = try BuiltinDialect.ModuleOp.create(&ctx, loc);
156 try func.getEntryBlock().addOperation(unsupported.op);
157
158 try std.testing.expectError(error.UnsupportedOperation, validateModule(module.op));
159 }