lib/accy/src/preparation/kernelization/lowering/common.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const alloc_fixed = @import("alloc_fixed");
4 const choir = @import("choir");
5 const accy_choir = @import("../../../choir/root.zig");
6
7 pub const ir = choir.ir;
8 pub const LoweringError = gpu.BackendError || error{ WorkExhausted, WorkOverflow };
9 pub const dialect_mod = accy_choir.dialect;
10 pub const kernel_root = @import("../../../kernel/model/root.zig");
11 pub const bufferization = @import("../../bufferization/root.zig");
12 pub const fusion = @import("../../fusion/root.zig");
13 const kernelization_model = @import("../model/root.zig");
14 pub const schedule_planning = @import("../../schedule/root.zig");
15 pub const shape_analysis = @import("../../shape/root.zig");
16
17 pub const i64_attr_list_stack_capacity: usize = 8;
18
19 pub fn dialectAttrPayload(
20 op: *ir.Operation,
21 attr_name: []const u8,
22 expected_dialect_attr_name: []const u8,
23 ) gpu.BackendError![]const u8 {
24 const attr = op.getAttr(attr_name) orelse return error.InvalidArtifact;
25 if (!std.mem.eql(u8, attr.abstract.name, expected_dialect_attr_name)) return error.InvalidArtifact;
26 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidArtifact;
27 return dialect_attr.payload;
28 }
29
30 pub fn readIntegerAttr(op: *const ir.Operation, attr_name: []const u8) gpu.BackendError!i64 {
31 const attr = op.getAttr(attr_name) orelse return error.InvalidArtifact;
32 if (!std.mem.eql(u8, attr.abstract.name, "builtin.integer")) return error.InvalidArtifact;
33 const int_attr = attr.cast(ir.Attribute.IntegerAttr) orelse return error.InvalidArtifact;
34 return int_attr.getValue();
35 }
36
37 pub fn readI64ListAttrBounded(
38 op: *ir.Operation,
39 attr_name: []const u8,
40 expected_dialect_attr_name: []const u8,
41 buffer: []i64,
42 ) gpu.BackendError![]const i64 {
43 const payload = try dialectAttrPayload(op, attr_name, expected_dialect_attr_name);
44 if (payload.len % @sizeOf(i64) != 0) return error.InvalidArtifact;
45 if (payload.len / @sizeOf(i64) > buffer.len) return error.CapabilityMismatch;
46 const values = buffer[0 .. payload.len / @sizeOf(i64)];
47 for (values, 0..) |*value, index| {
48 const start = index * @sizeOf(i64);
49 @memcpy(std.mem.asBytes(value), payload[start..][0..@sizeOf(i64)]);
50 }
51 return values;
52 }
53
54 pub fn constantPayloadForSlot(
55 slot: bufferization.BufferSlot,
56 ) gpu.BackendError![]const u8 {
57 if (!slot.role.constant) return &.{};
58 const producer = slot.producer orelse return error.InvalidArtifact;
59 if (!isName(producer.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
60 return error.InvalidArtifact;
61 }
62 const attr = producer.getAttr("payload") orelse return error.InvalidArtifact;
63 if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.ConstantOp.payload_attr_name)) {
64 return error.InvalidArtifact;
65 }
66 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidArtifact;
67 if (slot.byte_size) |byte_size| {
68 const expected = std.math.cast(usize, byte_size) orelse return error.InvalidArtifact;
69 if (dialect_attr.payload.len != expected) return error.InvalidArtifact;
70 }
71 return dialect_attr.payload;
72 }
73
74 pub fn constantPayloadEqualsF32(payload: []const u8, expected: f32) bool {
75 return payload.len == @sizeOf(f32) and std.mem.eql(u8, payload, std.mem.asBytes(&expected));
76 }
77
78 pub fn constantPayloadEqualsI32(payload: []const u8, expected: i32) bool {
79 return payload.len == @sizeOf(i32) and std.mem.eql(u8, payload, std.mem.asBytes(&expected));
80 }
81
82 pub fn constantPayloadEqualsF16(payload: []const u8, expected: f16) bool {
83 return payload.len == @sizeOf(f16) and std.mem.eql(u8, payload, std.mem.asBytes(&expected));
84 }
85
86 pub fn staticElementCountU32(slot: bufferization.BufferSlot) gpu.BackendError!u32 {
87 const count = slot.element_count orelse return error.CapabilityMismatch;
88 if (count == 0) return error.CapabilityMismatch;
89 return std.math.cast(u32, count) orelse return error.CapabilityMismatch;
90 }
91
92 pub fn staticPositiveDimU32(dim: i64) gpu.BackendError!u32 {
93 if (dim <= 0) return error.CapabilityMismatch;
94 return std.math.cast(u32, dim) orelse return error.CapabilityMismatch;
95 }
96
97 pub fn bufferSlotById(
98 buffer_plan: *const bufferization.BufferPlanAnalysis,
99 slot_id: usize,
100 ) ?*const bufferization.BufferSlot {
101 if (slot_id < buffer_plan.slots.items.len and buffer_plan.slots.items[slot_id].id == slot_id) {
102 return &buffer_plan.slots.items[slot_id];
103 }
104 for (buffer_plan.slots.items) |*slot| {
105 if (slot.id == slot_id) return slot;
106 }
107 return null;
108 }
109
110 pub fn externalInputIndex(
111 outline: kernelization_model.KernelOutline,
112 slot_id: usize,
113 ) ?usize {
114 for (outline.input_slot_ids, 0..) |input_slot_id, index| {
115 if (input_slot_id == slot_id) return index;
116 }
117 return null;
118 }
119
120 pub fn mapKernelBuildError(err: anyerror) LoweringError {
121 return switch (err) {
122 error.WorkExhausted,
123 error.ParameterCapacityExceeded,
124 error.KernelNameCapacityExceeded,
125 error.TemporaryValueCapacityExceeded,
126 error.TemporaryTypeCapacityExceeded,
127 error.ContextCapacityExceeded,
128 error.AxisCapacityExceeded,
129 error.StepCapacityExceeded,
130 error.NameCapacityExceeded,
131 error.AxisIdCapacityExceeded,
132 error.NestingLimitExceeded,
133 => error.WorkExhausted,
134 error.WorkOverflow,
135 error.CapacityOverflow,
136 error.ValueCountOverflow,
137 error.CountOverflow,
138 => error.WorkOverflow,
139 error.OutOfMemory => error.OutOfMemory,
140 error.UnsupportedDType => error.CapabilityMismatch,
141 error.VerificationFailed => error.InvalidArtifact,
142 else => error.InvalidArtifact,
143 };
144 }
145
146 pub fn mapGeneratedKernelError(err: anyerror) LoweringError {
147 return switch (err) {
148 error.UnsupportedOperation => error.UnsupportedOperation,
149 error.UnsupportedArtifactFormat => error.UnsupportedArtifactFormat,
150 error.CapabilityMismatch => error.CapabilityMismatch,
151 error.InvalidArtifact => error.InvalidArtifact,
152 error.OutOfMemory => error.OutOfMemory,
153 else => mapKernelBuildError(err),
154 };
155 }
156
157 pub fn isName(actual: []const u8, expected: []const u8) bool {
158 return std.mem.eql(u8, actual, expected);
159 }
160
161 pub fn optionalSplatConstantValue(
162 inner: anytype,
163 op: *ir.Operation,
164 ) LoweringError!?kernel_root.Value {
165 return splatConstantValue(inner, op) catch |err| switch (err) {
166 error.UnsupportedOperation => null,
167 else => return mapGeneratedKernelError(err),
168 };
169 }
170
171 pub fn splatConstantValue(inner: anytype, op: *ir.Operation) !kernel_root.Value {
172 const constant = dialect_mod.AccyDialect.ConstantOp{ .op = op };
173 const payload = constant.getPayload() orelse return error.InvalidArtifact;
174 const result = op.getResult(0) orelse return error.InvalidArtifact;
175 var dtype_arena_buffer: [160]u8 = undefined;
176 var dtype_arena = alloc_fixed.FixedBuffer.init(dtype_arena_buffer[0..]);
177 const result_type = dialect_mod.decodeTensorType(dtype_arena.allocator(), result.type) catch return error.InvalidArtifact;
178 switch (result_type.dtype) {
179 .f32 => {
180 if (payload.len < 4 or payload.len % 4 != 0) return error.UnsupportedOperation;
181 var first: f32 = undefined;
182 @memcpy(std.mem.asBytes(&first), payload[0..4]);
183 var offset: usize = 4;
184 while (offset < payload.len) : (offset += 4) {
185 var value: f32 = undefined;
186 @memcpy(std.mem.asBytes(&value), payload[offset..][0..4]);
187 if (@as(u32, @bitCast(value)) != @as(u32, @bitCast(first))) return error.UnsupportedOperation;
188 }
189 return inner.constantFloat(.f32, first);
190 },
191 .i32 => {
192 if (payload.len < 4 or payload.len % 4 != 0) return error.UnsupportedOperation;
193 var first: i32 = undefined;
194 @memcpy(std.mem.asBytes(&first), payload[0..4]);
195 var offset: usize = 4;
196 while (offset < payload.len) : (offset += 4) {
197 var value: i32 = undefined;
198 @memcpy(std.mem.asBytes(&value), payload[offset..][0..4]);
199 if (value != first) return error.UnsupportedOperation;
200 }
201 return inner.constantInt(.i32, first);
202 },
203 .i1 => {
204 if (payload.len == 0) return error.UnsupportedOperation;
205 const first = payload[0];
206 for (payload[1..]) |value| {
207 if (value != first) return error.UnsupportedOperation;
208 }
209 return inner.constantBool(first != 0);
210 },
211 else => return error.UnsupportedOperation,
212 }
213 }