authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2022-03-26 00:33:22-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2022-03-26 00:33:22-04:00
log88e98a0611b9fb41c1da026febac2467548bb129
tree038ccb88c632580be1898a536aa2433077c2e6d0
parentbae35bdf2d8919b60dee9a0af3afbdd93dd72b59
parentcd46daf7d047eeceb7690e2739af5952d60c3884
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #11289 from schmee/stage2-select

stage2: implement `@select`

12 files changed, 239 insertions(+), 23 deletions(-)

src/Air.zig+8-1
......@@ -344,7 +344,7 @@ pub const Inst = struct {
344344 /// to the storage for the variable. The local may be a const or a var.
345345 /// Result type is always void.
346346 /// Uses `pl_op`. The payload index is the variable name. It points to the extra
347 /// array, reinterpreting the bytes there as a null-terminated string.
347 /// array, reinterpreting the bytes there as a null-terminated string.
348348 dbg_var_ptr,
349349 /// Same as `dbg_var_ptr` except the local is a const, not a var, and the
350350 /// operand is the local's value.
......@@ -553,6 +553,9 @@ pub const Inst = struct {
553553 /// Constructs a vector by selecting elements from `a` and `b` based on `mask`.
554554 /// Uses the `ty_pl` field with payload `Shuffle`.
555555 shuffle,
556 /// Constructs a vector element-wise from `a` or `b` based on `pred`.
557 /// Uses the `pl_op` field with `pred` as operand, and payload `Bin`.
558 select,
556559
557560 /// Given dest ptr, value, and len, set all elements at dest to value.
558561 /// Result type is always void.
......@@ -1067,6 +1070,10 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
10671070 .reduce => return air.typeOf(datas[inst].reduce.operand).childType(),
10681071
10691072 .mul_add => return air.typeOf(datas[inst].pl_op.operand),
1073 .select => {
1074 const extra = air.extraData(Air.Bin, datas[inst].pl_op.payload).data;
1075 return air.typeOf(extra.lhs);
1076 },
10701077
10711078 .add_with_overflow,
10721079 .sub_with_overflow,
src/Liveness.zig+5
......@@ -433,6 +433,11 @@ fn analyzeInst(
433433 }
434434 return extra_tombs.finish();
435435 },
436 .select => {
437 const pl_op = inst_datas[inst].pl_op;
438 const extra = a.air.extraData(Air.Bin, pl_op.payload).data;
439 return trackOperands(a, new_set, inst, main_tomb, .{ pl_op.operand, extra.lhs, extra.rhs });
440 },
436441 .shuffle => {
437442 const extra = a.air.extraData(Air.Shuffle, inst_datas[inst].ty_pl.payload).data;
438443 return trackOperands(a, new_set, inst, main_tomb, .{ extra.a, extra.b, .none });
src/Sema.zig+85-2
......@@ -14890,8 +14890,91 @@ fn analyzeShuffle(
1489014890
1489114891fn zirSelect(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
1489214892 const inst_data = sema.code.instructions.items(.data)[inst].pl_node;
14893 const src = inst_data.src();
14894 return sema.fail(block, src, "TODO: Sema.zirSelect", .{});
14893 const extra = sema.code.extraData(Zir.Inst.Select, inst_data.payload_index).data;
14894 const target = sema.mod.getTarget();
14895
14896 const elem_ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
14897 const pred_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
14898 const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = inst_data.src_node };
14899 const b_src: LazySrcLoc = .{ .node_offset_builtin_call_arg3 = inst_data.src_node };
14900
14901 const elem_ty = try sema.resolveType(block, elem_ty_src, extra.elem_type);
14902 try sema.checkVectorElemType(block, elem_ty_src, elem_ty);
14903 const pred_uncoerced = sema.resolveInst(extra.pred);
14904 const pred_ty = sema.typeOf(pred_uncoerced);
14905
14906 const vec_len_u64 = switch (try pred_ty.zigTypeTagOrPoison()) {
14907 .Vector, .Array => pred_ty.arrayLen(),
14908 else => return sema.fail(block, pred_src, "expected vector or array, found '{}'", .{pred_ty.fmt(target)}),
14909 };
14910 const vec_len = try sema.usizeCast(block, pred_src, vec_len_u64);
14911
14912 const bool_vec_ty = try Type.vector(sema.arena, vec_len, Type.bool);
14913 const pred = try sema.coerce(block, bool_vec_ty, pred_uncoerced, pred_src);
14914
14915 const vec_ty = try Type.vector(sema.arena, vec_len, elem_ty);
14916 const a = try sema.coerce(block, vec_ty, sema.resolveInst(extra.a), a_src);
14917 const b = try sema.coerce(block, vec_ty, sema.resolveInst(extra.b), b_src);
14918
14919 const maybe_pred = try sema.resolveMaybeUndefVal(block, pred_src, pred);
14920 const maybe_a = try sema.resolveMaybeUndefVal(block, a_src, a);
14921 const maybe_b = try sema.resolveMaybeUndefVal(block, b_src, b);
14922
14923 const runtime_src = if (maybe_pred) |pred_val| rs: {
14924 if (pred_val.isUndef()) return sema.addConstUndef(vec_ty);
14925
14926 if (maybe_a) |a_val| {
14927 if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
14928
14929 if (maybe_b) |b_val| {
14930 if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
14931
14932 var buf: Value.ElemValueBuffer = undefined;
14933 const elems = try sema.gpa.alloc(Value, vec_len);
14934 for (elems) |*elem, i| {
14935 const pred_elem_val = pred_val.elemValueBuffer(i, &buf);
14936 const should_choose_a = pred_elem_val.toBool();
14937 if (should_choose_a) {
14938 elem.* = a_val.elemValueBuffer(i, &buf);
14939 } else {
14940 elem.* = b_val.elemValueBuffer(i, &buf);
14941 }
14942 }
14943
14944 return sema.addConstant(
14945 vec_ty,
14946 try Value.Tag.aggregate.create(sema.arena, elems),
14947 );
14948 } else {
14949 break :rs b_src;
14950 }
14951 } else {
14952 if (maybe_b) |b_val| {
14953 if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
14954 }
14955 break :rs a_src;
14956 }
14957 } else rs: {
14958 if (maybe_a) |a_val| {
14959 if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
14960 }
14961 if (maybe_b) |b_val| {
14962 if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
14963 }
14964 break :rs pred_src;
14965 };
14966
14967 try sema.requireRuntimeBlock(block, runtime_src);
14968 return block.addInst(.{
14969 .tag = .select,
14970 .data = .{ .pl_op = .{
14971 .operand = pred,
14972 .payload = try block.sema.addExtra(Air.Bin{
14973 .lhs = a,
14974 .rhs = b,
14975 }),
14976 } },
14977 });
1489514978}
1489614979
1489714980fn zirAtomicLoad(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
src/arch/aarch64/CodeGen.zig+8
......@@ -633,6 +633,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
633633 .tag_name => try self.airTagName(inst),
634634 .error_name => try self.airErrorName(inst),
635635 .splat => try self.airSplat(inst),
636 .select => try self.airSelect(inst),
636637 .shuffle => try self.airShuffle(inst),
637638 .reduce => try self.airReduce(inst),
638639 .aggregate_init => try self.airAggregateInit(inst),
......@@ -3666,6 +3667,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
36663667 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
36673668}
36683669
3670fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
3671 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
3672 const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
3673 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for {}", .{self.target.cpu.arch});
3674 return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
3675}
3676
36693677fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
36703678 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
36713679 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+8
......@@ -630,6 +630,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
630630 .tag_name => try self.airTagName(inst),
631631 .error_name => try self.airErrorName(inst),
632632 .splat => try self.airSplat(inst),
633 .select => try self.airSelect(inst),
633634 .shuffle => try self.airShuffle(inst),
634635 .reduce => try self.airReduce(inst),
635636 .aggregate_init => try self.airAggregateInit(inst),
......@@ -4323,6 +4324,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
43234324 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
43244325}
43254326
4327fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
4328 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
4329 const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
4330 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for arm", .{});
4331 return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
4332}
4333
43264334fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
43274335 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
43284336 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for arm", .{});
src/arch/riscv64/CodeGen.zig+8
......@@ -600,6 +600,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
600600 .tag_name => try self.airTagName(inst),
601601 .error_name => try self.airErrorName(inst),
602602 .splat => try self.airSplat(inst),
603 .select => try self.airSelect(inst),
603604 .shuffle => try self.airShuffle(inst),
604605 .reduce => try self.airReduce(inst),
605606 .aggregate_init => try self.airAggregateInit(inst),
......@@ -2396,6 +2397,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
23962397 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
23972398}
23982399
2400fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
2401 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
2402 const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
2403 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for riscv64", .{});
2404 return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
2405}
2406
23992407fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
24002408 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
24012409 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for riscv64", .{});
src/arch/wasm/CodeGen.zig+11
......@@ -1371,6 +1371,7 @@ fn genInst(self: *Self, inst: Air.Inst.Index) !WValue {
13711371 .ret_ptr => self.airRetPtr(inst),
13721372 .ret_load => self.airRetLoad(inst),
13731373 .splat => self.airSplat(inst),
1374 .select => self.airSelect(inst),
13741375 .shuffle => self.airShuffle(inst),
13751376 .reduce => self.airReduce(inst),
13761377 .aggregate_init => self.airAggregateInit(inst),
......@@ -3265,6 +3266,16 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
32653266 return self.fail("TODO: Implement wasm airSplat", .{});
32663267}
32673268
3269fn airSelect(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
3270 if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
3271
3272 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
3273 const operand = try self.resolveInst(pl_op.operand);
3274
3275 _ = operand;
3276 return self.fail("TODO: Implement wasm airSelect", .{});
3277}
3278
32683279fn airShuffle(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
32693280 if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
32703281
src/arch/x86_64/CodeGen.zig+8
......@@ -714,6 +714,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
714714 .tag_name => try self.airTagName(inst),
715715 .error_name => try self.airErrorName(inst),
716716 .splat => try self.airSplat(inst),
717 .select => try self.airSelect(inst),
717718 .shuffle => try self.airShuffle(inst),
718719 .reduce => try self.airReduce(inst),
719720 .aggregate_init => try self.airAggregateInit(inst),
......@@ -5624,6 +5625,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
56245625 return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
56255626}
56265627
5628fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
5629 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
5630 const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
5631 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for x86_64", .{});
5632 return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
5633}
5634
56275635fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
56285636 const ty_op = self.air.instructions.items(.data)[inst].ty_op;
56295637 const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for x86_64", .{});
src/codegen/c.zig+16
......@@ -1825,6 +1825,7 @@ fn genBody(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, OutO
18251825 .tag_name => try airTagName(f, inst),
18261826 .error_name => try airErrorName(f, inst),
18271827 .splat => try airSplat(f, inst),
1828 .select => try airSelect(f, inst),
18281829 .shuffle => try airShuffle(f, inst),
18291830 .reduce => try airReduce(f, inst),
18301831 .aggregate_init => try airAggregateInit(f, inst),
......@@ -3794,6 +3795,21 @@ fn airSplat(f: *Function, inst: Air.Inst.Index) !CValue {
37943795 return f.fail("TODO: C backend: implement airSplat", .{});
37953796}
37963797
3798fn airSelect(f: *Function, inst: Air.Inst.Index) !CValue {
3799 if (f.liveness.isUnused(inst)) return CValue.none;
3800
3801 const inst_ty = f.air.typeOfIndex(inst);
3802 const ty_pl = f.air.instructions.items(.data)[inst].ty_pl;
3803
3804 const writer = f.object.writer();
3805 const local = try f.allocLocal(inst_ty, .Const);
3806 try writer.writeAll(" = ");
3807
3808 _ = local;
3809 _ = ty_pl;
3810 return f.fail("TODO: C backend: implement airSelect", .{});
3811}
3812
37973813fn airShuffle(f: *Function, inst: Air.Inst.Index) !CValue {
37983814 if (f.liveness.isUnused(inst)) return CValue.none;
37993815
src/codegen/llvm.zig+13
......@@ -3444,6 +3444,7 @@ pub const FuncGen = struct {
34443444 .tag_name => try self.airTagName(inst),
34453445 .error_name => try self.airErrorName(inst),
34463446 .splat => try self.airSplat(inst),
3447 .select => try self.airSelect(inst),
34473448 .shuffle => try self.airShuffle(inst),
34483449 .reduce => try self.airReduce(inst),
34493450 .aggregate_init => try self.airAggregateInit(inst),
......@@ -6355,6 +6356,18 @@ pub const FuncGen = struct {
63556356 return self.builder.buildShuffleVector(op_vector, undef_vector, mask_llvm_ty.constNull(), "");
63566357 }
63576358
6359 fn airSelect(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
6360 if (self.liveness.isUnused(inst)) return null;
6361
6362 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
6363 const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
6364 const pred = try self.resolveInst(pl_op.operand);
6365 const a = try self.resolveInst(extra.lhs);
6366 const b = try self.resolveInst(extra.rhs);
6367
6368 return self.builder.buildSelect(pred, a, b, "");
6369 }
6370
63586371 fn airShuffle(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
63596372 if (self.liveness.isUnused(inst)) return null;
63606373
src/print_air.zig+14
......@@ -264,6 +264,7 @@ const Writer = struct {
264264 .wasm_memory_size => try w.writeWasmMemorySize(s, inst),
265265 .wasm_memory_grow => try w.writeWasmMemoryGrow(s, inst),
266266 .mul_add => try w.writeMulAdd(s, inst),
267 .select => try w.writeSelect(s, inst),
267268 .shuffle => try w.writeShuffle(s, inst),
268269 .reduce => try w.writeReduce(s, inst),
269270 .cmp_vector => try w.writeCmpVector(s, inst),
......@@ -396,6 +397,19 @@ const Writer = struct {
396397 try s.print(", mask {d}, len {d}", .{ extra.mask, extra.mask_len });
397398 }
398399
400 fn writeSelect(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
401 const pl_op = w.air.instructions.items(.data)[inst].pl_op;
402 const extra = w.air.extraData(Air.Bin, pl_op.payload).data;
403
404 const elem_ty = w.air.typeOfIndex(inst).childType();
405 try s.print("{}, ", .{elem_ty.fmtDebug()});
406 try w.writeOperand(s, inst, 0, pl_op.operand);
407 try s.writeAll(", ");
408 try w.writeOperand(s, inst, 1, extra.lhs);
409 try s.writeAll(", ");
410 try w.writeOperand(s, inst, 2, extra.rhs);
411 }
412
399413 fn writeReduce(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
400414 const reduce = w.air.instructions.items(.data)[inst].reduce;
401415
test/behavior/select.zig+55-20
......@@ -3,24 +3,59 @@ const builtin = @import("builtin");
33const mem = std.mem;
44const expect = std.testing.expect;
55
6test "@select" {
7 if (@import("builtin").zig_backend != .stage1) return error.SkipZigTest; // TODO
8
9 const S = struct {
10 fn doTheTest() !void {
11 var a: @Vector(4, bool) = [4]bool{ true, false, true, false };
12 var b: @Vector(4, i32) = [4]i32{ -1, 4, 999, -31 };
13 var c: @Vector(4, i32) = [4]i32{ -5, 1, 0, 1234 };
14 var abc = @select(i32, a, b, c);
15 try expect(mem.eql(i32, &@as([4]i32, abc), &[4]i32{ -1, 1, 999, 1234 }));
16
17 var x: @Vector(4, bool) = [4]bool{ false, false, false, true };
18 var y: @Vector(4, f32) = [4]f32{ 0.001, 33.4, 836, -3381.233 };
19 var z: @Vector(4, f32) = [4]f32{ 0.0, 312.1, -145.9, 9993.55 };
20 var xyz = @select(f32, x, y, z);
21 try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
22 }
23 };
24 try S.doTheTest();
25 comptime try S.doTheTest();
6test "@select vectors" {
7 if (builtin.zig_backend == .stage2_c) return error.SkipZigTest; // TODO
8 if (builtin.zig_backend == .stage2_wasm) return error.SkipZigTest; // TODO
9 if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest; // TODO
10 if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
11 if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
12
13 comptime try selectVectors();
14 try selectVectors();
15}
16
17fn selectVectors() !void {
18 var a = @Vector(4, bool){ true, false, true, false };
19 var b = @Vector(4, i32){ -1, 4, 999, -31 };
20 var c = @Vector(4, i32){ -5, 1, 0, 1234 };
21 var abc = @select(i32, a, b, c);
22 try expect(abc[0] == -1);
23 try expect(abc[1] == 1);
24 try expect(abc[2] == 999);
25 try expect(abc[3] == 1234);
26
27 var x = @Vector(4, bool){ false, false, false, true };
28 var y = @Vector(4, f32){ 0.001, 33.4, 836, -3381.233 };
29 var z = @Vector(4, f32){ 0.0, 312.1, -145.9, 9993.55 };
30 var xyz = @select(f32, x, y, z);
31 try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
32}
33
34test "@select arrays" {
35 if (builtin.zig_backend == .stage1) return error.SkipZigTest; // TODO
36 if (builtin.zig_backend == .stage2_c) return error.SkipZigTest; // TODO
37 if (builtin.zig_backend == .stage2_wasm) return error.SkipZigTest; // TODO
38 if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest; // TODO
39 if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
40 if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
41
42 comptime try selectArrays();
43 try selectArrays();
44}
45
46fn selectArrays() !void {
47 var a = [4]bool{ false, true, false, true };
48 var b = [4]usize{ 0, 1, 2, 3 };
49 var c = [4]usize{ 4, 5, 6, 7 };
50 var abc = @select(usize, a, b, c);
51 try expect(abc[0] == 4);
52 try expect(abc[1] == 1);
53 try expect(abc[2] == 6);
54 try expect(abc[3] == 3);
55
56 var x = [4]bool{ false, false, false, true };
57 var y = [4]f32{ 0.001, 33.4, 836, -3381.233 };
58 var z = [4]f32{ 0.0, 312.1, -145.9, 9993.55 };
59 var xyz = @select(f32, x, y, z);
60 try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
2661}