authorgravatar for git@vexu.euVeikka Tuominen <git@vexu.eu> 2022-03-11 13:12:32+02:00
committergravatar for git@vexu.euVeikka Tuominen <git@vexu.eu> 2022-03-11 13:12:32+02:00
logcba68090a60c3de8eadbf8eb53e37620a1d66683
tree19d972a0c3cab03307218e9fb757fd3fdc44c8a3
parentdb4262417009c0060e09a8cf8af088b76320aad0

stage2: implement `@shuffle` at runtime


11 files changed, 244 insertions(+), 2 deletions(-)

src/Air.zig+12
......@@ -520,6 +520,9 @@ pub const Inst = struct {
520520 /// equal to the scalar value.
521521 /// Uses the `ty_op` field.
522522 splat,
523 /// Constructs a vector by selecting elements from `a` and `b` based on `mask`.
524 /// Uses the `ty_pl` field with payload `Shuffle`.
525 shuffle,
523526
524527 /// Given dest ptr, value, and len, set all elements at dest to value.
525528 /// Result type is always void.
......@@ -740,6 +743,14 @@ pub const FieldParentPtr = struct {
740743 field_index: u32,
741744};
742745
746pub const Shuffle = struct {
747 a: Inst.Ref,
748 b: Inst.Ref,
749 // index to air_values
750 mask: u32,
751 mask_len: u32,
752};
753
743754/// Trailing:
744755/// 0. `Inst.Ref` for every outputs_len
745756/// 1. `Inst.Ref` for every inputs_len
......@@ -897,6 +908,7 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
897908 .cmpxchg_weak,
898909 .cmpxchg_strong,
899910 .slice,
911 .shuffle,
900912 .aggregate_init,
901913 .union_init,
902914 .field_parent_ptr,
src/Liveness.zig+4
......@@ -422,6 +422,10 @@ fn analyzeInst(
422422 }
423423 return extra_tombs.finish();
424424 },
425 .shuffle => {
426 const extra = a.air.extraData(Air.Shuffle, inst_datas[inst].ty_pl.payload).data;
427 return trackOperands(a, new_set, inst, main_tomb, .{ extra.a, extra.b, .none });
428 },
425429 .aggregate_init => {
426430 const ty_pl = inst_datas[inst].ty_pl;
427431 const aggregate_ty = a.air.getRefType(ty_pl.ty);
src/Sema.zig+123-2
......@@ -13452,8 +13452,129 @@ fn zirReduce(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.
1345213452
1345313453fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
1345413454 const inst_data = sema.code.instructions.items(.data)[inst].pl_node;
13455 const src = inst_data.src();
13456 return sema.fail(block, src, "TODO: Sema.zirShuffle", .{});
13455 const extra = sema.code.extraData(Zir.Inst.Shuffle, inst_data.payload_index).data;
13456 const elem_ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
13457 const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
13458 const b_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = inst_data.src_node };
13459 const mask_src: LazySrcLoc = .{ .node_offset_builtin_call_arg3 = inst_data.src_node };
13460
13461 const elem_ty = try sema.resolveType(block, elem_ty_src, extra.elem_type);
13462 try sema.checkVectorElemType(block, elem_ty_src, elem_ty);
13463 var a = sema.resolveInst(extra.a);
13464 var b = sema.resolveInst(extra.b);
13465 var mask = sema.resolveInst(extra.mask);
13466 var mask_ty = sema.typeOf(mask);
13467
13468 const mask_len = switch (sema.typeOf(mask).zigTypeTag()) {
13469 .Array, .Vector => sema.typeOf(mask).arrayLen(),
13470 else => return sema.fail(block, mask_src, "expected vector or array, found {}", .{sema.typeOf(mask)}),
13471 };
13472 mask_ty = try Type.Tag.vector.create(sema.arena, .{
13473 .len = mask_len,
13474 .elem_type = Type.@"i32",
13475 });
13476 mask = try sema.coerce(block, mask_ty, mask, mask_src);
13477
13478 const res_ty = try Type.Tag.vector.create(sema.arena, .{
13479 .len = mask_len,
13480 .elem_type = elem_ty,
13481 });
13482
13483 var maybe_a_len = switch (sema.typeOf(a).zigTypeTag()) {
13484 .Array, .Vector => sema.typeOf(a).arrayLen(),
13485 .Undefined => null,
13486 else => return sema.fail(block, a_src, "expected vector or array with element type {}, found {}", .{
13487 elem_ty,
13488 sema.typeOf(mask),
13489 }),
13490 };
13491 var maybe_b_len = switch (sema.typeOf(b).zigTypeTag()) {
13492 .Array, .Vector => sema.typeOf(b).arrayLen(),
13493 .Undefined => null,
13494 else => return sema.fail(block, b_src, "expected vector or array with element type {}, found {}", .{
13495 elem_ty,
13496 sema.typeOf(mask),
13497 }),
13498 };
13499 if (maybe_a_len == null and maybe_b_len == null) {
13500 return sema.addConstUndef(res_ty);
13501 }
13502 const a_len = maybe_a_len orelse maybe_b_len.?;
13503 const b_len = maybe_b_len orelse a_len;
13504
13505 const a_ty = try Type.Tag.vector.create(sema.arena, .{
13506 .len = a_len,
13507 .elem_type = elem_ty,
13508 });
13509 const b_ty = try Type.Tag.vector.create(sema.arena, .{
13510 .len = b_len,
13511 .elem_type = elem_ty,
13512 });
13513
13514 if (maybe_a_len == null) a = try sema.addConstUndef(a_ty);
13515 if (maybe_b_len == null) b = try sema.addConstUndef(b_ty);
13516
13517 const operand_info = [2]std.meta.Tuple(&.{ u64, LazySrcLoc, Type }){
13518 .{ a_len, a_src, a_ty },
13519 .{ b_len, b_src, b_ty },
13520 };
13521
13522 const mask_val = try sema.resolveConstMaybeUndefVal(block, mask_src, mask);
13523 var i: usize = 0;
13524 while (i < mask_len) : (i += 1) {
13525 var buf: Value.ElemValueBuffer = undefined;
13526 const elem = mask_val.elemValueBuffer(i, &buf);
13527 if (elem.isUndef()) continue;
13528 const int = elem.toSignedInt();
13529 var unsigned: u32 = undefined;
13530 var chosen: u32 = undefined;
13531 if (int >= 0) {
13532 unsigned = @intCast(u32, int);
13533 chosen = 0;
13534 } else {
13535 unsigned = @intCast(u32, ~int);
13536 chosen = 1;
13537 }
13538 if (unsigned >= operand_info[chosen][0]) {
13539 const msg = msg: {
13540 const msg = try sema.errMsg(block, mask_src, "mask index {d} has out-of-bounds selection", .{i});
13541 errdefer msg.destroy(sema.gpa);
13542
13543 try sema.errNote(block, operand_info[chosen][1], msg, "selected index {d} out of bounds of {}", .{
13544 unsigned,
13545 operand_info[chosen][2],
13546 });
13547
13548 if (chosen == 1) {
13549 try sema.errNote(block, b_src, msg, "selections from the second vector are specified with negative numbers", .{});
13550 }
13551
13552 break :msg msg;
13553 };
13554 return sema.failWithOwnedErrorMsg(block, msg);
13555 }
13556 }
13557
13558 // TODO at comptime
13559
13560 if (a_len != b_len) {
13561 return sema.fail(block, mask_src, "TODO handle shuffle a_len != b_len", .{});
13562 }
13563
13564 const mask_index = @intCast(u32, sema.air_values.items.len);
13565 try sema.air_values.append(sema.gpa, mask_val);
13566 return block.addInst(.{
13567 .tag = .shuffle,
13568 .data = .{ .ty_pl = .{
13569 .ty = try sema.addType(res_ty),
13570 .payload = try block.sema.addExtra(Air.Shuffle{
13571 .a = a,
13572 .b = b,
13573 .mask = mask_index,
13574 .mask_len = @intCast(u32, mask_len),
13575 }),
13576 } },
13577 });
1345713578}
1345813579
1345913580fn zirSelect(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
src/arch/aarch64/CodeGen.zig+7
......@@ -628,6 +628,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
628628 .tag_name => try self.airTagName(inst),
629629 .error_name => try self.airErrorName(inst),
630630 .splat => try self.airSplat(inst),
631 .shuffle => try self.airShuffle(inst),
631632 .aggregate_init => try self.airAggregateInit(inst),
632633 .union_init => try self.airUnionInit(inst),
633634 .prefetch => try self.airPrefetch(inst),
......@@ -3624,6 +3625,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
36243625 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
36253626}
36263627
3628fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
3629 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
3630 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for {}", .{self.target.cpu.arch});
3631 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
3632}
3633
36273634fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
36283635 const vector_ty = self.air.typeOfIndex(inst);
36293636 const len = vector_ty.vectorLen();
src/arch/arm/CodeGen.zig+7
......@@ -624,6 +624,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
624624 .tag_name => try self.airTagName(inst),
625625 .error_name => try self.airErrorName(inst),
626626 .splat => try self.airSplat(inst),
627 .shuffle => try self.airShuffle(inst),
627628 .aggregate_init => try self.airAggregateInit(inst),
628629 .union_init => try self.airUnionInit(inst),
629630 .prefetch => try self.airPrefetch(inst),
......@@ -4081,6 +4082,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
40814082 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
40824083}
40834084
4085fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
4086 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
4087 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for arm", .{});
4088 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
4089}
4090
40844091fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
40854092 const vector_ty = self.air.typeOfIndex(inst);
40864093 const len = vector_ty.vectorLen();
src/arch/riscv64/CodeGen.zig+7
......@@ -596,6 +596,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
596596 .tag_name => try self.airTagName(inst),
597597 .error_name => try self.airErrorName(inst),
598598 .splat => try self.airSplat(inst),
599 .shuffle => try self.airShuffle(inst),
599600 .aggregate_init => try self.airAggregateInit(inst),
600601 .union_init => try self.airUnionInit(inst),
601602 .prefetch => try self.airPrefetch(inst),
......@@ -2174,6 +2175,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
21742175 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
21752176}
21762177
2178fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
2179 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
2180 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for riscv64", .{});
2181 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
2182}
2183
21772184fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
21782185 const vector_ty = self.air.typeOfIndex(inst);
21792186 const len = vector_ty.vectorLen();
src/arch/wasm/CodeGen.zig+12
......@@ -1255,6 +1255,7 @@ fn genInst(self: *Self, inst: Air.Inst.Index) !WValue {
12551255 .ret_ptr => self.airRetPtr(inst),
12561256 .ret_load => self.airRetLoad(inst),
12571257 .splat => self.airSplat(inst),
1258 .shuffle => self.airShuffle(inst),
12581259 .aggregate_init => self.airAggregateInit(inst),
12591260 .union_init => self.airUnionInit(inst),
12601261 .prefetch => self.airPrefetch(inst),
......@@ -2985,6 +2986,17 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
29852986 return self.fail("TODO: Implement wasm airSplat", .{});
29862987}
29872988
2989fn airShuffle(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
2990 if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
2991
2992 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
2993 const operand = try self.resolveInst(ty_op.operand);
2994
2995 _ = ty_op;
2996 _ = operand;
2997 return self.fail("TODO: Implement wasm airShuffle", .{});
2998}
2999
29883000fn airAggregateInit(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
29893001 if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
29903002
src/arch/x86_64/CodeGen.zig+7
......@@ -713,6 +713,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
713713 .tag_name => try self.airTagName(inst),
714714 .error_name => try self.airErrorName(inst),
715715 .splat => try self.airSplat(inst),
716 .shuffle => try self.airShuffle(inst),
716717 .aggregate_init => try self.airAggregateInit(inst),
717718 .union_init => try self.airUnionInit(inst),
718719 .prefetch => try self.airPrefetch(inst),
......@@ -5528,6 +5529,12 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
55285529 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
55295530}
55305531
5532fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
5533 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
5534 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for x86_64", .{});
5535 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
5536}
5537
55315538fn airAggregateInit(self: *Self, inst: Air.Inst.Index) !void {
55325539 const vector_ty = self.air.typeOfIndex(inst);
55335540 const len = vector_ty.vectorLen();
src/codegen/c.zig+16
......@@ -1716,6 +1716,7 @@ fn genBody(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, OutO
17161716 .tag_name => try airTagName(f, inst),
17171717 .error_name => try airErrorName(f, inst),
17181718 .splat => try airSplat(f, inst),
1719 .shuffle => try airShuffle(f, inst),
17191720 .aggregate_init => try airAggregateInit(f, inst),
17201721 .union_init => try airUnionInit(f, inst),
17211722 .prefetch => try airPrefetch(f, inst),
......@@ -3557,6 +3558,21 @@ fn airSplat(f: *Function, inst: Air.Inst.Index) !CValue {
35573558 return f.fail("TODO: C backend: implement airSplat", .{});
35583559}
35593560
3561fn airShuffle(f: *Function, inst: Air.Inst.Index) !CValue {
3562 if (f.liveness.isUnused(inst)) return CValue.none;
3563
3564 const inst_ty = f.air.typeOfIndex(inst);
3565 const ty_op = f.air.instructions.items(.data)[inst].ty_op;
3566 const operand = try f.resolveInst(ty_op.operand);
3567 const writer = f.object.writer();
3568 const local = try f.allocLocal(inst_ty, .Const);
3569 try writer.writeAll(" = ");
3570
3571 _ = operand;
3572 _ = local;
3573 return f.fail("TODO: C backend: implement airShuffle", .{});
3574}
3575
35603576fn airAggregateInit(f: *Function, inst: Air.Inst.Index) !CValue {
35613577 if (f.liveness.isUnused(inst)) return CValue.none;
35623578
src/codegen/llvm.zig+38
......@@ -3204,6 +3204,7 @@ pub const FuncGen = struct {
32043204 .tag_name => try self.airTagName(inst),
32053205 .error_name => try self.airErrorName(inst),
32063206 .splat => try self.airSplat(inst),
3207 .shuffle => try self.airShuffle(inst),
32073208 .aggregate_init => try self.airAggregateInit(inst),
32083209 .union_init => try self.airUnionInit(inst),
32093210 .prefetch => try self.airPrefetch(inst),
......@@ -5850,6 +5851,43 @@ pub const FuncGen = struct {
58505851 return self.builder.buildShuffleVector(op_vector, undef_vector, mask_llvm_ty.constNull(), "");
58515852 }
58525853
5854 fn airShuffle(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
5855 if (self.liveness.isUnused(inst)) return null;
5856
5857 const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
5858 const extra = self.air.extraData(Air.Shuffle, ty_pl.payload).data;
5859 const a = try self.resolveInst(extra.a);
5860 const b = try self.resolveInst(extra.b);
5861 const mask = self.air.values[extra.mask];
5862 const mask_len = extra.mask_len;
5863 const a_len = self.air.typeOf(extra.a).vectorLen();
5864
5865 // LLVM uses integers larger than the length of the first array to
5866 // index into the second array. This was deemed unnecessarily fragile
5867 // when changing code, so Zig uses negative numbers to index the
5868 // second vector. These start at -1 and go down, and are easiest to use
5869 // with the ~ operator. Here we convert between the two formats.
5870 const values = try self.gpa.alloc(*const llvm.Value, mask_len);
5871 defer self.gpa.free(values);
5872
5873 const llvm_i32 = self.context.intType(32);
5874
5875 for (values) |*val, i| {
5876 var buf: Value.ElemValueBuffer = undefined;
5877 const elem = mask.elemValueBuffer(i, &buf);
5878 if (elem.isUndef()) {
5879 val.* = llvm_i32.getUndef();
5880 } else {
5881 const int = elem.toSignedInt();
5882 const unsigned = if (int >= 0) @intCast(u32, int) else @intCast(u32, ~int + a_len);
5883 val.* = llvm_i32.constInt(unsigned, .False);
5884 }
5885 }
5886
5887 const llvm_mask_value = llvm.constVector(values.ptr, mask_len);
5888 return self.builder.buildShuffleVector(a, b, llvm_mask_value, "");
5889 }
5890
58535891 fn airAggregateInit(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
58545892 if (self.liveness.isUnused(inst)) return null;
58555893
src/print_air.zig+11
......@@ -258,6 +258,7 @@ const Writer = struct {
258258 .wasm_memory_size => try w.writeWasmMemorySize(s, inst),
259259 .wasm_memory_grow => try w.writeWasmMemoryGrow(s, inst),
260260 .mul_add => try w.writeMulAdd(s, inst),
261 .shuffle => try w.writeShuffle(s, inst),
261262
262263 .add_with_overflow,
263264 .sub_with_overflow,
......@@ -375,6 +376,16 @@ const Writer = struct {
375376 try w.writeOperand(s, inst, 2, pl_op.operand);
376377 }
377378
379 fn writeShuffle(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
380 const pl_op = w.air.instructions.items(.data)[inst].pl_op;
381 const extra = w.air.extraData(Air.Shuffle, pl_op.payload).data;
382
383 try w.writeOperand(s, inst, 0, extra.a);
384 try s.writeAll(", ");
385 try w.writeOperand(s, inst, 1, extra.b);
386 try s.print(", mask {d}, len {d}", .{ extra.mask, extra.mask_len });
387 }
388
378389 fn writeFence(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
379390 const atomic_order = w.air.instructions.items(.data)[inst].fence;
380391