stage2: implement @shuffle at runtime

src/Air.zig

@@ -520,6 +520,9 @@ pub const Inst = struct {
         /// equal to the scalar value.
         /// Uses the `ty_op` field.
         splat,
+        /// Constructs a vector by selecting elements from `a` and `b` based on `mask`.
+        /// Uses the `ty_pl` field with payload `Shuffle`.
+        shuffle,
 
         /// Given dest ptr, value, and len, set all elements at dest to value.
         /// Result type is always void.
@@ -740,6 +743,14 @@ pub const FieldParentPtr = struct {
     field_index: u32,
 };
 
+pub const Shuffle = struct {
+    a: Inst.Ref,
+    b: Inst.Ref,
+    // index to air_values
+    mask: u32,
+    mask_len: u32,
+};
+
 /// Trailing:
 /// 0. `Inst.Ref` for every outputs_len
 /// 1. `Inst.Ref` for every inputs_len
@@ -897,6 +908,7 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
         .cmpxchg_weak,
         .cmpxchg_strong,
         .slice,
+        .shuffle,
         .aggregate_init,
         .union_init,
         .field_parent_ptr,

src/Liveness.zig

@@ -422,6 +422,10 @@ fn analyzeInst(
             }
             return extra_tombs.finish();
         },
+        .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 });
+        },
         .aggregate_init => {
             const ty_pl = inst_datas[inst].ty_pl;
             const aggregate_ty = a.air.getRefType(ty_pl.ty);

src/Sema.zig

@@ -13452,8 +13452,129 @@ fn zirReduce(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.
 
 fn zirShuffle(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();
-    return sema.fail(block, src, "TODO: Sema.zirShuffle", .{});
+    const extra = sema.code.extraData(Zir.Inst.Shuffle, inst_data.payload_index).data;
+    const elem_ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
+    const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
+    const b_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = inst_data.src_node };
+    const mask_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);
+    var a = sema.resolveInst(extra.a);
+    var b = sema.resolveInst(extra.b);
+    var mask = sema.resolveInst(extra.mask);
+    var mask_ty = sema.typeOf(mask);
+
+    const mask_len = switch (sema.typeOf(mask).zigTypeTag()) {
+        .Array, .Vector => sema.typeOf(mask).arrayLen(),
+        else => return sema.fail(block, mask_src, "expected vector or array, found {}", .{sema.typeOf(mask)}),
+    };
+    mask_ty = try Type.Tag.vector.create(sema.arena, .{
+        .len = mask_len,
+        .elem_type = Type.@"i32",
+    });
+    mask = try sema.coerce(block, mask_ty, mask, mask_src);
+
+    const res_ty = try Type.Tag.vector.create(sema.arena, .{
+        .len = mask_len,
+        .elem_type = elem_ty,
+    });
+
+    var maybe_a_len = switch (sema.typeOf(a).zigTypeTag()) {
+        .Array, .Vector => sema.typeOf(a).arrayLen(),
+        .Undefined => null,
+        else => return sema.fail(block, a_src, "expected vector or array with element type {}, found {}", .{
+            elem_ty,
+            sema.typeOf(mask),
+        }),
+    };
+    var maybe_b_len = switch (sema.typeOf(b).zigTypeTag()) {
+        .Array, .Vector => sema.typeOf(b).arrayLen(),
+        .Undefined => null,
+        else => return sema.fail(block, b_src, "expected vector or array with element type {}, found {}", .{
+            elem_ty,
+            sema.typeOf(mask),
+        }),
+    };
+    if (maybe_a_len == null and maybe_b_len == null) {
+        return sema.addConstUndef(res_ty);
+    }
+    const a_len = maybe_a_len orelse maybe_b_len.?;
+    const b_len = maybe_b_len orelse a_len;
+
+    const a_ty = try Type.Tag.vector.create(sema.arena, .{
+        .len = a_len,
+        .elem_type = elem_ty,
+    });
+    const b_ty = try Type.Tag.vector.create(sema.arena, .{
+        .len = b_len,
+        .elem_type = elem_ty,
+    });
+
+    if (maybe_a_len == null) a = try sema.addConstUndef(a_ty);
+    if (maybe_b_len == null) b = try sema.addConstUndef(b_ty);
+
+    const operand_info = [2]std.meta.Tuple(&.{ u64, LazySrcLoc, Type }){
+        .{ a_len, a_src, a_ty },
+        .{ b_len, b_src, b_ty },
+    };
+
+    const mask_val = try sema.resolveConstMaybeUndefVal(block, mask_src, mask);
+    var i: usize = 0;
+    while (i < mask_len) : (i += 1) {
+        var buf: Value.ElemValueBuffer = undefined;
+        const elem = mask_val.elemValueBuffer(i, &buf);
+        if (elem.isUndef()) continue;
+        const int = elem.toSignedInt();
+        var unsigned: u32 = undefined;
+        var chosen: u32 = undefined;
+        if (int >= 0) {
+            unsigned = @intCast(u32, int);
+            chosen = 0;
+        } else {
+            unsigned = @intCast(u32, ~int);
+            chosen = 1;
+        }
+        if (unsigned >= operand_info[chosen][0]) {
+            const msg = msg: {
+                const msg = try sema.errMsg(block, mask_src, "mask index {d} has out-of-bounds selection", .{i});
+                errdefer msg.destroy(sema.gpa);
+
+                try sema.errNote(block, operand_info[chosen][1], msg, "selected index {d} out of bounds of {}", .{
+                    unsigned,
+                    operand_info[chosen][2],
+                });
+
+                if (chosen == 1) {
+                    try sema.errNote(block, b_src, msg, "selections from the second vector are specified with negative numbers", .{});
+                }
+
+                break :msg msg;
+            };
+            return sema.failWithOwnedErrorMsg(block, msg);
+        }
+    }
+
+    // TODO at comptime
+
+    if (a_len != b_len) {
+        return sema.fail(block, mask_src, "TODO handle shuffle a_len != b_len", .{});
+    }
+
+    const mask_index = @intCast(u32, sema.air_values.items.len);
+    try sema.air_values.append(sema.gpa, mask_val);
+    return block.addInst(.{
+        .tag = .shuffle,
+        .data = .{ .ty_pl = .{
+            .ty = try sema.addType(res_ty),
+            .payload = try block.sema.addExtra(Air.Shuffle{
+                .a = a,
+                .b = b,
+                .mask = mask_index,
+                .mask_len = @intCast(u32, mask_len),
+            }),
+        } },
+    });
 }
 
 fn zirSelect(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {

src/arch/aarch64/CodeGen.zig

@@ -628,6 +628,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),
+            .shuffle         => try self.airShuffle(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
             .union_init      => try self.airUnionInit(inst),
             .prefetch        => try self.airPrefetch(inst),
@@ -3624,6 +3625,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+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});
+    return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
+}
+
 fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
     const vector_ty = self.air.typeOfIndex(inst);
     const len = vector_ty.vectorLen();

