stage2: implement @select

src/Air.zig

@@ -344,7 +344,7 @@ pub const Inst = struct {
         /// to the storage for the variable. The local may be a const or a var.
         /// Result type is always void.
         /// Uses `pl_op`. The payload index is the variable name. It points to the extra
-        /// array, reinterpreting the bytes there as a null-terminated string. 
+        /// array, reinterpreting the bytes there as a null-terminated string.
         dbg_var_ptr,
         /// Same as `dbg_var_ptr` except the local is a const, not a var, and the
         /// operand is the local's value.
@@ -553,6 +553,9 @@ pub const Inst = struct {
         /// Constructs a vector by selecting elements from `a` and `b` based on `mask`.
         /// Uses the `ty_pl` field with payload `Shuffle`.
         shuffle,
+        /// Constructs a vector element-wise from `a` or `b` based on `pred`.
+        /// Uses the `ty_pl` field with payload `Select`.
+        select,
 
         /// Given dest ptr, value, and len, set all elements at dest to value.
         /// Result type is always void.
@@ -785,6 +788,12 @@ pub const Shuffle = struct {
     mask_len: u32,
 };
 
+pub const Select = struct {
+    pred: Inst.Ref,
+    a: Inst.Ref,
+    b: Inst.Ref,
+};
+
 pub const VectorCmp = struct {
     lhs: Inst.Ref,
     rhs: Inst.Ref,
@@ -956,6 +965,7 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
         .cmpxchg_weak,
         .cmpxchg_strong,
         .slice,
+        .select,
         .shuffle,
         .aggregate_init,
         .union_init,

src/Liveness.zig

@@ -433,6 +433,10 @@ fn analyzeInst(
             }
             return extra_tombs.finish();
         },
+        .select => {
+            const extra = a.air.extraData(Air.Select, inst_datas[inst].ty_pl.payload).data;
+            return trackOperands(a, new_set, inst, main_tomb, .{ extra.pred, extra.a, extra.b });
+        },
         .shuffle => {
             const extra = a.air.extraData(Air.Shuffle, inst_datas[inst].ty_pl.payload).data;
             return trackOperands(a, new_set, inst, main_tomb, .{ extra.a, extra.b, .none });

src/Sema.zig

@@ -14804,8 +14804,105 @@ fn analyzeShuffle(
 
 fn zirSelect(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
     const inst_data = sema.code.instructions.items(.data)[inst].pl_node;
-    const src = inst_data.src();
+    const extra = sema.code.extraData(Zir.Inst.Select, inst_data.payload_index).data;
-    return sema.fail(block, src, "TODO: Sema.zirSelect", .{});
+
+    const elem_ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
+    const pred_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
+    const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = inst_data.src_node };
+    const b_src: LazySrcLoc = .{ .node_offset_builtin_call_arg3 = inst_data.src_node };
+
+    const elem_ty = try sema.resolveType(block, elem_ty_src, extra.elem_type);
+    try sema.checkVectorElemType(block, elem_ty_src, elem_ty);
+    const pred = sema.resolveInst(extra.pred);
+    const a = sema.resolveInst(extra.a);
+    const b = sema.resolveInst(extra.b);
+    const target = sema.mod.getTarget();
+
+    const pred_ty = sema.typeOf(pred);
+    switch (try pred_ty.zigTypeTagOrPoison()) {
+        .Vector => {
+            const scalar_ty = pred_ty.childType();
+            if (!scalar_ty.eql(Type.bool, target)) {
+                const bool_vec_ty = try Type.vector(sema.arena, pred_ty.vectorLen(), Type.bool);
+                return sema.fail(block, pred_src, "Expected '{}', found '{}'", .{ bool_vec_ty.fmt(target), pred_ty.fmt(target) });
+            }
+        },
+        else => return sema.fail(block, pred_src, "Expected vector type, found '{}'", .{pred_ty.fmt(target)}),
+    }
+
+    const vec_len = pred_ty.vectorLen();
+    const vec_ty = try Type.vector(sema.arena, vec_len, elem_ty);
+
+    const a_ty = sema.typeOf(a);
+    if (!a_ty.eql(vec_ty, target)) {
+        return sema.fail(block, a_src, "Expected '{}', found '{}'", .{ vec_ty.fmt(target), a_ty.fmt(target) });
+    }
+
+    const b_ty = sema.typeOf(b);
+    if (!b_ty.eql(vec_ty, target)) {
+        return sema.fail(block, b_src, "Expected '{}', found '{}'", .{ vec_ty.fmt(target), b_ty.fmt(target) });
+    }
+
+    const maybe_pred = try sema.resolveMaybeUndefVal(block, pred_src, pred);
+    const maybe_a = try sema.resolveMaybeUndefVal(block, a_src, a);
+    const maybe_b = try sema.resolveMaybeUndefVal(block, b_src, b);
+
+    const runtime_src = if (maybe_pred) |pred_val| rs: {
+        if (pred_val.isUndef()) return sema.addConstUndef(vec_ty);
+
+        if (maybe_a) |a_val| {
+            if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
+
+            if (maybe_b) |b_val| {
+                if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
+
+                var buf: Value.ElemValueBuffer = undefined;
+                const elems = try sema.gpa.alloc(Value, vec_len);
+                for (elems) |*elem, i| {
+                    const pred_elem_val = pred_val.elemValueBuffer(i, &buf);
+                    const should_choose_a = pred_elem_val.toBool();
+                    if (should_choose_a) {
+                        elem.* = a_val.elemValueBuffer(i, &buf);
+                    } else {
+                        elem.* = b_val.elemValueBuffer(i, &buf);
+                    }
+                }
+
+                return sema.addConstant(
+                    vec_ty,
+                    try Value.Tag.aggregate.create(sema.arena, elems),
+                );
+            } else {
+                break :rs b_src;
+            }
+        } else {
+            if (maybe_b) |b_val| {
+                if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
+            }
+            break :rs a_src;
+        }
+    } else rs: {
+        if (maybe_a) |a_val| {
+            if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
+        }
+        if (maybe_b) |b_val| {
+            if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
+        }
+        break :rs pred_src;
+    };
+
+    try sema.requireRuntimeBlock(block, runtime_src);
+    return block.addInst(.{
+        .tag = .select,
+        .data = .{ .ty_pl = .{
+            .ty = try block.sema.addType(vec_ty),
+            .payload = try block.sema.addExtra(Air.Select{
+                .pred = pred,
+                .a = a,
+                .b = b,
+            }),
+        } },
+    });
 }
 
 fn zirAtomicLoad(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {

src/arch/aarch64/CodeGen.zig

@@ -640,6 +640,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -3746,6 +3747,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+    const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for {}", .{self.target.cpu.arch});
+    return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for {}", .{self.target.cpu.arch});

src/arch/arm/CodeGen.zig

@@ -630,6 +630,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -4323,6 +4324,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+    const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for arm", .{});
+    return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for arm", .{});

src/arch/riscv64/CodeGen.zig

@@ -600,6 +600,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -2396,6 +2397,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+    const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for riscv64", .{});
+    return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for riscv64", .{});

src/arch/wasm/CodeGen.zig

@@ -1371,6 +1371,7 @@ fn genInst(self: *Self, inst: Air.Inst.Index) !WValue {
         .ret_ptr => self.airRetPtr(inst),
         .ret_load => self.airRetLoad(inst),
         .splat => self.airSplat(inst),
+        .select => self.airSelect(inst),
         .shuffle => self.airShuffle(inst),
         .reduce => self.airReduce(inst),
         .aggregate_init => self.airAggregateInit(inst),
@@ -3265,6 +3266,16 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
     return self.fail("TODO: Implement wasm airSplat", .{});
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
+    if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
+
+    const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+    const ty = try self.resolveInst(ty_pl.ty);
+
+    _ = ty;
+    return self.fail("TODO: Implement wasm airSelect", .{});
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
     if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
 

src/arch/x86_64/CodeGen.zig

@@ -721,6 +721,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -5678,6 +5679,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+    const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for x86_64", .{});
+    return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for x86_64", .{});

src/codegen/c.zig

@@ -1825,6 +1825,7 @@ fn genBody(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, OutO
             .tag_name         => try airTagName(f, inst),
             .error_name       => try airErrorName(f, inst),
             .splat            => try airSplat(f, inst),
+            .select           => try airSelect(f, inst),
             .shuffle          => try airShuffle(f, inst),
             .reduce           => try airReduce(f, inst),
             .aggregate_init   => try airAggregateInit(f, inst),
@@ -3794,6 +3795,21 @@ fn airSplat(f: *Function, inst: Air.Inst.Index) !CValue {
     return f.fail("TODO: C backend: implement airSplat", .{});
 }
 
+fn airSelect(f: *Function, inst: Air.Inst.Index) !CValue {
+    if (f.liveness.isUnused(inst)) return CValue.none;
+
+    const inst_ty = f.air.typeOfIndex(inst);
+    const ty_pl = f.air.instructions.items(.data)[inst].ty_pl;
+
+    const writer = f.object.writer();
+    const local = try f.allocLocal(inst_ty, .Const);
+    try writer.writeAll(" = ");
+
+    _ = local;
+    _ = ty_pl;
+    return f.fail("TODO: C backend: implement airSelect", .{});
+}
+
 fn airShuffle(f: *Function, inst: Air.Inst.Index) !CValue {
     if (f.liveness.isUnused(inst)) return CValue.none;
 

src/codegen/llvm.zig

@@ -3444,6 +3444,7 @@ pub const FuncGen = struct {
                 .tag_name       => try self.airTagName(inst),
                 .error_name     => try self.airErrorName(inst),
                 .splat          => try self.airSplat(inst),
+                .select         => try self.airSelect(inst),
                 .shuffle        => try self.airShuffle(inst),
                 .reduce         => try self.airReduce(inst),
                 .aggregate_init => try self.airAggregateInit(inst),
@@ -6355,6 +6356,18 @@ pub const FuncGen = struct {
         return self.builder.buildShuffleVector(op_vector, undef_vector, mask_llvm_ty.constNull(), "");
     }
 
+    fn airSelect(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
+        if (self.liveness.isUnused(inst)) return null;
+
+        const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+        const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
+        const pred = try self.resolveInst(extra.pred);
+        const a = try self.resolveInst(extra.a);
+        const b = try self.resolveInst(extra.b);
+
+        return self.builder.buildSelect(pred, a, b, "");
+    }
+
     fn airShuffle(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
         if (self.liveness.isUnused(inst)) return null;
 

src/print_air.zig

@@ -264,6 +264,7 @@ const Writer = struct {
             .wasm_memory_size => try w.writeWasmMemorySize(s, inst),
             .wasm_memory_grow => try w.writeWasmMemoryGrow(s, inst),
             .mul_add => try w.writeMulAdd(s, inst),
+            .select => try w.writeSelect(s, inst),
             .shuffle => try w.writeShuffle(s, inst),
             .reduce => try w.writeReduce(s, inst),
             .cmp_vector => try w.writeCmpVector(s, inst),
@@ -396,6 +397,18 @@ const Writer = struct {
         try s.print(", mask {d}, len {d}", .{ extra.mask, extra.mask_len });
     }
 
+    fn writeSelect(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
+        const ty_pl = w.air.instructions.items(.data)[inst].ty_pl;
+        const extra = w.air.extraData(Air.Select, ty_pl.payload).data;
+
+        try s.print("{}, ", .{w.air.getRefType(ty_pl.ty).fmtDebug()});
+        try w.writeOperand(s, inst, 0, extra.pred);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 1, extra.a);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 2, extra.b);
+    }
+
     fn writeReduce(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
         const reduce = w.air.instructions.items(.data)[inst].reduce;
 

test/behavior/select.zig

@@ -4,23 +4,29 @@ const mem = std.mem;
 const expect = std.testing.expect;
 
 test "@select" {
-    if (@import("builtin").zig_backend != .stage1) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_c) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_wasm) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
 
-    const S = struct {
+    try doTheTest();
-        fn doTheTest() !void {
+    comptime try doTheTest();
-            var a: @Vector(4, bool) = [4]bool{ true, false, true, false };
+}
-            var b: @Vector(4, i32) = [4]i32{ -1, 4, 999, -31 };
+
-            var c: @Vector(4, i32) = [4]i32{ -5, 1, 0, 1234 };
+fn doTheTest() !void {
-            var abc = @select(i32, a, b, c);
+    var a = @Vector(4, bool){ true, false, true, false };
-            try expect(mem.eql(i32, &@as([4]i32, abc), &[4]i32{ -1, 1, 999, 1234 }));
+    var b = @Vector(4, i32){ -1, 4, 999, -31 };
-
+    var c = @Vector(4, i32){ -5, 1, 0, 1234 };
-            var x: @Vector(4, bool) = [4]bool{ false, false, false, true };
+    var abc = @select(i32, a, b, c);
-            var y: @Vector(4, f32) = [4]f32{ 0.001, 33.4, 836, -3381.233 };
+    try expect(abc[0] == -1);
-            var z: @Vector(4, f32) = [4]f32{ 0.0, 312.1, -145.9, 9993.55 };
+    try expect(abc[1] == 1);
-            var xyz = @select(f32, x, y, z);
+    try expect(abc[2] == 999);
-            try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
+    try expect(abc[3] == 1234);
-        }
+
-    };
+    var x = @Vector(4, bool){ false, false, false, true };
-    try S.doTheTest();
+    var y = @Vector(4, f32){ 0.001, 33.4, 836, -3381.233 };
-    comptime try S.doTheTest();
+    var z = @Vector(4, f32){ 0.0, 312.1, -145.9, 9993.55 };
+    var xyz = @select(f32, x, y, z);
+    try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
 }