lib/choir/src/backends/aarch64/control.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const ir = @import("../../core/root.zig");
 3 const dialects = @import("../../dialects/root.zig");
 4 const scf = dialects.ScfDialect;
 5 
 6 pub const max_depth: usize = 32;
 7 pub const max_carried_values: usize = 16;
 8 pub const max_transfers: usize = max_carried_values + 1;
 9 pub const Kind = enum { conditional, counted, repeated };
10 pub const Error = error{ ControlShape, ControlType, CarriedCapacity, DepthCapacity };
11 
12 pub fn kind(op: *const ir.Operation) ?Kind {
13     if (is(op, scf.IfOp.operation_name)) return .conditional;
14     if (is(op, scf.ForOp.operation_name)) return .counted;
15     if (is(op, scf.WhileOp.operation_name)) return .repeated;
16     return null;
17 }
18 
19 pub fn integer(typ: ir.Type) bool {
20     const name = typ.getDialectTypeName() orelse return false;
21     const scalar = dialects.arith.scalarKindFromTypeName(name) orelse return false;
22     return switch (scalar) {
23         .i8, .i16, .i32, .i64, .u8, .u16, .u32, .u64, .index, .bool => true,
24         else => false,
25     };
26 }
27 
28 fn named(typ: ir.Type, name: []const u8) bool {
29     return std.mem.eql(u8, typ.getDialectTypeName() orelse "", name);
30 }
31 
32 /// Admission establishes all region/edge shapes before placement or emission.
33 pub fn validate(op: *ir.Operation, construct: Kind) Error!void {
34     const count = op.results.items.len;
35     if (count > max_carried_values) return error.CarriedCapacity;
36     for (op.results.items) |result| if (!integer(result.type)) return error.ControlType;
37     const regions = op.regions.items;
38     const expected_regions: usize = if (construct == .counted) 1 else 2;
39     if (regions.len != expected_regions) {
40         if (!(construct == .conditional and count == 0 and regions.len == 1)) return error.ControlShape;
41     }
42     for (regions) |*region| {
43         if (region.blocks.size != 1) return error.ControlShape;
44         const block = region.getEntryBlock() orelse return error.ControlShape;
45         if (block.operations.tail == null) return error.ControlShape;
46         for (block.arguments.items) |arg| if (!integer(arg.type)) return error.ControlType;
47     }
48     switch (construct) {
49         .conditional => {
50             if (op.operands.items.len != 1) return error.ControlShape;
51             if (!named(op.operands.items[0].value.type, "arith.bool")) return error.ControlType;
52             for (regions) |*region| {
53                 const block = region.getEntryBlock().?;
54                 if (block.arguments.items.len != 0) return error.ControlShape;
55                 try yieldTypes(last(block), op.results.items);
56             }
57         },
58         .counted => {
59             if (op.operands.items.len != count + 3) return error.ControlShape;
60             for (op.operands.items[0..3]) |operand| if (!named(operand.value.type, "arith.index")) return error.ControlType;
61             const block = regions[0].getEntryBlock().?;
62             if (block.arguments.items.len != count + 1) return error.ControlShape;
63             if (!named(block.arguments.items[0].type, "arith.index")) return error.ControlType;
64             for (op.results.items, op.operands.items[3..], block.arguments.items[1..]) |result, init, arg| {
65                 if (!result.type.eql(init.value.type) or !result.type.eql(arg.type)) return error.ControlType;
66             }
67             try yieldTypes(last(block), op.results.items);
68         },
69         .repeated => {
70             if (op.operands.items.len != count) return error.ControlShape;
71             for (regions) |*region| if (region.getEntryBlock().?.arguments.items.len != count) return error.ControlShape;
72             const before = regions[0].getEntryBlock().?;
73             const after = regions[1].getEntryBlock().?;
74             for (op.results.items, op.operands.items, before.arguments.items, after.arguments.items) |result, init, lhs, rhs| {
75                 if (!result.type.eql(init.value.type) or !result.type.eql(lhs.type) or !result.type.eql(rhs.type)) return error.ControlType;
76             }
77             const condition = last(before);
78             if (!is(condition, scf.ConditionOp.operation_name) or condition.operands.items.len != count + 1) return error.ControlShape;
79             if (!named(condition.operands.items[0].value.type, "arith.bool")) return error.ControlType;
80             for (op.results.items, condition.operands.items[1..]) |result, forwarded| if (!result.type.eql(forwarded.value.type)) return error.ControlType;
81             try yieldTypes(last(after), op.results.items);
82         },
83     }
84 }
85 
86 fn yieldTypes(op: *ir.Operation, results: []const ir.Value) Error!void {
87     if (!is(op, scf.YieldOp.operation_name) or op.operands.items.len != results.len) return error.ControlShape;
88     for (results, op.operands.items) |result, yielded| if (!result.type.eql(yielded.value.type)) return error.ControlType;
89 }
90 
91 /// Compare an operation's registered spelling without requiring registration.
92 pub fn is(op: *const ir.Operation, name: []const u8) bool {
93     return std.mem.eql(u8, op.name.name, name);
94 }
95 
96 pub fn last(block: *ir.Block) *ir.Operation {
97     return @ptrCast(@alignCast(block.operations.tail.?));
98 }