src/arch/arm/CodeGen.zig

@@ -624,6 +624,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),
+            .shuffle         => try self.airShuffle(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
             .union_init      => try self.airUnionInit(inst),
             .prefetch        => try self.airPrefetch(inst),
@@ -4081,6 +4082,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+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", .{});
+    return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
+}
+
 fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
     const vector_ty = self.air.typeOfIndex(inst);
     const len = vector_ty.vectorLen();

src/arch/riscv64/CodeGen.zig

@@ -596,6 +596,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),
+            .shuffle         => try self.airShuffle(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
             .union_init      => try self.airUnionInit(inst),
             .prefetch        => try self.airPrefetch(inst),
@@ -2174,6 +2175,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+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", .{});
+    return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
+}
+
 fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
     const vector_ty = self.air.typeOfIndex(inst);
     const len = vector_ty.vectorLen();

src/arch/wasm/CodeGen.zig

@@ -1255,6 +1255,7 @@ fn genInst(self: *Self, inst: Air.Inst.Index) !WValue {
         .ret_ptr => self.airRetPtr(inst),
         .ret_load => self.airRetLoad(inst),
         .splat => self.airSplat(inst),
+        .shuffle => self.airShuffle(inst),
         .aggregate_init => self.airAggregateInit(inst),
         .union_init => self.airUnionInit(inst),
         .prefetch => self.airPrefetch(inst),
@@ -2985,6 +2986,17 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
     return self.fail("TODO: Implement wasm airSplat", .{});
 }
 
+fn airShuffle(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
+    if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
+
+    const ty_op = self.air.instructions.items(.data)[inst].ty_op;
+    const operand = try self.resolveInst(ty_op.operand);
+
+    _ = ty_op;
+    _ = operand;
+    return self.fail("TODO: Implement wasm airShuffle", .{});
+}
+
 fn airAggregateInit(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
     if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
 

src/arch/x86_64/CodeGen.zig

@@ -713,6 +713,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),
+            .shuffle         => try self.airShuffle(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
             .union_init      => try self.airUnionInit(inst),
             .prefetch        => try self.airPrefetch(inst),
@@ -5528,6 +5529,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+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", .{});
+    return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
+}
+
 fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
     const vector_ty = self.air.typeOfIndex(inst);
     const len = vector_ty.vectorLen();

