lib/accy/src/preparation/activation.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const choir = @import("choir");
4 const accy_choir = @import("../choir/root.zig");
5 const kernel_library = @import("../kernel/library/root.zig");
6 const kernel_selection = @import("../kernel/logical/selection/root.zig");
7 const call_preparation = @import("call.zig");
8 const library_preparation = @import("library.zig");
9
10 const ir = choir.ir;
11 const rewrite = ir.rewrite;
12 const passes = choir.passes;
13 const work = passes.pass.work;
14 const dialect_mod = accy_choir.dialect;
15 const semantics = accy_choir.semantics;
16
17 pub const Options = struct {
18 kernel_library: library_preparation.KernelLibraryLowering = .disabled,
19
20 pub fn eql(self: Options, other: Options) bool {
21 return self.kernel_library == other.kernel_library;
22 }
23 };
24
25 pub const activation_lowering_pass_name = "accy-choir-activation-lower";
26 pub const activation_lowering_pass_description =
27 "Lower semantic Accy activation patterns into selected tensor kernels";
28
29 const kernel_library_option_choices = [_]passes.PassOptionChoice{
30 .{ .name = "disabled" },
31 .{ .name = "enabled" },
32 };
33
34 pub const activation_lowering_pass_options = [_]passes.PassOptionSpec{
35 .{
36 .name = "kernel-library",
37 .description = "Use kernel library calls for supported activation patterns",
38 .kind = .choice,
39 .choices = &kernel_library_option_choices,
40 .default_value = "disabled",
41 },
42 };
43
44 pub fn activationLoweringPass() passes.Pass {
45 return .{
46 .name = activation_lowering_pass_name,
47 .description = activation_lowering_pass_description,
48 .run_fn = runActivationLoweringPass,
49 .work_contract = activation_work_contract,
50 };
51 }
52
53 pub fn activationLoweringPassWithOptions(options: *const Options) passes.Pass {
54 return .{
55 .name = activation_lowering_pass_name,
56 .description = activation_lowering_pass_description,
57 .state = @constCast(options),
58 .run_with_state_fn = runActivationLoweringPassWithState,
59 .work_contract = activation_work_contract,
60 };
61 }
62
63 pub fn activationLoweringPassFromOptions(allocator: std.mem.Allocator, set: passes.PassOptionSet) anyerror!passes.Pass {
64 const options = try allocator.create(Options);
65 errdefer allocator.destroy(options);
66 options.* = .{
67 .kernel_library = try kernelLibraryLoweringFromText(set.choiceValue("kernel-library", "disabled")),
68 };
69 var pass = activationLoweringPassWithOptions(options);
70 pass.state_deinit_fn = destroyOptions;
71 return pass;
72 }
73
74 const activation_work_contract: work.Contract = .{
75 .identity = .{ .name = activation_lowering_pass_name, .version = 1 },
76 .estimate = activationWork,
77 };
78
79 const ActivationWork = struct {
80 options: Options,
81 created: u64 = 0,
82 erased: u64 = 0,
83 payload: u64 = 0,
84 temporary: u64 = 0,
85 metadata: u64 = 0,
86
87 fn visit(self: *ActivationWork, op: *ir.Operation) !ir.WalkResult {
88 const activation = std.mem.eql(
89 u8,
90 op.name.name,
91 dialect_mod.AccyDialect.ActivationOp.operation_name,
92 );
93 const maximum = std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MaxOp.operation_name);
94 if (!activation and !(maximum and self.options.kernel_library == .enabled)) return .advance;
95 if (op.getNumResults() != 1) return .advance;
96 const kind = if (activation) activationKindFromOp(op) orelse return .advance else .relu;
97 const shape = try ActivationShape.inspect(op.getResult(0).?.type);
98 const decoding = try work.multiply(2, try work.add(try work.multiply(shape.rank, 8), 128));
99 self.temporary = try work.add(self.temporary, try work.multiply(4, decoding));
100 if (self.options.kernel_library == .enabled and shape.rank == 1) {
101 if (kernel_selection.selectActivationCatalog(.{
102 .dtype = shape.dtype,
103 .kind = kind,
104 .extent = shape.elements,
105 })) |selected| {
106 self.created = try work.add(self.created, 1);
107 self.erased = try work.add(self.erased, if (maximum) 2 else 1);
108 self.temporary = try work.add(self.temporary, 64);
109 self.metadata = try work.add(
110 self.metadata,
111 selected.descriptor.metadata.target.len + 64,
112 );
113 return .advance;
114 }
115 }
116 if (!activation or !shape.dtype.isFloat()) return .advance;
117 const constants: u64 = if (kind == .gelu) 4 else 1;
118 const nodes: u64 = switch (kind) {
119 .relu => 2,
120 .silu => 5,
121 .gelu => 13,
122 };
123 self.created = try work.add(self.created, nodes);
124 self.erased = try work.add(self.erased, 1);
125 self.payload = try work.add(self.payload, try work.multiply(
126 constants,
127 try work.multiply(shape.elements, shape.dtype.sizeOf()),
128 ));
129 self.temporary = try work.add(self.temporary, try work.multiply(constants, decoding));
130 return .advance;
131 }
132 };
133
134 const ActivationShape = struct {
135 dtype: semantics.DType,
136 rank: u64 = 0,
137 elements: u64 = 1,
138
139 fn inspect(typ: ir.Type) !ActivationShape {
140 const name = typ.getDialectTypeName() orelse return error.ExpectedAccyTensorType;
141 if (!std.mem.eql(u8, name, dialect_mod.tensor_type_name)) {
142 return error.ExpectedAccyTensorType;
143 }
144 const key = typ.getDialectParamKey() orelse return error.MalformedAccyTensorType;
145 const comma = std.mem.indexOfScalar(u8, key, ',') orelse
146 return error.MalformedAccyTensorType;
147 var shape: ActivationShape = .{
148 .dtype = semantics.DType.fromName(key[0..comma]) orelse return error.UnknownAccyDType,
149 };
150 const dims = key[comma + 1 ..];
151 if (dims.len == 0) return shape;
152 var iter = std.mem.splitScalar(u8, dims, 'x');
153 while (iter.next()) |part| {
154 const dim = std.fmt.parseInt(i64, part, 10) catch return error.MalformedAccyTensorType;
155 if (dim < 0) return error.MalformedAccyTensorType;
156 shape.rank = try work.add(shape.rank, 1);
157 shape.elements = try work.multiply(shape.elements, @intCast(dim));
158 }
159 return shape;
160 }
161 };
162
163 fn activationWork(input: work.Input) !work.Bounds {
164 const options: *const Options = if (input.state) |state| @ptrCast(@alignCast(state)) else &.{};
165 const counts = try work.Census.inspect(input.operation);
166 var facts: ActivationWork = .{ .options = options.* };
167 _ = try input.operation.walk(.{ .order = .pre_order }, &facts, ActivationWork.visit);
168 const queues = try work.add(
169 try work.arrayListGrowth(*ir.Operation, facts.created),
170 try work.arrayListGrowth(*ir.Operation, facts.erased),
171 );
172 const bytes = try work.add(queues, try work.add(facts.temporary, facts.payload));
173 const units = try work.add(try work.add(counts.atoms, counts.input_bytes), 1);
174 const uses = try work.add(try work.add(counts.values, counts.operands), facts.created);
175 const traversal = try work.multiply(128, try work.multiply(units, try work.add(uses, 1)));
176 const nodes = try work.multiply(facts.created, @sizeOf(ir.Operation) + @sizeOf(ir.Value) +
177 2 * @sizeOf(ir.OpOperand) + 8 * @sizeOf(ir.NamedAttribute) + 256);
178 return .{
179 .work = .{
180 .input_bytes = counts.input_bytes,
181 .output_bytes = try work.add(nodes, try work.add(facts.metadata, facts.payload)),
182 .structural_visits = try work.add(traversal, try work.multiply(16, facts.payload)),
183 .allocation_capacity = bytes,
184 },
185 .workspace = bytes,
186 };
187 }
188
189 fn destroyOptions(raw: ?*anyopaque, allocator: std.mem.Allocator) void {
190 const options: *Options = @ptrCast(@alignCast(raw orelse return));
191 allocator.destroy(options);
192 }
193
194 fn kernelLibraryLoweringFromText(value: []const u8) !library_preparation.KernelLibraryLowering {
195 if (std.mem.eql(u8, value, "disabled")) return .disabled;
196 if (std.mem.eql(u8, value, "enabled")) return .enabled;
197 return error.InvalidPassOptionValue;
198 }
199
200 fn runActivationLoweringPass(pass_ctx: *passes.PassContext) passes.PassResult {
201 return runActivationLoweringWithOptions(pass_ctx, .{});
202 }
203
204 fn runActivationLoweringPassWithState(raw: ?*anyopaque, pass_ctx: *passes.PassContext) passes.PassResult {
205 const options: *const Options = @ptrCast(@alignCast(raw orelse return .failure));
206 return runActivationLoweringWithOptions(pass_ctx, options.*);
207 }
208
209 fn runActivationLoweringWithOptions(pass_ctx: *passes.PassContext, options: Options) passes.PassResult {
210 var rewriter = rewrite.PatternRewriter.init(pass_ctx.allocator, pass_ctx.ir_ctx);
211 defer rewriter.deinit();
212
213 var lowered_count: usize = 0;
214 lowerOnOp(pass_ctx.op, &rewriter, options, &lowered_count) catch return .failure;
215 if (lowered_count == 0) {
216 pass_ctx.preserveAllAnalyses();
217 } else {
218 rewriter.finalize(pass_ctx.op);
219 pass_ctx.markModified();
220 }
221 return .success;
222 }
223
224 fn lowerOnOp(
225 op: *ir.Operation,
226 rewriter: *rewrite.PatternRewriter,
227 options: Options,
228 lowered_count: *usize,
229 ) !void {
230 for (op.regions.items) |*region| {
231 var block_iter = region.getBlocks();
232 while (block_iter.next()) |block| {
233 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
234 while (current) |current_op| {
235 const next = current_op.next_op;
236 if (current_op.regions.items.len > 0) {
237 try lowerOnOp(current_op, rewriter, options, lowered_count);
238 }
239 if (std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.ActivationOp.operation_name)) {
240 var guard = rewriter.insertionGuard();
241 defer guard.deinit();
242 rewriter.setInsertionPointBefore(current_op);
243 if (try lowerActivationOp(current_op, rewriter, options)) {
244 lowered_count.* += 1;
245 }
246 } else if (options.kernel_library == .enabled and std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.MaxOp.operation_name)) {
247 var guard = rewriter.insertionGuard();
248 defer guard.deinit();
249 rewriter.setInsertionPointBefore(current_op);
250 if (try lowerReluMaxOp(current_op, rewriter)) |lowered| {
251 lowered_count.* += 1;
252 if (lowered.zero_op) |erased| try rewriter.eraseOp(erased);
253 }
254 }
255 current = next;
256 }
257 }
258 }
259 }
260
261 fn lowerActivationOp(
262 op: *ir.Operation,
263 rewriter: *rewrite.PatternRewriter,
264 options: Options,
265 ) !bool {
266 if (op.getNumResults() != 1) return false;
267 const result = op.getResult(0) orelse return false;
268 const operands = op.getOperandValues();
269 if (operands.len != 1) return false;
270 const input = operands[0];
271 if (!input.type.eql(result.type)) return false;
272 const kind = activationKindFromOp(op) orelse return false;
273
274 if (options.kernel_library == .enabled) {
275 if (try selectedActivationDescriptor(rewriter.allocator, result.type, kind)) |descriptor| {
276 const result_types = [_]ir.Type{result.type};
277 const call = try call_preparation.insertCatalogCall(rewriter, .{
278 .descriptor = descriptor,
279 .operands = &.{input},
280 .result_types = &result_types,
281 });
282 try rewriter.replaceOpWithValue(op, call.getFirstResult());
283 return true;
284 }
285 }
286
287 const replacement = try expandActivationToPrimitives(rewriter, result.type, input, kind) orelse return false;
288 try rewriter.replaceOpWithValue(op, replacement);
289 return true;
290 }
291
292 const LoweredRelu = struct {
293 zero_op: ?*ir.Operation,
294 };
295
296 fn lowerReluMaxOp(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) !?LoweredRelu {
297 if (!std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MaxOp.operation_name)) return null;
298 if (op.getNumResults() != 1) return null;
299 const result = op.getResult(0) orelse return null;
300 const operands = op.getOperandValues();
301 if (operands.len != 2) return null;
302
303 const match = try matchReluOperands(rewriter.allocator, op, result.type, operands[0], operands[1]) orelse return null;
304
305 const result_types = [_]ir.Type{result.type};
306 const call = try call_preparation.insertCatalogCall(rewriter, .{
307 .descriptor = match.descriptor,
308 .operands = &.{match.input},
309 .result_types = &result_types,
310 });
311 try rewriter.replaceOpWithValue(op, call.getFirstResult());
312 return .{ .zero_op = if (match.erase_zero) match.zero_op else null };
313 }
314
315 const ReluMatch = struct {
316 input: *ir.Value,
317 zero_op: *ir.Operation,
318 erase_zero: bool,
319 descriptor: kernel_library.CatalogDescriptor,
320 };
321
322 fn matchReluOperands(
323 allocator: std.mem.Allocator,
324 op: *ir.Operation,
325 result_type: ir.Type,
326 lhs: *ir.Value,
327 rhs: *ir.Value,
328 ) !?ReluMatch {
329 if (try matchReluOperandOrder(allocator, op, result_type, lhs, rhs)) |match| return match;
330 return try matchReluOperandOrder(allocator, op, result_type, rhs, lhs);
331 }
332
333 fn matchReluOperandOrder(
334 allocator: std.mem.Allocator,
335 op: *ir.Operation,
336 result_type: ir.Type,
337 input: *ir.Value,
338 zero: *ir.Value,
339 ) !?ReluMatch {
340 const zero_op = constantDefiningOp(zero) orelse return null;
341 if (constantDefiningOp(input) != null) return null;
342 if (!input.type.eql(result_type) or !zero.type.eql(result_type)) return null;
343
344 var arena_state = alloc_arena.Arena.init(allocator);
345 defer arena_state.deinit();
346 const arena = arena_state.allocator();
347 const tensor_type = try dialect_mod.decodeTensorType(arena, result_type);
348 if (tensor_type.dtype != .f32 or tensor_type.dims.len != 1) return null;
349 if (tensor_type.dims[0] <= 0) return null;
350
351 const payload = (dialect_mod.AccyDialect.ConstantOp{ .op = zero_op }).getPayload() orelse return null;
352 if (!payloadIsZeroF32(payload, @intCast(tensor_type.dims[0]))) return null;
353
354 const descriptor = try selectedActivationDescriptor(allocator, result_type, .relu) orelse return null;
355
356 return .{
357 .input = input,
358 .zero_op = zero_op,
359 .erase_zero = hasOnlyUseBy(zero, op),
360 .descriptor = descriptor,
361 };
362 }
363
364 fn selectedActivationDescriptor(
365 allocator: std.mem.Allocator,
366 result_type: ir.Type,
367 kind: semantics.ActivationKind,
368 ) !?kernel_library.CatalogDescriptor {
369 var arena_state = alloc_arena.Arena.init(allocator);
370 defer arena_state.deinit();
371 const arena = arena_state.allocator();
372 const tensor_type = try dialect_mod.decodeTensorType(arena, result_type);
373 if (tensor_type.dtype != .f32 or tensor_type.dims.len != 1) return null;
374 if (tensor_type.dims[0] <= 0) return null;
375 const extent = std.math.cast(u64, tensor_type.dims[0]) orelse return null;
376 const selected = kernel_selection.selectActivationCatalog(.{
377 .dtype = .f32,
378 .kind = kind,
379 .extent = extent,
380 }) orelse return null;
381 return selected.descriptor;
382 }
383
384 fn activationKindFromOp(op: *ir.Operation) ?semantics.ActivationKind {
385 const attr = op.getAttr("activation_kind") orelse return null;
386 if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.ActivationOp.activation_kind_attr_name)) return null;
387 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
388 return semantics.ActivationKind.fromName(dialect_attr.payload);
389 }
390
391 fn expandActivationToPrimitives(
392 rewriter: *rewrite.PatternRewriter,
393 result_type: ir.Type,
394 input: *ir.Value,
395 kind: semantics.ActivationKind,
396 ) !?*ir.Value {
397 return switch (kind) {
398 .relu => try expandRelu(rewriter, result_type, input),
399 .silu => try expandSilu(rewriter, result_type, input),
400 .gelu => try expandGelu(rewriter, result_type, input),
401 };
402 }
403
404 fn expandRelu(
405 rewriter: *rewrite.PatternRewriter,
406 result_type: ir.Type,
407 input: *ir.Value,
408 ) !?*ir.Value {
409 const zero = try splatFloatConstant(rewriter, result_type, 0.0) orelse return null;
410 const max = try dialect_mod.AccyDialect.MaxOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), input, zero);
411 _ = try rewriter.insert(max.op);
412 return max.getResult();
413 }
414
415 fn expandSilu(
416 rewriter: *rewrite.PatternRewriter,
417 result_type: ir.Type,
418 input: *ir.Value,
419 ) !?*ir.Value {
420 const one = try splatFloatConstant(rewriter, result_type, 1.0) orelse return null;
421 const neg = try dialect_mod.AccyDialect.NegOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), input);
422 _ = try rewriter.insert(neg.op);
423 const exponent = try dialect_mod.AccyDialect.ExpOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), neg.getResult());
424 _ = try rewriter.insert(exponent.op);
425 const denominator = try dialect_mod.AccyDialect.AddOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), one, exponent.getResult());
426 _ = try rewriter.insert(denominator.op);
427 const out = try dialect_mod.AccyDialect.DivOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), input, denominator.getResult());
428 _ = try rewriter.insert(out.op);
429 return out.getResult();
430 }
431
432 fn expandGelu(
433 rewriter: *rewrite.PatternRewriter,
434 result_type: ir.Type,
435 input: *ir.Value,
436 ) !?*ir.Value {
437 const c044715 = try splatFloatConstant(rewriter, result_type, 0.044715) orelse return null;
438 const c079788 = try splatFloatConstant(rewriter, result_type, 0.7978845608028654) orelse return null;
439 const one = try splatFloatConstant(rewriter, result_type, 1.0) orelse return null;
440 const half = try splatFloatConstant(rewriter, result_type, 0.5) orelse return null;
441
442 const x2 = try dialect_mod.AccyDialect.MulOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), input, input);
443 _ = try rewriter.insert(x2.op);
444 const x3 = try dialect_mod.AccyDialect.MulOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), x2.getResult(), input);
445 _ = try rewriter.insert(x3.op);
446 const scaled_cubic = try dialect_mod.AccyDialect.MulOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), c044715, x3.getResult());
447 _ = try rewriter.insert(scaled_cubic.op);
448 const inner_sum = try dialect_mod.AccyDialect.AddOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), input, scaled_cubic.getResult());
449 _ = try rewriter.insert(inner_sum.op);
450 const scaled = try dialect_mod.AccyDialect.MulOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), c079788, inner_sum.getResult());
451 _ = try rewriter.insert(scaled.op);
452 const activated = try dialect_mod.AccyDialect.TanhOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), scaled.getResult());
453 _ = try rewriter.insert(activated.op);
454 const bracket = try dialect_mod.AccyDialect.AddOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), one, activated.getResult());
455 _ = try rewriter.insert(bracket.op);
456 const half_x = try dialect_mod.AccyDialect.MulOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), half, input);
457 _ = try rewriter.insert(half_x.op);
458 const out = try dialect_mod.AccyDialect.MulOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), half_x.getResult(), bracket.getResult());
459 _ = try rewriter.insert(out.op);
460 return out.getResult();
461 }
462
463 fn splatFloatConstant(
464 rewriter: *rewrite.PatternRewriter,
465 result_type: ir.Type,
466 value: f64,
467 ) !?*ir.Value {
468 var arena_state = alloc_arena.Arena.init(rewriter.allocator);
469 defer arena_state.deinit();
470 const arena = arena_state.allocator();
471 const tensor_type = try dialect_mod.decodeTensorType(arena, result_type);
472 if (!tensor_type.dtype.isFloat()) return null;
473 const count = elementCount(tensor_type) orelse return null;
474 const payload_len = std.math.mul(usize, count, tensor_type.dtype.sizeOf()) catch return null;
475 const payload = try rewriter.allocator.alloc(u8, payload_len);
476 defer rewriter.allocator.free(payload);
477 fillFloatPayload(payload, tensor_type.dtype, value);
478 const constant = try dialect_mod.AccyDialect.ConstantOp.create(rewriter.ir_ctx, ir.Location.getUnknown(), payload, result_type);
479 _ = try rewriter.insert(constant.op);
480 return constant.getResult();
481 }
482
483 fn elementCount(tensor_type: semantics.Type) ?usize {
484 var count: usize = 1;
485 for (tensor_type.dims) |dim| {
486 const dim_usize = std.math.cast(usize, dim) orelse return null;
487 count = std.math.mul(usize, count, dim_usize) catch return null;
488 }
489 return count;
490 }
491
492 fn fillFloatPayload(payload: []u8, dtype: semantics.DType, value: f64) void {
493 const width = dtype.sizeOf();
494 var offset: usize = 0;
495 while (offset < payload.len) : (offset += width) {
496 switch (dtype) {
497 .f16 => {
498 var converted: f16 = @floatCast(value);
499 @memcpy(payload[offset..][0..@sizeOf(f16)], std.mem.asBytes(&converted));
500 },
501 .bf16 => {
502 const Bf16 = @as(semantics.DType, .bf16).ZigType();
503 var converted = Bf16.fromF32(@floatCast(value));
504 @memcpy(payload[offset..][0..@sizeOf(Bf16)], std.mem.asBytes(&converted));
505 },
506 .f32 => {
507 var converted: f32 = @floatCast(value);
508 @memcpy(payload[offset..][0..@sizeOf(f32)], std.mem.asBytes(&converted));
509 },
510 .f64 => {
511 var converted = value;
512 @memcpy(payload[offset..][0..@sizeOf(f64)], std.mem.asBytes(&converted));
513 },
514 else => unreachable,
515 }
516 }
517 }
518
519 fn constantDefiningOp(value: *ir.Value) ?*ir.Operation {
520 const def_any = value.getDefiningOp() orelse return null;
521 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
522 if (!std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) return null;
523 return def_op;
524 }
525
526 fn hasOnlyUseBy(value: *ir.Value, op: *ir.Operation) bool {
527 if (!value.hasOneUse()) return false;
528 const first = value.first_use orelse return false;
529 const owner: *ir.Operation = @ptrCast(@alignCast(first.owner));
530 return owner == op;
531 }
532
533 fn payloadIsZeroF32(payload: []const u8, extent: usize) bool {
534 if (payload.len != extent * @sizeOf(f32)) return false;
535 for (0..extent) |index| {
536 var value: f32 = undefined;
537 const start = index * @sizeOf(f32);
538 @memcpy(std.mem.asBytes(&value), payload[start..][0..@sizeOf(f32)]);
539 if (value != 0.0) return false;
540 }
541 return true;
542 }
543
544 const testing = std.testing;
545 const semantic = accy_choir.semantic;
546
547 const ActivationCase = struct {
548 dims: []const i64 = &.{8},
549 dtype: semantics.DType = .f32,
550 kind: semantics.ActivationKind = .gelu,
551 copies: u32 = 1,
552 library: library_preparation.KernelLibraryLowering = .disabled,
553 calls: bool = false,
554 maximum: bool = false,
555
556 fn module(self: ActivationCase) !*semantic.SemanticModule {
557 var limits = semantic.Builder.ContextLimits.standard;
558 var elements: usize = 1;
559 for (self.dims) |dim| elements = try std.math.mul(usize, elements, @intCast(dim));
560 const payload = try std.math.mul(usize, elements, self.dtype.sizeOf());
561 limits.attributes.payload_bytes += try std.math.mul(usize, payload, 4 * self.copies);
562 var builder = try semantic.Builder.init(testing.allocator, limits);
563 defer builder.deinit();
564 const tensor = try builder.tensor(self.dtype, self.dims);
565 var function = try builder.beginFunction("activation_accounting", &.{tensor}, &.{tensor});
566 var value = function.parameter(0);
567 for (0..self.copies) |index| {
568 if (self.maximum) {
569 const bytes = try testing.allocator.alloc(u8, payload);
570 defer testing.allocator.free(bytes);
571 @memset(bytes, 0);
572 const zero = try function.constant(tensor, bytes);
573 value = if (index % 2 == 0)
574 try function.max(value, zero)
575 else
576 try function.max(zero, value);
577 } else {
578 value = try function.activation(value, self.kind);
579 }
580 }
581 try function.return_(&.{value});
582 try function.finish();
583 return builder.finish();
584 }
585
586 fn check(self: ActivationCase, module_: *semantic.SemanticModule) !void {
587 try module_.verify();
588 const root = module_.choir_module;
589 try testing.expectEqual(
590 @as(usize, 0),
591 ir.inspection.countOperationsNamed(root, "accy.activation"),
592 );
593 try testing.expectEqual(
594 @as(usize, if (self.calls) self.copies else 0),
595 ir.inspection.countOperationsNamed(root, "accy.kernel_call"),
596 );
597 var constants: ConstantWitness = .{ .case = self };
598 _ = try root.walk(.{ .order = .pre_order }, &constants, ConstantWitness.visit);
599 const per_activation: u64 = if (self.kind == .gelu) 4 else 1;
600 const count: u64 = if (self.calls) 0 else self.copies * per_activation;
601 try testing.expectEqual(count, constants.count);
602 }
603 };
604
605 const ConstantWitness = struct {
606 case: ActivationCase,
607 count: u64 = 0,
608
609 fn visit(self: *ConstantWitness, op: *ir.Operation) !ir.WalkResult {
610 if (!std.mem.eql(u8, op.name.name, "accy.constant")) return .advance;
611 const payload = (dialect_mod.AccyDialect.ConstantOp{ .op = op }).getPayload().?;
612 var elements: usize = 1;
613 for (self.case.dims) |dim| elements *= @intCast(dim);
614 try testing.expectEqual(elements * self.case.dtype.sizeOf(), payload.len);
615 const value: f64 = switch (self.case.kind) {
616 .relu => 0,
617 .silu => 1,
618 .gelu => ([_]f64{ 0.044715, 0.7978845608028654, 1, 0.5 })[self.count % 4],
619 };
620 var bytes: [@sizeOf(f64)]u8 = undefined;
621 const scalar = bytes[0..self.case.dtype.sizeOf()];
622 fillFloatPayload(scalar, self.case.dtype, value);
623 for (0..elements) |index| {
624 try testing.expectEqualSlices(
625 u8,
626 scalar,
627 payload[index * scalar.len ..][0..scalar.len],
628 );
629 }
630 self.count += 1;
631 return .advance;
632 }
633 };
634
635 fn checkActivationAccounting(admitted: bool, constructor: u32) !void {
636 const allocator = testing.allocator;
637 const revision = choir.product.revision;
638 const fixture: ActivationCase = .{
639 .library = if (constructor == 0) .disabled else .enabled,
640 .calls = constructor != 0,
641 };
642 const module = try fixture.module();
643 defer module.deinit();
644 const root = module.choir_module;
645 const before = try choir.bytecode.encodeModule(allocator, root);
646 defer allocator.free(before);
647 const options: Options = .{ .kernel_library = fixture.library };
648 const bounds = try activationWork(.{ .operation = root, .state = &options });
649 var allowance = revision.WorkVector.uniform(1 << 40);
650 if (!admitted) allowance.structural_visits = bounds.work.structural_visits - 1;
651 const ledger = try revision.AccountingV1.create(allocator, .{
652 .allowance = allowance,
653 .workspace = 1 << 24,
654 .events = 4,
655 }, &.{.{ .name = activation_lowering_pass_name, .version = 1 }});
656 defer ledger.destroy();
657 var cache = try passes.AnalysisCache.initAccounted(
658 allocator,
659 null,
660 ledger,
661 .{ .context = module.context() },
662 0,
663 );
664 defer cache.deinit();
665 var manager = passes.PassManager.init(allocator);
666 defer manager.deinit();
667 try manager.addPass(switch (constructor) {
668 0 => activationLoweringPass(),
669 1 => activationLoweringPassWithOptions(&options),
670 2 => try activationLoweringPassFromOptions(allocator, .{
671 .assignments = &.{.{ .name = "kernel-library", .value = "enabled" }},
672 }),
673 else => unreachable,
674 });
675 const result = manager.runWithAnalysisCache(root, module.context(), &cache, .{});
676 try testing.expectEqual(if (admitted) passes.PassResult.success else .failure, result);
677 if (admitted) {
678 try ledger.producersComplete();
679 try testing.expect(!ledger.view().missing_work_contract);
680 try testing.expectEqual(@as(u64, 1), ledger.view().executed.counters.pass_runs);
681 try fixture.check(module);
682 } else {
683 try testing.expectEqual(.exhausted, ledger.view().outcome);
684 try testing.expectEqual(@as(u64, 0), manager.stats.pass_runs);
685 const after = try choir.bytecode.encodeModule(allocator, root);
686 defer allocator.free(after);
687 try testing.expectEqualSlices(u8, before, after);
688 }
689 }
690
691 test "activation lowering accounts each constructor and refuses before mutation" {
692 for (0..3) |constructor| {
693 try checkActivationAccounting(false, @intCast(constructor));
694 try checkActivationAccounting(true, @intCast(constructor));
695 }
696 }
697
698 fn checkActivationStorage(fixture: ActivationCase) !void {
699 const allocator = testing.allocator;
700 const module = try fixture.module();
701 defer module.deinit();
702 const options: Options = .{ .kernel_library = fixture.library };
703 const bounds = try activationWork(.{ .operation = module.choir_module, .state = &options });
704 const bytes = try allocator.alloc(u8, @intCast(bounds.workspace));
705 defer allocator.free(bytes);
706 var storage = @import("alloc_fixed").Tracked.init(bytes);
707 var cache = passes.AnalysisCache.init(allocator, null);
708 defer cache.deinit();
709 var context = passes.PassContext.init(module.choir_module, module.context(), allocator, &cache);
710 defer context.deinit();
711 context.allocator = storage.allocator();
712 defer context.allocator = allocator;
713 const result = runActivationLoweringWithOptions(&context, options);
714 try testing.expect(!storage.exhausted);
715 try testing.expectEqual(null, module.context().exhaustedSegment());
716 try testing.expectEqual(.success, result);
717 try testing.expect(storage.status().high_water_bytes <= bounds.workspace);
718 try testing.expect(storage.status().high_water_bytes > 0);
719 try fixture.check(module);
720 }
721
722 test "activation lowering scratch covers full tensor payloads and repeated expansion" {
723 const shapes = [_][]const i64{
724 &.{}, &.{0}, &.{20000}, &.{ 2, 3 }, &.{ 1, 1, 1, 1, 1, 1, 1, 1 },
725 };
726 for (shapes) |shape| {
727 for ([_]semantics.ActivationKind{ .relu, .silu, .gelu }) |kind| {
728 for ([_]u32{ 1, 17 }) |copies| {
729 try checkActivationStorage(.{ .dims = shape, .kind = kind, .copies = copies });
730 }
731 }
732 }
733 for ([_]semantics.DType{ .f16, .bf16, .f64 }) |dtype| {
734 try checkActivationStorage(.{ .dims = &.{20000}, .dtype = dtype });
735 }
736 }
737
738 test "activation lowering scratch covers selected catalog and primitive fallback" {
739 for ([_]semantics.ActivationKind{ .relu, .silu, .gelu }) |kind| {
740 try checkActivationStorage(.{
741 .kind = kind,
742 .copies = 17,
743 .library = .enabled,
744 .calls = true,
745 });
746 try checkActivationStorage(.{ .dims = &.{20000}, .kind = kind, .library = .enabled });
747 }
748 try checkActivationStorage(.{
749 .kind = .relu,
750 .copies = 17,
751 .library = .enabled,
752 .calls = true,
753 .maximum = true,
754 });
755 try checkActivationStorage(.{
756 .dims = &.{16},
757 .kind = .relu,
758 .library = .enabled,
759 .maximum = true,
760 });
761 }
762
763 test "activation lowering pass selects kernel library relu" {
764 const allocator = testing.allocator;
765
766 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
767 defer builder.deinit();
768 const f32_8 = try builder.tensor(.f32, &.{8});
769 const zeros = @as([8]f32, @splat(0.0));
770 var fb = try builder.beginFunction("activation_kernel_library_relu_lowering_pass", &.{f32_8}, &.{f32_8});
771 const zero = try fb.constant(f32_8, std.mem.sliceAsBytes(zeros[0..]));
772 const out = try fb.max(fb.parameter(0), zero);
773 try fb.return_(&.{out});
774 try fb.finish();
775 const module = try builder.finish();
776 defer module.deinit();
777
778 var options = Options{ .kernel_library = .enabled };
779 var pm = passes.PassManager.init(allocator);
780 defer pm.deinit();
781 try pm.addPass(activationLoweringPassWithOptions(&options));
782
783 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
784 try module.verify();
785 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MaxOp.operation_name));
786 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ConstantOp.operation_name));
787 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
788 }
789
790 test "activation lowering pass selects kernel library gelu activation op" {
791 const allocator = testing.allocator;
792
793 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
794 defer builder.deinit();
795 const f32_8 = try builder.tensor(.f32, &.{8});
796 var fb = try builder.beginFunction("activation_kernel_library_gelu_op_lowering_pass", &.{f32_8}, &.{f32_8});
797 const out = try fb.activation(fb.parameter(0), .gelu);
798 try fb.return_(&.{out});
799 try fb.finish();
800 const module = try builder.finish();
801 defer module.deinit();
802
803 var options = Options{ .kernel_library = .enabled };
804 var pm = passes.PassManager.init(allocator);
805 defer pm.deinit();
806 try pm.addPass(activationLoweringPassWithOptions(&options));
807
808 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
809 try module.verify();
810 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ActivationOp.operation_name));
811 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
812 }
813
814 test "activation lowering pass expands silu activation op by default" {
815 const allocator = testing.allocator;
816
817 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
818 defer builder.deinit();
819 const f32_8 = try builder.tensor(.f32, &.{8});
820 var fb = try builder.beginFunction("activation_default_silu_op_lowering_pass", &.{f32_8}, &.{f32_8});
821 const out = try fb.activation(fb.parameter(0), .silu);
822 try fb.return_(&.{out});
823 try fb.finish();
824 const module = try builder.finish();
825 defer module.deinit();
826
827 var pm = passes.PassManager.init(allocator);
828 defer pm.deinit();
829 try pm.addPass(activationLoweringPass());
830
831 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
832 try module.verify();
833 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ActivationOp.operation_name));
834 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
835 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.NegOp.operation_name));
836 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ExpOp.operation_name));
837 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DivOp.operation_name));
838 }
839
840 test "activation lowering pass keeps relu generic by default" {
841 const allocator = testing.allocator;
842
843 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
844 defer builder.deinit();
845 const f32_8 = try builder.tensor(.f32, &.{8});
846 const zeros = @as([8]f32, @splat(0.0));
847 var fb = try builder.beginFunction("activation_registered_relu_generic_lowering_pass", &.{f32_8}, &.{f32_8});
848 const zero = try fb.constant(f32_8, std.mem.sliceAsBytes(zeros[0..]));
849 const out = try fb.max(zero, fb.parameter(0));
850 try fb.return_(&.{out});
851 try fb.finish();
852 const module = try builder.finish();
853 defer module.deinit();
854
855 var pm = passes.PassManager.init(allocator);
856 defer pm.deinit();
857 try pm.addPass(activationLoweringPass());
858
859 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
860 try module.verify();
861 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MaxOp.operation_name));
862 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ConstantOp.operation_name));
863 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
864 }
865
866 test "activation lowering pass falls back when activation catalog shape is unavailable" {
867 const allocator = testing.allocator;
868
869 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
870 defer builder.deinit();
871 const f32_16 = try builder.tensor(.f32, &.{16});
872 var fb = try builder.beginFunction("activation_unavailable_gelu_op_lowering_pass", &.{f32_16}, &.{f32_16});
873 const out = try fb.activation(fb.parameter(0), .gelu);
874 try fb.return_(&.{out});
875 try fb.finish();
876 const module = try builder.finish();
877 defer module.deinit();
878
879 var options = Options{ .kernel_library = .enabled };
880 var pm = passes.PassManager.init(allocator);
881 defer pm.deinit();
882 try pm.addPass(activationLoweringPassWithOptions(&options));
883
884 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
885 try module.verify();
886 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ActivationOp.operation_name));
887 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
888 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.TanhOp.operation_name));
889 }
890
891 test "activation lowering pass rejects unavailable relu catalog shape" {
892 const allocator = testing.allocator;
893
894 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
895 defer builder.deinit();
896 const f32_16 = try builder.tensor(.f32, &.{16});
897 const zeros = @as([16]f32, @splat(0.0));
898 var fb = try builder.beginFunction("activation_unavailable_relu_lowering_pass", &.{f32_16}, &.{f32_16});
899 const zero = try fb.constant(f32_16, std.mem.sliceAsBytes(zeros[0..]));
900 const out = try fb.max(fb.parameter(0), zero);
901 try fb.return_(&.{out});
902 try fb.finish();
903 const module = try builder.finish();
904 defer module.deinit();
905
906 var options = Options{ .kernel_library = .enabled };
907 var pm = passes.PassManager.init(allocator);
908 defer pm.deinit();
909 try pm.addPass(activationLoweringPassWithOptions(&options));
910
911 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
912 try module.verify();
913 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MaxOp.operation_name));
914 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ConstantOp.operation_name));
915 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
916 }