Air.Legalize: revert to loops for scalarizations

I had tried unrolling the loops to avoid requiring the
`vector_store_elem` instruction, but it's arguably a problem to generate
O(N) code for an operation on `@Vector(N, T)`. In addition, that
lowering emitted a lot of `.aggregate_init` instructions, which is
itself a quite difficult operation to codegen.

This requires reintroducing runtime vector indexing internally. However,
I've put it in a couple of instructions which are intended only for use
by `Air.Legalize`, named `legalize_vec_elem_val` (like `array_elem_val`,
but for indexing a vector with a runtime-known index) and
`legalize_vec_store_elem` (like the old `vector_store_elem`
instruction). These are explicitly documented as *not* being emitted by
Sema, so need only be implemented by backends if they actually use an
`Air.Legalize.Feature` which emits them (otherwise they can be marked as
`unreachable`).

Author: mlugg

src/Air.zig

@@ -660,8 +660,8 @@ pub const Inst = struct {
         /// Given a pointer to a slice, return a pointer to the pointer of the slice.
         /// Uses the `ty_op` field.
         ptr_slice_ptr_ptr,
-        /// Given an (array value or vector value) and element index,
-        /// return the element value at that index.
+        /// Given an (array value or vector value) and element index, return the element value at
+        /// that index. If the lhs is a vector value, the index is guaranteed to be comptime-known.
         /// Result type is the element type of the array operand.
         /// Uses the `bin_op` field.
         array_elem_val,
@@ -915,6 +915,26 @@ pub const Inst = struct {
         /// Operand is unused and set to Ref.none
         work_group_id,
 
+        // The remaining instructions are not emitted by Sema. They are only emitted by `Legalize`,
+        // depending on the enabled features. As such, backends can consider them `unreachable` if
+        // they do not enable the relevant legalizations.
+
+        /// Given a pointer to a vector, a runtime-known index, and a scalar value, store the value
+        /// into the vector at the given index. Zig does not support this operation, but `Legalize`
+        /// may emit it when scalarizing vector operations.
+        ///
+        /// Uses the `pl_op` field with payload `Bin`. `operand` is the vector pointer. `lhs` is the
+        /// element index of type `usize`. `rhs` is the element value. Result is always void.
+        legalize_vec_store_elem,
+        /// Given a vector value and a runtime-known index, return the element value at that index.
+        /// This instruction is similar to `array_elem_val`; the only difference is that the index
+        /// here is runtime-known, which is usually not allowed for vectors. `Legalize` may emit
+        /// this instruction when scalarizing vector operations.
+        ///
+        /// Uses the `bin_op` field. `lhs` is the vector pointer. `rhs` is the element index. Result
+        /// type is the vector element type.
+        legalize_vec_elem_val,
+
         pub fn fromCmpOp(op: std.math.CompareOperator, optimized: bool) Tag {
             switch (op) {
                 .lt => return if (optimized) .cmp_lt_optimized else .cmp_lt,
@@ -1681,6 +1701,7 @@ pub fn typeOfIndex(air: *const Air, inst: Air.Inst.Index, ip: *const InternPool)
         .prefetch,
         .set_err_return_trace,
         .c_va_end,
+        .legalize_vec_store_elem,
         => return .void,
 
         .slice_len,
@@ -1699,7 +1720,7 @@ pub fn typeOfIndex(air: *const Air, inst: Air.Inst.Index, ip: *const InternPool)
             return .fromInterned(ip.funcTypeReturnType(callee_ty.toIntern()));
         },
 
-        .slice_elem_val, .ptr_elem_val, .array_elem_val => {
+        .slice_elem_val, .ptr_elem_val, .array_elem_val, .legalize_vec_elem_val => {
             const ptr_ty = air.typeOf(datas[@intFromEnum(inst)].bin_op.lhs, ip);
             return ptr_ty.childTypeIp(ip);
         },
@@ -1857,6 +1878,7 @@ pub fn mustLower(air: Air, inst: Air.Inst.Index, ip: *const InternPool) bool {
         .intcast_safe,
         .int_from_float_safe,
         .int_from_float_optimized_safe,
+        .legalize_vec_store_elem,
         => true,
 
         .add,
@@ -2002,6 +2024,7 @@ pub fn mustLower(air: Air, inst: Air.Inst.Index, ip: *const InternPool) bool {
         .work_item_id,
         .work_group_size,
         .work_group_id,
+        .legalize_vec_elem_val,
         => false,
 
         .is_non_null_ptr, .is_null_ptr, .is_non_err_ptr, .is_err_ptr => air.typeOf(data.un_op, ip).isVolatilePtrIp(ip),

src/Air/Legalize.zig

@@ -458,6 +458,7 @@ fn analyzeInst(
         .memset_safe,
         .memcpy,
         .memmove,
+        .legalize_vec_elem_val,
         => {
             const o = inst_datas[@intFromEnum(inst)].bin_op;
             return analyzeOperands(a, pass, data, inst, .{ o.lhs, o.rhs, .none });
@@ -769,6 +770,12 @@ fn analyzeInst(
             const pl_op = inst_datas[@intFromEnum(inst)].pl_op;
             return analyzeOperands(a, pass, data, inst, .{ pl_op.operand, .none, .none });
         },
+
+        .legalize_vec_store_elem => {
+            const pl_op = inst_datas[@intFromEnum(inst)].pl_op;
+            const bin = a.air.extraData(Air.Bin, pl_op.payload).data;
+            return analyzeOperands(a, pass, data, inst, .{ pl_op.operand, bin.lhs, bin.rhs });
+        },
     }
 }
 

src/Air/Liveness.zig

@@ -272,6 +272,7 @@ fn verifyBody(self: *Verify, body: []const Air.Inst.Index) Error!void {
             .memset_safe,
             .memcpy,
             .memmove,
+            .legalize_vec_elem_val,
             => {
                 const bin_op = data[@intFromEnum(inst)].bin_op;
                 try self.verifyInstOperands(inst, .{ bin_op.lhs, bin_op.rhs, .none });
@@ -577,6 +578,11 @@ fn verifyBody(self: *Verify, body: []const Air.Inst.Index) Error!void {
 
                 try self.verifyInst(inst);
             },
+            .legalize_vec_store_elem => {
+                const pl_op = data[@intFromEnum(inst)].pl_op;
+                const bin = self.air.extraData(Air.Bin, pl_op.payload).data;
+                try self.verifyInstOperands(inst, .{ pl_op.operand, bin.lhs, bin.rhs });
+            },
         }
     }
 }

src/Air/Liveness/Verify.zig

@@ -171,6 +171,7 @@ const Writer = struct {
             .memmove,
             .memset,
             .memset_safe,
+            .legalize_vec_elem_val,
             => try w.writeBinOp(s, inst),
 
             .is_null,
@@ -331,6 +332,7 @@ const Writer = struct {
             .reduce, .reduce_optimized => try w.writeReduce(s, inst),
             .cmp_vector, .cmp_vector_optimized => try w.writeCmpVector(s, inst),
             .runtime_nav_ptr => try w.writeRuntimeNavPtr(s, inst),
+            .legalize_vec_store_elem => try w.writeLegalizeVecStoreElem(s, inst),
 
             .work_item_id,
             .work_group_size,
@@ -508,6 +510,18 @@ const Writer = struct {
         try w.writeOperand(s, inst, 2, pl_op.operand);
     }
 
+    fn writeLegalizeVecStoreElem(w: *Writer, s: *std.Io.Writer, inst: Air.Inst.Index) Error!void {
+        const pl_op = w.air.instructions.items(.data)[@intFromEnum(inst)].pl_op;
+        const bin = w.air.extraData(Air.Bin, pl_op.payload).data;
+
+        try w.writeOperand(s, inst, 0, pl_op.operand);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 1, bin.lhs);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 2, bin.rhs);
+        try s.writeAll(", ");
+    }
+
     fn writeShuffleOne(w: *Writer, s: *std.Io.Writer, inst: Air.Inst.Index) Error!void {
         const unwrapped = w.air.unwrapShuffleOne(w.pt.zcu, inst);
         try w.writeType(s, unwrapped.result_ty);

src/Air/print.zig

@@ -88,6 +88,7 @@ fn checkBody(air: Air, body: []const Air.Inst.Index, zcu: *Zcu) bool {
             .atomic_store_monotonic,
             .atomic_store_release,
             .atomic_store_seq_cst,
+            .legalize_vec_elem_val,
             => {
                 if (!checkRef(data.bin_op.lhs, zcu)) return false;
                 if (!checkRef(data.bin_op.rhs, zcu)) return false;
@@ -322,6 +323,7 @@ fn checkBody(air: Air, body: []const Air.Inst.Index, zcu: *Zcu) bool {
 
             .select,
             .mul_add,
+            .legalize_vec_store_elem,
             => {
                 const bin = air.extraData(Air.Bin, data.pl_op.payload).data;
                 if (!checkRef(data.pl_op.operand, zcu)) return false;

src/Air/types_resolved.zig

@@ -15930,16 +15930,21 @@ fn zirOverflowArithmetic(
                     }
                 }
                 // If either of the arguments is one, the result is the other and no overflow occured.
-                const scalar_one = try pt.intValue(dest_ty.scalarType(zcu), 1);
-                const vec_one = try sema.splat(dest_ty, scalar_one);
-                if (maybe_lhs_val) |lhs_val| {
-                    if (!lhs_val.isUndef(zcu) and try sema.compareAll(lhs_val, .eq, vec_one, dest_ty)) {
-                        break :result .{ .overflow_bit = try sema.splat(overflow_ty, .zero_u1), .inst = rhs };
+                const dest_scalar_ty = dest_ty.scalarType(zcu);
+                const dest_scalar_int = dest_scalar_ty.intInfo(zcu);
+                // We could still be working with i1, where '1' is not a legal value!
+                if (!(dest_scalar_int.bits == 1 and dest_scalar_int.signedness == .signed)) {
+                    const scalar_one = try pt.intValue(dest_scalar_ty, 1);
+                    const vec_one = try sema.splat(dest_ty, scalar_one);
+                    if (maybe_lhs_val) |lhs_val| {
+                        if (!lhs_val.isUndef(zcu) and try sema.compareAll(lhs_val, .eq, vec_one, dest_ty)) {
+                            break :result .{ .overflow_bit = try sema.splat(overflow_ty, .zero_u1), .inst = rhs };
+                        }
                     }
-                }
-                if (maybe_rhs_val) |rhs_val| {
-                    if (!rhs_val.isUndef(zcu) and try sema.compareAll(rhs_val, .eq, vec_one, dest_ty)) {
-                        break :result .{ .overflow_bit = try sema.splat(overflow_ty, .zero_u1), .inst = lhs };
+                    if (maybe_rhs_val) |rhs_val| {
+                        if (!rhs_val.isUndef(zcu) and try sema.compareAll(rhs_val, .eq, vec_one, dest_ty)) {
+                            break :result .{ .overflow_bit = try sema.splat(overflow_ty, .zero_u1), .inst = lhs };
+                        }
                     }
                 }
 

src/Sema.zig

@@ -134,6 +134,10 @@ pub fn analyze(isel: *Select, air_body: []const Air.Inst.Index) !void {
     var air_inst_index = air_body[air_body_index];
     const initial_def_order_len = isel.def_order.count();
     air_tag: switch (air_tags[@intFromEnum(air_inst_index)]) {
+        // No "scalarize" legalizations are enabled, so these instructions never appear.
+        .legalize_vec_elem_val => unreachable,
+        .legalize_vec_store_elem => unreachable,
+
         .arg,
         .ret_addr,
         .frame_addr,
@@ -950,6 +954,11 @@ pub fn body(isel: *Select, air_body: []const Air.Inst.Index) error{ OutOfMemory,
     };
     air_tag: switch (air.next().?) {
         else => |air_tag| return isel.fail("unimplemented {t}", .{air_tag}),
+
+        // No "scalarize" legalizations are enabled, so these instructions never appear.
+        .legalize_vec_elem_val => unreachable,
+        .legalize_vec_store_elem => unreachable,
+
         .arg => {
             const arg_vi = isel.live_values.fetchRemove(air.inst_index).?.value;
             defer arg_vi.deref(isel);

src/codegen/aarch64/Select.zig

@@ -3325,6 +3325,10 @@ fn genBodyInner(f: *Function, body: []const Air.Inst.Index) Error!void {
             // zig fmt: off
             .inferred_alloc, .inferred_alloc_comptime => unreachable,
 
+            // No "scalarize" legalizations are enabled, so these instructions never appear.
+            .legalize_vec_elem_val   => unreachable,
+            .legalize_vec_store_elem => unreachable,
+
             .arg      => try airArg(f, inst),
 
             .breakpoint => try airBreakpoint(f),

src/codegen/c.zig

@@ -4886,6 +4886,11 @@ pub const FuncGen = struct {
 
             const val: Builder.Value = switch (air_tags[@intFromEnum(inst)]) {
                 // zig fmt: off
+
+                // No "scalarize" legalizations are enabled, so these instructions never appear.
+                .legalize_vec_elem_val   => unreachable,
+                .legalize_vec_store_elem => unreachable,
+
                 .add            => try self.airAdd(inst, .normal),
                 .add_optimized  => try self.airAdd(inst, .fast),
                 .add_wrap       => try self.airAddWrap(inst),