src/codegen/c.zig

@@ -1716,6 +1716,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),
+            .shuffle          => try airShuffle(f, inst),
             .aggregate_init   => try airAggregateInit(f, inst),
             .union_init       => try airUnionInit(f, inst),
             .prefetch         => try airPrefetch(f, inst),
@@ -3557,6 +3558,21 @@ fn airSplat(f: *Function, inst: Air.Inst.Index) !CValue {
     return f.fail("TODO: C backend: implement airSplat", .{});
 }
 
+fn airShuffle(f: *Function, inst: Air.Inst.Index) !CValue {
+    if (f.liveness.isUnused(inst)) return CValue.none;
+
+    const inst_ty = f.air.typeOfIndex(inst);
+    const ty_op = f.air.instructions.items(.data)[inst].ty_op;
+    const operand = try f.resolveInst(ty_op.operand);
+    const writer = f.object.writer();
+    const local = try f.allocLocal(inst_ty, .Const);
+    try writer.writeAll(" = ");
+
+    _ = operand;
+    _ = local;
+    return f.fail("TODO: C backend: implement airShuffle", .{});
+}
+
 fn airAggregateInit(f: *Function, inst: Air.Inst.Index) !CValue {
     if (f.liveness.isUnused(inst)) return CValue.none;
 

src/codegen/llvm.zig

@@ -3204,6 +3204,7 @@ pub const FuncGen = struct {
                 .tag_name       => try self.airTagName(inst),
                 .error_name     => try self.airErrorName(inst),
                 .splat          => try self.airSplat(inst),
+                .shuffle        => try self.airShuffle(inst),
                 .aggregate_init => try self.airAggregateInit(inst),
                 .union_init     => try self.airUnionInit(inst),
                 .prefetch       => try self.airPrefetch(inst),
@@ -5850,6 +5851,43 @@ pub const FuncGen = struct {
         return self.builder.buildShuffleVector(op_vector, undef_vector, mask_llvm_ty.constNull(), "");
     }
 
+    fn airShuffle(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.Shuffle, ty_pl.payload).data;
+        const a = try self.resolveInst(extra.a);
+        const b = try self.resolveInst(extra.b);
+        const mask = self.air.values[extra.mask];
+        const mask_len = extra.mask_len;
+        const a_len = self.air.typeOf(extra.a).vectorLen();
+
+        // LLVM uses integers larger than the length of the first array to
+        // index into the second array. This was deemed unnecessarily fragile
+        // when changing code, so Zig uses negative numbers to index the
+        // second vector. These start at -1 and go down, and are easiest to use
+        // with the ~ operator. Here we convert between the two formats.
+        const values = try self.gpa.alloc(*const llvm.Value, mask_len);
+        defer self.gpa.free(values);
+
+        const llvm_i32 = self.context.intType(32);
+
+        for (values) |*val, i| {
+            var buf: Value.ElemValueBuffer = undefined;
+            const elem = mask.elemValueBuffer(i, &buf);
+            if (elem.isUndef()) {
+                val.* = llvm_i32.getUndef();
+            } else {
+                const int = elem.toSignedInt();
+                const unsigned = if (int >= 0) @intCast(u32, int) else @intCast(u32, ~int + a_len);
+                val.* = llvm_i32.constInt(unsigned, .False);
+            }
+        }
+
+        const llvm_mask_value = llvm.constVector(values.ptr, mask_len);
+        return self.builder.buildShuffleVector(a, b, llvm_mask_value, "");
+    }
+
     fn airAggregateInit(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
         if (self.liveness.isUnused(inst)) return null;
 

src/print_air.zig

@@ -258,6 +258,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),
+            .shuffle => try w.writeShuffle(s, inst),
 
             .add_with_overflow,
             .sub_with_overflow,
@@ -375,6 +376,16 @@ const Writer = struct {
         try w.writeOperand(s, inst, 2, pl_op.operand);
     }
 
+    fn writeShuffle(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
+        const pl_op = w.air.instructions.items(.data)[inst].pl_op;
+        const extra = w.air.extraData(Air.Shuffle, pl_op.payload).data;
+
+        try w.writeOperand(s, inst, 0, extra.a);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 1, extra.b);
+        try s.print(", mask {d}, len {d}", .{ extra.mask, extra.mask_len });
+    }
+
     fn writeFence(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
         const atomic_order = w.air.instructions.items(.data)[inst].fence;