Legalize: introduce a new pass before liveness

Each target can opt into different sets of legalize features.
By performing these transformations before liveness, instructions
that become unreferenced will have up-to-date liveness information.

Author: jacobly0

CMakeLists.txt

@@ -512,13 +512,15 @@ set(ZIG_STAGE2_SOURCES
     lib/std/zig/llvm/bitcode_writer.zig
     lib/std/zig/llvm/ir.zig
     src/Air.zig
+    src/Air/Legalize.zig
+    src/Air/Liveness.zig
+    src/Air/Liveness/Verify.zig
+    src/Air/types_resolved.zig
     src/Builtin.zig
     src/Compilation.zig
     src/Compilation/Config.zig
     src/DarwinPosixSpawn.zig
     src/InternPool.zig
-    src/Liveness.zig
-    src/Liveness/Verify.zig
     src/Package.zig
     src/Package/Fetch.zig
     src/Package/Fetch/git.zig

src/Air.zig

@@ -9,16 +9,19 @@ const builtin = @import("builtin");
 const assert = std.debug.assert;
 
 const Air = @This();
-const Value = @import("Value.zig");
-const Type = @import("Type.zig");
 const InternPool = @import("InternPool.zig");
+const Type = @import("Type.zig");
+const Value = @import("Value.zig");
 const Zcu = @import("Zcu.zig");
 const types_resolved = @import("Air/types_resolved.zig");
 
+pub const Legalize = @import("Air/Legalize.zig");
+pub const Liveness = @import("Air/Liveness.zig");
+
 instructions: std.MultiArrayList(Inst).Slice,
 /// The meaning of this data is determined by `Inst.Tag` value.
 /// The first few indexes are reserved. See `ExtraIndex` for the values.
-extra: []const u32,
+extra: std.ArrayListUnmanaged(u32),
 
 pub const ExtraIndex = enum(u32) {
     /// Payload index of the main `Block` in the `extra` array.
@@ -244,22 +247,27 @@ pub const Inst = struct {
         /// Uses the `bin_op` field.
         bit_or,
         /// Shift right. `>>`
+        /// The rhs type may be a scalar version of the lhs type.
         /// Uses the `bin_op` field.
         shr,
         /// Shift right. The shift produces a poison value if it shifts out any non-zero bits.
+        /// The rhs type may be a scalar version of the lhs type.
         /// Uses the `bin_op` field.
         shr_exact,
         /// Shift left. `<<`
+        /// The rhs type may be a scalar version of the lhs type.
         /// Uses the `bin_op` field.
         shl,
         /// Shift left; For unsigned integers, the shift produces a poison value if it shifts
         /// out any non-zero bits. For signed integers, the shift produces a poison value if
         /// it shifts out any bits that disagree with the resultant sign bit.
+        /// The rhs type may be a scalar version of the lhs type.
         /// Uses the `bin_op` field.
         shl_exact,
         /// Saturating integer shift left. `<<|`. The result is the same type as the `lhs`.
         /// The `rhs` must have the same vector shape as the `lhs`, but with any unsigned
         /// integer as the scalar type.
+        /// The rhs type may be a scalar version of the lhs type.
         /// Uses the `bin_op` field.
         shl_sat,
         /// Bitwise XOR. `^`
@@ -1378,9 +1386,9 @@ pub const UnionInit = struct {
 };
 
 pub fn getMainBody(air: Air) []const Air.Inst.Index {
-    const body_index = air.extra[@intFromEnum(ExtraIndex.main_block)];
+    const body_index = air.extra.items[@intFromEnum(ExtraIndex.main_block)];
     const extra = air.extraData(Block, body_index);
-    return @ptrCast(air.extra[extra.end..][0..extra.data.body_len]);
+    return @ptrCast(air.extra.items[extra.end..][0..extra.data.body_len]);
 }
 
 pub fn typeOf(air: *const Air, inst: Air.Inst.Ref, ip: *const InternPool) Type {
@@ -1656,9 +1664,9 @@ pub fn extraData(air: Air, comptime T: type, index: usize) struct { data: T, end
     var result: T = undefined;
     inline for (fields) |field| {
         @field(result, field.name) = switch (field.type) {
-            u32 => air.extra[i],
-            InternPool.Index, Inst.Ref => @enumFromInt(air.extra[i]),
-            i32, CondBr.BranchHints => @bitCast(air.extra[i]),
+            u32 => air.extra.items[i],
+            InternPool.Index, Inst.Ref => @enumFromInt(air.extra.items[i]),
+            i32, CondBr.BranchHints => @bitCast(air.extra.items[i]),
             else => @compileError("bad field type: " ++ @typeName(field.type)),
         };
         i += 1;
@@ -1671,7 +1679,7 @@ pub fn extraData(air: Air, comptime T: type, index: usize) struct { data: T, end
 
 pub fn deinit(air: *Air, gpa: std.mem.Allocator) void {
     air.instructions.deinit(gpa);
-    gpa.free(air.extra);
+    air.extra.deinit(gpa);
     air.* = undefined;
 }
 
@@ -1700,7 +1708,7 @@ pub const NullTerminatedString = enum(u32) {
 
     pub fn toSlice(nts: NullTerminatedString, air: Air) [:0]const u8 {
         if (nts == .none) return "";
-        const bytes = std.mem.sliceAsBytes(air.extra[@intFromEnum(nts)..]);
+        const bytes = std.mem.sliceAsBytes(air.extra.items[@intFromEnum(nts)..]);
         return bytes[0..std.mem.indexOfScalar(u8, bytes, 0).? :0];
     }
 };
@@ -1943,7 +1951,7 @@ pub const UnwrappedSwitch = struct {
         return us.getHintInner(us.cases_len);
     }
     fn getHintInner(us: UnwrappedSwitch, idx: u32) std.builtin.BranchHint {
-        const bag = us.air.extra[us.branch_hints_start..][idx / 10];
+        const bag = us.air.extra.items[us.branch_hints_start..][idx / 10];
         const bits: u3 = @truncate(bag >> @intCast(3 * (idx % 10)));
         return @enumFromInt(bits);
     }
@@ -1971,13 +1979,13 @@ pub const UnwrappedSwitch = struct {
 
             const extra = it.air.extraData(SwitchBr.Case, it.extra_index);
             var extra_index = extra.end;
-            const items: []const Inst.Ref = @ptrCast(it.air.extra[extra_index..][0..extra.data.items_len]);
+            const items: []const Inst.Ref = @ptrCast(it.air.extra.items[extra_index..][0..extra.data.items_len]);
             extra_index += items.len;
             // TODO: ptrcast from []const Inst.Ref to []const [2]Inst.Ref when supported
-            const ranges_ptr: [*]const [2]Inst.Ref = @ptrCast(it.air.extra[extra_index..]);
+            const ranges_ptr: [*]const [2]Inst.Ref = @ptrCast(it.air.extra.items[extra_index..]);
             const ranges: []const [2]Inst.Ref = ranges_ptr[0..extra.data.ranges_len];
             extra_index += ranges.len * 2;
-            const body: []const Inst.Index = @ptrCast(it.air.extra[extra_index..][0..extra.data.body_len]);
+            const body: []const Inst.Index = @ptrCast(it.air.extra.items[extra_index..][0..extra.data.body_len]);
             extra_index += body.len;
             it.extra_index = @intCast(extra_index);
 
@@ -1992,7 +2000,7 @@ pub const UnwrappedSwitch = struct {
         /// Returns the body of the "default" (`else`) case.
         pub fn elseBody(it: *CaseIterator) []const Inst.Index {
             assert(it.next_case == it.cases_len);
-            return @ptrCast(it.air.extra[it.extra_index..][0..it.else_body_len]);
+            return @ptrCast(it.air.extra.items[it.extra_index..][0..it.else_body_len]);
         }
         pub const Case = struct {
             idx: u32,
@@ -2025,6 +2033,7 @@ pub fn unwrapSwitch(air: *const Air, switch_inst: Inst.Index) UnwrappedSwitch {
 pub const typesFullyResolved = types_resolved.typesFullyResolved;
 pub const typeFullyResolved = types_resolved.checkType;
 pub const valFullyResolved = types_resolved.checkVal;
+pub const legalize = Legalize.legalize;
 
 pub const CoveragePoint = enum(u1) {
     /// Indicates the block is not a place of interest corresponding to

src/Air/Legalize.zig

@@ -0,0 +1,147 @@
+zcu: *const Zcu,
+air: Air,
+features: std.enums.EnumSet(Feature),
+
+pub const Feature = enum {
+    /// Legalize (shift lhs, (splat rhs)) -> (shift lhs, rhs)
+    remove_shift_vector_rhs_splat,
+    /// Legalize reduce of a one element vector to a bitcast
+    reduce_one_elem_to_bitcast,
+};
+
+pub const Features = std.enums.EnumFieldStruct(Feature, bool, false);
+
+pub fn legalize(air: *Air, backend: std.builtin.CompilerBackend, zcu: *const Zcu) std.mem.Allocator.Error!void {
+    var l: Legalize = .{
+        .zcu = zcu,
+        .air = air.*,
+        .features = features: switch (backend) {
+            .other, .stage1 => unreachable,
+            inline .stage2_llvm,
+            .stage2_c,
+            .stage2_wasm,
+            .stage2_arm,
+            .stage2_x86_64,
+            .stage2_aarch64,
+            .stage2_x86,
+            .stage2_riscv64,
+            .stage2_sparc64,
+            .stage2_spirv64,
+            .stage2_powerpc,
+            => |ct_backend| {
+                const Backend = codegen.importBackend(ct_backend) orelse break :features .initEmpty();
+                break :features if (@hasDecl(Backend, "legalize_features"))
+                    .init(Backend.legalize_features)
+                else
+                    .initEmpty();
+            },
+            _ => unreachable,
+        },
+    };
+    defer air.* = l.air;
+    if (!l.features.bits.eql(.initEmpty())) try l.legalizeBody(l.air.getMainBody());
+}
+
+fn legalizeBody(l: *Legalize, body: []const Air.Inst.Index) std.mem.Allocator.Error!void {
+    const zcu = l.zcu;
+    const ip = &zcu.intern_pool;
+    const tags = l.air.instructions.items(.tag);
+    const data = l.air.instructions.items(.data);
+    for (body) |inst| inst: switch (tags[@intFromEnum(inst)]) {
+        else => {},
+
+        .shl,
+        .shl_exact,
+        .shl_sat,
+        .shr,
+        .shr_exact,
+        => |air_tag| if (l.features.contains(.remove_shift_vector_rhs_splat)) done: {
+            const bin_op = data[@intFromEnum(inst)].bin_op;
+            const ty = l.air.typeOf(bin_op.rhs, ip);
+            if (!ty.isVector(zcu)) break :done;
+            if (bin_op.rhs.toInterned()) |rhs_ip_index| switch (ip.indexToKey(rhs_ip_index)) {
+                else => {},
+                .aggregate => |aggregate| switch (aggregate.storage) {
+                    else => {},
+                    .repeated_elem => |splat| continue :inst l.replaceInst(inst, air_tag, .{ .bin_op = .{
+                        .lhs = bin_op.lhs,
+                        .rhs = Air.internedToRef(splat),
+                    } }),
+                },
+            } else {
+                const rhs_inst = bin_op.rhs.toIndex().?;
+                switch (tags[@intFromEnum(rhs_inst)]) {
+                    else => {},
+                    .splat => continue :inst l.replaceInst(inst, air_tag, .{ .bin_op = .{
+                        .lhs = bin_op.lhs,
+                        .rhs = data[@intFromEnum(rhs_inst)].ty_op.operand,
+                    } }),
+                }
+            }
+        },
+
+        .reduce,
+        .reduce_optimized,
+        => if (l.features.contains(.reduce_one_elem_to_bitcast)) done: {
+            const reduce = data[@intFromEnum(inst)].reduce;
+            const vector_ty = l.air.typeOf(reduce.operand, ip);
+            switch (vector_ty.vectorLen(zcu)) {
+                0 => unreachable,
+                1 => continue :inst l.replaceInst(inst, .bitcast, .{ .ty_op = .{
+                    .ty = Air.internedToRef(vector_ty.scalarType(zcu).toIntern()),
+                    .operand = reduce.operand,
+                } }),
+                else => break :done,
+            }
+        },
+
+        .@"try", .try_cold => {
+            const pl_op = data[@intFromEnum(inst)].pl_op;
+            const extra = l.air.extraData(Air.Try, pl_op.payload);
+            try l.legalizeBody(@ptrCast(l.air.extra.items[extra.end..][0..extra.data.body_len]));
+        },
+        .try_ptr, .try_ptr_cold => {
+            const ty_pl = data[@intFromEnum(inst)].ty_pl;
+            const extra = l.air.extraData(Air.TryPtr, ty_pl.payload);
+            try l.legalizeBody(@ptrCast(l.air.extra.items[extra.end..][0..extra.data.body_len]));
+        },
+        .block, .loop => {
+            const ty_pl = data[@intFromEnum(inst)].ty_pl;
+            const extra = l.air.extraData(Air.Block, ty_pl.payload);
+            try l.legalizeBody(@ptrCast(l.air.extra.items[extra.end..][0..extra.data.body_len]));
+        },
+        .dbg_inline_block => {
+            const ty_pl = data[@intFromEnum(inst)].ty_pl;
+            const extra = l.air.extraData(Air.DbgInlineBlock, ty_pl.payload);
+            try l.legalizeBody(@ptrCast(l.air.extra.items[extra.end..][0..extra.data.body_len]));
+        },
+        .cond_br => {
+            const pl_op = data[@intFromEnum(inst)].pl_op;
+            const extra = l.air.extraData(Air.CondBr, pl_op.payload);
+            try l.legalizeBody(@ptrCast(l.air.extra.items[extra.end..][0..extra.data.then_body_len]));
+            try l.legalizeBody(@ptrCast(l.air.extra.items[extra.end + extra.data.then_body_len ..][0..extra.data.else_body_len]));
+        },
+        .switch_br, .loop_switch_br => {
+            const switch_br = l.air.unwrapSwitch(inst);
+            var it = switch_br.iterateCases();
+            while (it.next()) |case| try l.legalizeBody(case.body);
+            try l.legalizeBody(it.elseBody());
+        },
+    };
+}
+
+// inline to propagate comptime `tag`s
+inline fn replaceInst(l: *Legalize, inst: Air.Inst.Index, tag: Air.Inst.Tag, data: Air.Inst.Data) Air.Inst.Tag {
+    const ip = &l.zcu.intern_pool;
+    const orig_ty = if (std.debug.runtime_safety) l.air.typeOfIndex(inst, ip) else {};
+    l.air.instructions.items(.tag)[@intFromEnum(inst)] = tag;
+    l.air.instructions.items(.data)[@intFromEnum(inst)] = data;
+    if (std.debug.runtime_safety) std.debug.assert(l.air.typeOfIndex(inst, ip).toIntern() == orig_ty.toIntern());
+    return tag;
+}
+
+const Air = @import("../Air.zig");
+const codegen = @import("../codegen.zig");
+const Legalize = @This();
+const std = @import("std");
+const Zcu = @import("../Zcu.zig");