lib/choir/src/backends/gpu/calls.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 //! Reachable ordinary functions in callee-first order for GPU stage backends.
2 //!
3 //! Every stage backend plans its module here before it checks anything of its own, so a module
4 //! all of them refuse is refused here by one name. A stage body, and every helper a stage reaches,
5 //! holds only the memory `gpu.stage.admitsMemory` admits. A helper only kernels reach keeps the
6 //! kernel vocabulary.
7
8 const std = @import("std");
9 const choir = @import("../../root.zig");
10 const gpu = @import("../../dialects/gpu/root.zig");
11
12 const ir = choir.ir;
13 const FuncDialect = choir.dialects.FuncDialect;
14
15 const Mark = enum { visiting, done };
16 const max_call_depth = 64;
17 pub const SignatureError = error{UnsupportedHelperSignature};
18 const PlanError = error{ OutOfMemory, UnsupportedOperation } || gpu.stage.MemoryError || SignatureError;
19
20 pub const Plan = struct {
21 allocator: std.mem.Allocator,
22 functions: std.StringHashMapUnmanaged(*ir.Operation) = .{},
23 marks: std.AutoHashMapUnmanaged(*ir.Operation, Mark) = .{},
24 helpers: std.ArrayListUnmanaged(*ir.Operation) = .empty,
25 /// Receives the name of the op a stage may not hold, when init refuses one.
26 refused: ?*[]const u8 = null,
27
28 /// Plans the stages first, so a helper both stages and kernels reach meets the stage rules.
29 pub fn init(allocator: std.mem.Allocator, module: *ir.Operation, refused: ?*[]const u8) PlanError!Plan {
30 var plan = Plan{ .allocator = allocator, .refused = refused };
31 errdefer plan.deinit();
32 const region = module.getRegion(0) orelse return error.UnsupportedOperation;
33 const body = region.getEntryBlock() orelse return error.UnsupportedOperation;
34 var ops = body.getOperations();
35 while (ops.next()) |op| {
36 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;
37 const name = (FuncDialect.FuncOp{ .op = op }).getName() orelse return error.UnsupportedOperation;
38 const entry = try plan.functions.getOrPut(allocator, name);
39 if (entry.found_existing) return error.UnsupportedOperation;
40 entry.value_ptr.* = op;
41 }
42 for ([_]bool{ true, false }) |stage| {
43 ops = body.getOperations();
44 while (ops.next()) |op| {
45 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;
46 const entry = if (stage) gpu.stage.stageOf(op) != null else op.getAttr("kernel") != null;
47 if (!entry) continue;
48 try plan.walkCalls((FuncDialect.FuncOp{ .op = op }).getEntryBlock(), 0, stage);
49 }
50 }
51 return plan;
52 }
53
54 pub fn deinit(self: *Plan) void {
55 self.functions.deinit(self.allocator);
56 self.marks.deinit(self.allocator);
57 self.helpers.deinit(self.allocator);
58 }
59
60 pub fn function(self: *const Plan, name: []const u8) ?*ir.Operation {
61 return self.functions.get(name);
62 }
63
64 fn visitHelper(self: *Plan, func: *ir.Operation, depth: usize, stage: bool) PlanError!void {
65 if (depth >= max_call_depth) return error.UnsupportedOperation;
66 if (self.marks.get(func)) |mark| {
67 if (mark == .visiting) return error.UnsupportedOperation;
68 return;
69 }
70 if (func.getAttr("kernel") != null or gpu.stage.stageOf(func) != null) return error.UnsupportedOperation;
71 if (stage and !admittedSignature(FuncDialect.FuncOp{ .op = func })) {
72 if (self.refused) |name| name.* = (FuncDialect.FuncOp{ .op = func }).getName() orelse "";
73 return error.UnsupportedHelperSignature;
74 }
75 const region = func.getRegion(0) orelse return error.UnsupportedOperation;
76 if (!region.hasOneBlock()) return error.UnsupportedOperation;
77 try validateHelperBlock(region.getEntryBlock().?);
78 try self.marks.put(self.allocator, func, .visiting);
79 try self.walkCalls(region.getEntryBlock().?, depth + 1, stage);
80 try self.marks.put(self.allocator, func, .done);
81 try self.helpers.append(self.allocator, func);
82 }
83
84 fn admittedSignature(func: FuncDialect.FuncOp) bool {
85 const results = func.getResultTypes();
86 if (results.len > 1) return false;
87 if (results.len == 1 and !admittedScalar(results[0])) return false;
88 for (func.getArguments()) |arg| if (!admittedScalar(arg.type)) return false;
89 return true;
90 }
91
92 fn admittedScalar(typ: ir.Type) bool {
93 const name = typ.getDialectTypeName() orelse return false;
94 const kind = choir.dialects.arith.scalarKindFromTypeName(name) orelse return false;
95 return switch (kind) {
96 .bool, .index, .i8, .i16, .i32, .u32, .i64, .f16, .f32 => true,
97 else => false,
98 };
99 }
100
101 fn validateHelperBlock(block: *ir.Block) PlanError!void {
102 var ops = block.getOperations();
103 while (ops.next()) |op| {
104 if (std.mem.startsWith(u8, op.name.name, "gpu.")) return error.UnsupportedOperation;
105 for (0..op.getNumRegions()) |index| {
106 const region = op.getRegion(index) orelse continue;
107 var blocks = region.getBlocks();
108 while (blocks.next()) |nested| try validateHelperBlock(nested);
109 }
110 }
111 }
112
113 fn walkCalls(self: *Plan, block: *ir.Block, depth: usize, stage: bool) PlanError!void {
114 var ops = block.getOperations();
115 while (ops.next()) |op| {
116 if (stage and !gpu.stage.admitsMemory(op)) {
117 if (self.refused) |name| name.* = op.name.name;
118 return error.UnsupportedStageMemory;
119 }
120 if (std.mem.eql(u8, op.name.name, FuncDialect.CallOp.operation_name)) {
121 const call = FuncDialect.CallOp{ .op = op };
122 const name = call.getCallee() orelse return error.UnsupportedOperation;
123 const callee = self.function(name) orelse return error.UnsupportedOperation;
124 try self.visitHelper(callee, depth, stage);
125 }
126 for (0..op.getNumRegions()) |index| {
127 const region = op.getRegion(index) orelse continue;
128 var blocks = region.getBlocks();
129 while (blocks.next()) |nested| try self.walkCalls(nested, depth, stage);
130 }
131 }
132 }
133 };