lib/accy/src/preparation/call.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const choir = @import("choir");
4 const accy_choir = @import("../choir/root.zig");
5 const kernel_library = @import("../kernel/library/root.zig");
6
7 const ir = choir.ir;
8 const rewrite = ir.rewrite;
9 const dialects = choir.dialects;
10 const dialect_mod = accy_choir.dialect;
11 const semantic = accy_choir.semantic;
12
13 pub const CatalogCallOptions = struct {
14 has_side_effects: bool = false,
15 operand_effects: ?[]const semantic.KernelOperandEffect = null,
16 result_aliases: ?[]const ?usize = null,
17 runtime_scalars: ?[]const dialect_mod.AccyDialect.KernelCallScalar = null,
18 };
19
20 pub fn catalogCallScalars(
21 comptime count: usize,
22 arguments: [count]choir_abi.ScalarArgument,
23 ) [count]dialect_mod.AccyDialect.KernelCallScalar {
24 var scalars: [count]dialect_mod.AccyDialect.KernelCallScalar = undefined;
25 for (arguments, 0..) |argument, index| {
26 scalars[index] = switch (argument) {
27 .i32 => |value| .{ .kind = .i32, .bits = @as(u32, @bitCast(value)) },
28 .u32 => |value| .{ .kind = .u32, .bits = value },
29 .i64 => |value| .{ .kind = .i64, .bits = @bitCast(value) },
30 .u64 => |value| .{ .kind = .u64, .bits = value },
31 .f32 => |value| .{ .kind = .f32, .bits = @as(u32, @bitCast(value)) },
32 .f64 => |value| .{ .kind = .f64, .bits = @bitCast(value) },
33 };
34 }
35 return scalars;
36 }
37
38 pub const CatalogCallRequest = struct {
39 descriptor: kernel_library.CatalogDescriptor,
40 operands: []const *ir.Value,
41 result_types: []const ir.Type,
42 options: CatalogCallOptions = .{},
43 };
44
45 pub fn insertCatalogCall(
46 rewriter: *rewrite.PatternRewriter,
47 request: CatalogCallRequest,
48 ) !dialect_mod.AccyDialect.KernelCallOp {
49 const operand_effects = try kernelOperandEffects(rewriter.allocator, request.operands.len, request.options.operand_effects);
50 defer if (request.options.operand_effects == null) rewriter.allocator.free(operand_effects);
51
52 const result_aliases = try kernelResultAliases(rewriter.allocator, request.result_types.len, request.options.result_aliases);
53 defer if (request.options.result_aliases == null) rewriter.allocator.free(result_aliases);
54
55 const metadata = request.descriptor.metadata;
56 const call = try dialect_mod.AccyDialect.KernelCallOp.create(
57 rewriter.ir_ctx,
58 ir.Location.getUnknown(),
59 request.operands,
60 request.result_types,
61 metadata.target,
62 metadata.version,
63 request.options.has_side_effects,
64 operand_effects,
65 result_aliases,
66 );
67 if (request.options.runtime_scalars) |scalars| {
68 try dialect_mod.AccyDialect.setKernelCallRuntimeScalars(rewriter.ir_ctx, call.op, scalars);
69 }
70 _ = try rewriter.insert(call.op);
71 return call;
72 }
73
74 fn kernelOperandEffects(
75 allocator: std.mem.Allocator,
76 operand_count: usize,
77 provided: ?[]const semantic.KernelOperandEffect,
78 ) ![]const semantic.KernelOperandEffect {
79 if (provided) |effects| {
80 if (effects.len != operand_count) return error.InvalidKernelCallContract;
81 return effects;
82 }
83
84 const effects = try allocator.alloc(semantic.KernelOperandEffect, operand_count);
85 @memset(effects, .read);
86 return effects;
87 }
88
89 fn kernelResultAliases(
90 allocator: std.mem.Allocator,
91 result_count: usize,
92 provided: ?[]const ?usize,
93 ) ![]const ?usize {
94 if (provided) |aliases| {
95 if (aliases.len != result_count) return error.InvalidKernelCallContract;
96 return aliases;
97 }
98
99 const aliases = try allocator.alloc(?usize, result_count);
100 @memset(aliases, null);
101 return aliases;
102 }
103
104 fn findOpNamed(op: *ir.Operation, name: []const u8) ?*ir.Operation {
105 if (std.mem.eql(u8, op.name.name, name)) return op;
106 for (op.regions.items) |*region| {
107 var block_iter = region.getBlocks();
108 while (block_iter.next()) |block| {
109 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
110 while (current) |current_op| {
111 if (findOpNamed(current_op, name)) |found| return found;
112 current = current_op.next_op;
113 }
114 }
115 }
116 return null;
117 }
118
119 const testing = std.testing;
120
121 test "catalog call insertion emits descriptor target and default contract" {
122 const allocator = testing.allocator;
123
124 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
125 defer builder.deinit();
126 const f32_8 = try builder.tensor(.f32, &.{8});
127 var fb = try builder.beginFunction("catalog_call_insertion", &.{ f32_8, f32_8 }, &.{f32_8});
128 const lhs = fb.parameter(0);
129 const rhs = fb.parameter(1);
130 const out = try fb.add(lhs, rhs);
131 try fb.return_(&.{out});
132 try fb.finish();
133 const module = try builder.finish();
134 defer module.deinit();
135
136 const return_op = findOpNamed(module.choir_module, dialects.FuncDialect.ReturnOp.operation_name) orelse {
137 return error.TestExpectedReturn;
138 };
139 var rewriter = rewrite.PatternRewriter.init(allocator, module.context());
140 defer rewriter.deinit();
141 rewriter.setInsertionPointBefore(return_op);
142
143 const result_types = [_]ir.Type{f32_8};
144 const descriptor = kernel_library.findEntry(
145 kernel_library.elementwise.VectorAdd8F32.target,
146 kernel_library.elementwise.VectorAdd8F32.version,
147 ) orelse return error.TestExpectedCatalogDescriptor;
148 const call = try insertCatalogCall(&rewriter, .{
149 .descriptor = descriptor,
150 .operands = &.{ lhs, rhs },
151 .result_types = &result_types,
152 });
153
154 try module.verify();
155 const target_attr = call.op.getAttr("target") orelse return error.TestExpectedTarget;
156 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedTarget;
157 try testing.expectEqualStrings(dialect_mod.AccyDialect.KernelCallOp.target_attr_name, target_attr.abstract.name);
158 try testing.expectEqualStrings(kernel_library.elementwise.VectorAdd8F32.target, target.payload);
159
160 const version = call.op.getAttrAs(ir.Attribute.IntegerAttr, "version") orelse return error.TestExpectedVersion;
161 try testing.expectEqual(@as(i64, kernel_library.elementwise.VectorAdd8F32.version), version.getValue());
162
163 const effects_attr = call.op.getAttr("operand_effects") orelse {
164 return error.TestExpectedOperandEffects;
165 };
166 const effects = effects_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedOperandEffects;
167 try testing.expectEqualStrings(dialect_mod.AccyDialect.KernelCallOp.operand_effects_attr_name, effects_attr.abstract.name);
168 try testing.expectEqual(@as(usize, 2), effects.payload.len);
169 try testing.expectEqual(@backingInt(semantic.KernelOperandEffect.read), effects.payload[0]);
170 try testing.expectEqual(@backingInt(semantic.KernelOperandEffect.read), effects.payload[1]);
171
172 const aliases_attr = call.op.getAttr("result_aliases") orelse {
173 return error.TestExpectedResultAliases;
174 };
175 const aliases = aliases_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedResultAliases;
176 try testing.expectEqualStrings(dialect_mod.AccyDialect.KernelCallOp.result_aliases_attr_name, aliases_attr.abstract.name);
177 try testing.expectEqual(@as(usize, @sizeOf(i64)), aliases.payload.len);
178 const values = std.mem.bytesAsSlice(i64, aliases.payload);
179 try testing.expectEqual(@as(i64, -1), values[0]);
180 }