From cba68090a60c3de8eadbf8eb53e37620a1d66683 Mon Sep 17 00:00:00 2001 From: Veikka Tuominen Date: Fri, 11 Mar 2022 13:12:32 +0200 Subject: [PATCH] stage2: implement `@shuffle` at runtime --- src/Air.zig | 12 ++++ src/Liveness.zig | 4 ++ src/Sema.zig | 125 ++++++++++++++++++++++++++++++++++- src/arch/aarch64/CodeGen.zig | 7 ++ src/arch/arm/CodeGen.zig | 7 ++ src/arch/riscv64/CodeGen.zig | 7 ++ src/arch/wasm/CodeGen.zig | 12 ++++ src/arch/x86_64/CodeGen.zig | 7 ++ src/codegen/c.zig | 16 +++++ src/codegen/llvm.zig | 38 +++++++++++ src/print_air.zig | 11 +++ 11 files changed, 244 insertions(+), 2 deletions(-) diff --git a/src/Air.zig b/src/Air.zig index cc15113653..a1d5e1e8d9 100644 --- a/src/Air.zig +++ b/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, diff --git a/src/Liveness.zig b/src/Liveness.zig index 077d2f2888..c91288f354 100644 --- a/src/Liveness.zig +++ b/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); diff --git a/src/Sema.zig b/src/Sema.zig index 6f2dd9cac3..f631bf36b5 100644 --- a/src/Sema.zig +++ b/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 { diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig index 0bd549a45d..d50d0b2fdc 100644 --- a/src/arch/aarch64/CodeGen.zig +++ b/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(); diff --git a/src/arch/arm/CodeGen.zig b/src/arch/arm/CodeGen.zig index 3a247c09f2..ef8e612709 100644 --- a/src/arch/arm/CodeGen.zig +++ b/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(); diff --git a/src/arch/riscv64/CodeGen.zig b/src/arch/riscv64/CodeGen.zig index b38ec6ed93..326abdabdc 100644 --- a/src/arch/riscv64/CodeGen.zig +++ b/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(); diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig index 3f28b87b55..6c54699831 100644 --- a/src/arch/wasm/CodeGen.zig +++ b/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 = {} }; diff --git a/src/arch/x86_64/CodeGen.zig b/src/arch/x86_64/CodeGen.zig index be3d09bc8f..a665f9c1c8 100644 --- a/src/arch/x86_64/CodeGen.zig +++ b/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(); diff --git a/src/codegen/c.zig b/src/codegen/c.zig index 4a30bf023b..c179b65046 100644 --- a/src/codegen/c.zig +++ b/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; diff --git a/src/codegen/llvm.zig b/src/codegen/llvm.zig index d7df161b00..190fcd0ee6 100644 --- a/src/codegen/llvm.zig +++ b/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; diff --git a/src/print_air.zig b/src/print_air.zig index 56d0cb9623..749e4751fb 100644 --- a/src/print_air.zig +++ b/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;