lib/choir/src/backends/signature.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 //! Boundary signatures of IR function definitions and of Zig function pointer types.
2
3 const std = @import("std");
4 const alloc_arena = @import("alloc_arena");
5 const ir = @import("../core/root.zig");
6 const dialects = @import("../dialects/root.zig");
7 const artifact = @import("artifact/root.zig");
8
9 const Signature = artifact.Signature;
10 const ScalarType = artifact.ScalarType;
11 const ValueType = artifact.ValueType;
12
13 pub const Error = error{ TooManyParameters, TooManyResults, UnsupportedType };
14
15 /// Records the parameter and result types of the function definition `func`.
16 /// Fails with `error.TooManyParameters` or `error.TooManyResults` when `func` has more parameters
17 /// or results than a `Signature` records, and with `error.UnsupportedType` when a type has no
18 /// boundary equivalent, such as `arith.f16`.
19 pub fn ofFunction(func: *ir.Operation) Error!Signature {
20 std.debug.assert(std.mem.eql(u8, func.name.name, dialects.FuncDialect.FuncOp.operation_name));
21 const arguments = func.getRegion(0).?.getEntryBlock().?.arguments.items;
22 const results = func.results.items;
23 if (arguments.len > Signature.max_parameters) return error.TooManyParameters;
24 if (results.len > Signature.max_results) return error.TooManyResults;
25 var parameter_types: [Signature.max_parameters]ValueType = undefined;
26 for (arguments, parameter_types[0..arguments.len]) |argument, *value_type| {
27 value_type.* = try valueTypeOf(argument.type);
28 }
29 var result_types: [Signature.max_results]ValueType = undefined;
30 for (results, result_types[0..results.len]) |result, *value_type| {
31 value_type.* = try valueTypeOf(result.type);
32 }
33 const signature = Signature.init(
34 parameter_types[0..arguments.len],
35 result_types[0..results.len],
36 ) catch unreachable;
37 std.debug.assert(signature.parameters().len == arguments.len);
38 std.debug.assert(signature.results().len == results.len);
39 return signature;
40 }
41
42 /// Maps the IR type `ty` to its boundary type.
43 /// Fails with `error.UnsupportedType` when `ty` has none.
44 pub fn valueTypeOf(ty: ir.Type) error{UnsupportedType}!ValueType {
45 const name = ty.getDialectTypeName() orelse return error.UnsupportedType;
46 if (std.mem.eql(u8, name, dialects.MemrefDialect.name)) return .memref;
47 if (dialects.arith.parseVectorTypeName(name)) |vector| {
48 const kind = dialects.arith.scalarKindFromTypeName(vector.elem_type_name).?;
49 const lanes = std.math.cast(u8, vector.width) orelse return error.UnsupportedType;
50 const element = try scalarOf(kind);
51 const value_type = ValueType{ .vector = .{ .element = element, .lanes = lanes } };
52 if (!value_type.isValid()) return error.UnsupportedType;
53 return value_type;
54 }
55 const kind = dialects.arith.scalarKindFromTypeName(name) orelse return error.UnsupportedType;
56 return .{ .scalar = try scalarOf(kind) };
57 }
58
59 fn scalarOf(kind: dialects.arith.ScalarKind) error{UnsupportedType}!ScalarType {
60 return switch (kind) {
61 .f16, .bf16 => error.UnsupportedType,
62 inline else => |tag| @field(ScalarType, @tagName(tag)),
63 };
64 }
65
66 /// Derives the signature of the C calling-convention function that `FunctionPointer` points to.
67 /// Each parameter, and a result other than `void` or a struct, maps to one boundary type.
68 /// `i8`, `i16`, `i32`, `i64`, and their unsigned forms map to the arith integer of the same name.
69 /// `usize` maps to `arith.index`. `bool`, `f32`, and `f64` map to the arith type of the same name.
70 /// A pointer that is not a slice, optional or not, maps to `memref`.
71 /// A `void` result maps to no results.
72 /// An extern struct result of `n` fields maps to `n` results. Field `i` starts at byte `8 * i`,
73 /// and the struct spans `8 * n` bytes.
74 /// Any other type, `isize` included, is a compile error.
75 pub fn ofZigFunction(comptime FunctionPointer: type) Signature {
76 return comptime derive(FunctionPointer);
77 }
78
79 fn derive(comptime FunctionPointer: type) Signature {
80 const name = @typeName(FunctionPointer);
81 const function = switch (@typeInfo(FunctionPointer)) {
82 .pointer => |pointer| switch (@typeInfo(pointer.child)) {
83 .@"fn" => |function| function,
84 else => @compileError(name ++ " does not point to a function"),
85 },
86 else => @compileError(name ++ " is not a function pointer"),
87 };
88 if (!function.attrs.@"callconv".eql(.c)) @compileError(name ++ " is not callconv(.c)");
89 if (function.attrs.varargs) @compileError(name ++ " is variadic");
90 if (function.is_generic) @compileError(name ++ " is generic");
91 var parameter_types: [function.param_types.len]ValueType = undefined;
92 for (function.param_types, ¶meter_types) |parameter_type, *value_type| {
93 value_type.* = zigValueType(parameter_type.?);
94 }
95 const Return = function.return_type.?;
96 const result_types: []const ValueType = switch (@typeInfo(Return)) {
97 .void => &.{},
98 .@"struct" => &zigProductTypes(Return),
99 else => &.{zigValueType(Return)},
100 };
101 return Signature.init(¶meter_types, result_types) catch |err| {
102 @compileError(name ++ ": " ++ @errorName(err));
103 };
104 }
105
106 fn zigValueType(comptime T: type) ValueType {
107 const missing = @typeName(T) ++ " has no boundary type";
108 if (T == usize) return .{ .scalar = .index };
109 if (T == isize) @compileError(missing ++ "; use i64");
110 return switch (@typeInfo(T)) {
111 .bool => .{ .scalar = .bool },
112 .int => |int| .{ .scalar = switch (int.bits) {
113 8 => if (int.signedness == .signed) .i8 else .u8,
114 16 => if (int.signedness == .signed) .i16 else .u16,
115 32 => if (int.signedness == .signed) .i32 else .u32,
116 64 => if (int.signedness == .signed) .i64 else .u64,
117 else => @compileError(missing),
118 } },
119 .float => |float| .{ .scalar = switch (float.bits) {
120 32 => .f32,
121 64 => .f64,
122 else => @compileError(missing),
123 } },
124 .pointer => |pointer| if (pointer.size == .slice) @compileError(missing) else .memref,
125 .optional => |optional| switch (@typeInfo(optional.child)) {
126 .pointer => |pointer| if (pointer.size == .slice) @compileError(missing) else .memref,
127 else => @compileError(missing),
128 },
129 else => @compileError(missing),
130 };
131 }
132
133 fn zigProductTypes(
134 comptime Product: type,
135 ) [@typeInfo(Product).@"struct".field_names.len]ValueType {
136 if (comptime productLayoutError(Product)) |message| @compileError(message);
137 const product = @typeInfo(Product).@"struct";
138 var types: [product.field_names.len]ValueType = undefined;
139 inline for (product.field_types, &types) |Field, *slot| slot.* = zigValueType(Field);
140 return types;
141 }
142
143 fn productLayoutError(comptime Product: type) ?[]const u8 {
144 const product = @typeInfo(Product).@"struct";
145 const name = @typeName(Product);
146 if (product.layout != .@"extern") return name ++ " is not extern";
147 inline for (product.field_names, 0..) |field, index| {
148 const offset = @offsetOf(Product, field);
149 if (offset != 8 * index) return std.fmt.comptimePrint(
150 "{s}.{s} starts at byte {d}, not {d}",
151 .{ name, field, offset, 8 * index },
152 );
153 }
154 const size = 8 * product.field_names.len;
155 if (@sizeOf(Product) == size) return null;
156 return std.fmt.comptimePrint("{s} spans {d} bytes, not {d}", .{ name, @sizeOf(Product), size });
157 }
158
159 const SampleProduct = extern struct {
160 sum: i64,
161 scale: f64,
162 narrow: i32 align(8),
163 below: bool align(8),
164 };
165
166 const OverAlignedProduct = extern struct {
167 value: i64 align(32),
168 };
169
170 test "Zig function pointer types map to boundary signatures" {
171 const Function = *const fn (
172 u8,
173 i16,
174 u32,
175 i64,
176 usize,
177 bool,
178 f32,
179 f64,
180 [*]const f64,
181 ?*anyopaque,
182 ) callconv(.c) SampleProduct;
183 const expected = try Signature.init(&.{
184 .{ .scalar = .u8 }, .{ .scalar = .i16 }, .{ .scalar = .u32 }, .{ .scalar = .i64 },
185 .{ .scalar = .index }, .{ .scalar = .bool }, .{ .scalar = .f32 }, .{ .scalar = .f64 },
186 .memref, .memref,
187 }, &.{ .{ .scalar = .i64 }, .{ .scalar = .f64 }, .{ .scalar = .i32 }, .{ .scalar = .bool } });
188 try std.testing.expect(ofZigFunction(Function).eql(&expected));
189
190 const Nullary = *const fn () callconv(.c) void;
191 const empty = ofZigFunction(Nullary);
192 try std.testing.expectEqual(@as(usize, 0), empty.parameters().len);
193 try std.testing.expectEqual(@as(usize, 0), empty.results().len);
194 const Address = *const fn (*u8) callconv(.c) [*]u8;
195 const address = try Signature.init(&.{.memref}, &.{.memref});
196 try std.testing.expect(ofZigFunction(Address).eql(&address));
197 }
198
199 test "Zig result structs span eight bytes per field" {
200 try std.testing.expect(comptime productLayoutError(SampleProduct) == null);
201 const message = comptime productLayoutError(OverAlignedProduct).?;
202 const expected = "OverAlignedProduct spans 32 bytes, not 8";
203 try std.testing.expect(std.mem.endsWith(u8, message, expected));
204 }
205
206 test "function definitions record boundary signatures" {
207 var arena = alloc_arena.Arena.init(std.testing.allocator);
208 defer arena.deinit();
209 const allocator = arena.allocator();
210 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
211 defer ctx.deinit(allocator);
212 try dialects.registerAllDialects(&ctx);
213 const Arith = dialects.ArithDialect;
214 const Func = dialects.FuncDialect;
215 const loc = ir.Location.getUnknown();
216 const i32_type = try Arith.getScalarType(&ctx, .i32);
217 const f16_type = try Arith.getScalarType(&ctx, .f16);
218 const memref = try dialects.MemrefDialect.getMemrefTypeDynamic(&ctx, i32_type, .host);
219 const vector = (try Arith.getVecType(&ctx, 4, dialects.arith.type_names.float32)).?;
220
221 const typed = try Func.FuncOp.create(&ctx, loc, "typed", &.{ i32_type, memref }, &.{vector});
222 const lanes = ValueType{ .vector = .{ .element = .f32, .lanes = 4 } };
223 const expected = try Signature.init(&.{ .{ .scalar = .i32 }, .memref }, &.{lanes});
224 try std.testing.expect((try ofFunction(typed.op)).eql(&expected));
225
226 const half = try Func.FuncOp.create(&ctx, loc, "half", &.{f16_type}, &.{});
227 try std.testing.expectError(error.UnsupportedType, ofFunction(half.op));
228 var wide: [Signature.max_parameters + 1]ir.Type = undefined;
229 @memset(&wide, i32_type);
230 const many = try Func.FuncOp.create(&ctx, loc, "many", &wide, &.{});
231 try std.testing.expectError(error.TooManyParameters, ofFunction(many.op));
232 const results = wide[0 .. Signature.max_results + 1];
233 const product = try Func.FuncOp.create(&ctx, loc, "product", &.{}, results);
234 try std.testing.expectError(error.TooManyResults, ofFunction(product.op));
235 }