authorgravatar for git@vexu.euVeikka Tuominen <git@vexu.eu> 2022-03-11 14:25:15+02:00
committergravatar for git@vexu.euVeikka Tuominen <git@vexu.eu> 2022-03-11 14:25:15+02:00
log01cd4119b032f13899e9b7b30c7e093620058dfd
treebfafdc11324d3b77ff5af8b59b5ff324db66335b
parentcba68090a60c3de8eadbf8eb53e37620a1d66683

Sema: implement `@shuffle` at comptime and for differing lengths

3 files changed, 110 insertions(+), 34 deletions(-)

src/Sema.zig+73-10
......@@ -13454,8 +13454,6 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1345413454 const inst_data = sema.code.instructions.items(.data)[inst].pl_node;
1345513455 const extra = sema.code.extraData(Zir.Inst.Shuffle, inst_data.payload_index).data;
1345613456 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 };
1345913457 const mask_src: LazySrcLoc = .{ .node_offset_builtin_call_arg3 = inst_data.src_node };
1346013458
1346113459 const elem_ty = try sema.resolveType(block, elem_ty_src, extra.elem_type);
......@@ -13474,6 +13472,25 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1347413472 .elem_type = Type.@"i32",
1347513473 });
1347613474 mask = try sema.coerce(block, mask_ty, mask, mask_src);
13475 const mask_val = try sema.resolveConstMaybeUndefVal(block, mask_src, mask);
13476 return sema.analyzeShuffle(block, inst_data.src_node, elem_ty, a, b, mask_val, @intCast(u32, mask_len));
13477}
13478
13479fn analyzeShuffle(
13480 sema: *Sema,
13481 block: *Block,
13482 src_node: i32,
13483 elem_ty: Type,
13484 a_arg: Air.Inst.Ref,
13485 b_arg: Air.Inst.Ref,
13486 mask: Value,
13487 mask_len: u32,
13488) CompileError!Air.Inst.Ref {
13489 const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = src_node };
13490 const b_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = src_node };
13491 const mask_src: LazySrcLoc = .{ .node_offset_builtin_call_arg3 = src_node };
13492 var a = a_arg;
13493 var b = b_arg;
1347713494
1347813495 const res_ty = try Type.Tag.vector.create(sema.arena, .{
1347913496 .len = mask_len,
......@@ -13485,7 +13502,7 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1348513502 .Undefined => null,
1348613503 else => return sema.fail(block, a_src, "expected vector or array with element type {}, found {}", .{
1348713504 elem_ty,
13488 sema.typeOf(mask),
13505 sema.typeOf(a),
1348913506 }),
1349013507 };
1349113508 var maybe_b_len = switch (sema.typeOf(b).zigTypeTag()) {
......@@ -13493,7 +13510,7 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1349313510 .Undefined => null,
1349413511 else => return sema.fail(block, b_src, "expected vector or array with element type {}, found {}", .{
1349513512 elem_ty,
13496 sema.typeOf(mask),
13513 sema.typeOf(b),
1349713514 }),
1349813515 };
1349913516 if (maybe_a_len == null and maybe_b_len == null) {
......@@ -13519,11 +13536,10 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1351913536 .{ b_len, b_src, b_ty },
1352013537 };
1352113538
13522 const mask_val = try sema.resolveConstMaybeUndefVal(block, mask_src, mask);
1352313539 var i: usize = 0;
1352413540 while (i < mask_len) : (i += 1) {
1352513541 var buf: Value.ElemValueBuffer = undefined;
13526 const elem = mask_val.elemValueBuffer(i, &buf);
13542 const elem = mask.elemValueBuffer(i, &buf);
1352713543 if (elem.isUndef()) continue;
1352813544 const int = elem.toSignedInt();
1352913545 var unsigned: u32 = undefined;
......@@ -13555,14 +13571,61 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1355513571 }
1355613572 }
1355713573
13558 // TODO at comptime
13574 if (try sema.resolveMaybeUndefVal(block, a_src, a)) |a_val| {
13575 if (try sema.resolveMaybeUndefVal(block, b_src, b)) |b_val| {
13576 const values = try sema.arena.alloc(Value, mask_len);
13577
13578 i = 0;
13579 while (i < mask_len) : (i += 1) {
13580 var buf: Value.ElemValueBuffer = undefined;
13581 const mask_elem_val = mask.elemValueBuffer(i, &buf);
13582 if (mask_elem_val.isUndef()) {
13583 values[i] = Value.undef;
13584 continue;
13585 }
13586 const int = mask_elem_val.toSignedInt();
13587 const unsigned = if (int >= 0) @intCast(u32, int) else @intCast(u32, ~int);
13588 if (int >= 0) {
13589 values[i] = try a_val.elemValue(sema.arena, unsigned);
13590 } else {
13591 values[i] = try b_val.elemValue(sema.arena, unsigned);
13592 }
13593 }
13594 const res_val = try Value.Tag.array.create(sema.arena, values);
13595 return sema.addConstant(res_ty, res_val);
13596 }
13597 }
1355913598
13599 // All static analysis passed, and not comptime.
13600 // For runtime codegen, vectors a and b must be the same length. Here we
13601 // recursively @shuffle the smaller vector to append undefined elements
13602 // to it up to the length of the longer vector. This recursion terminates
13603 // in 1 call because these calls to analyzeShuffle guarantee a_len == b_len.
1356013604 if (a_len != b_len) {
13561 return sema.fail(block, mask_src, "TODO handle shuffle a_len != b_len", .{});
13605 const min_len = std.math.min(a_len, b_len);
13606 const max_len = std.math.max(a_len, b_len);
13607
13608 const expand_mask_values = try sema.arena.alloc(Value, max_len);
13609 i = 0;
13610 while (i < min_len) : (i += 1) {
13611 expand_mask_values[i] = try Value.Tag.int_u64.create(sema.arena, i);
13612 }
13613 while (i < max_len) : (i += 1) {
13614 expand_mask_values[i] = Value.negative_one;
13615 }
13616 const expand_mask = try Value.Tag.array.create(sema.arena, expand_mask_values);
13617
13618 if (a_len < b_len) {
13619 const undef = try sema.addConstUndef(a_ty);
13620 a = try sema.analyzeShuffle(block, src_node, elem_ty, a, undef, expand_mask, @intCast(u32, max_len));
13621 } else {
13622 const undef = try sema.addConstUndef(b_ty);
13623 b = try sema.analyzeShuffle(block, src_node, elem_ty, b, undef, expand_mask, @intCast(u32, max_len));
13624 }
1356213625 }
1356313626
1356413627 const mask_index = @intCast(u32, sema.air_values.items.len);
13565 try sema.air_values.append(sema.gpa, mask_val);
13628 try sema.air_values.append(sema.gpa, mask);
1356613629 return block.addInst(.{
1356713630 .tag = .shuffle,
1356813631 .data = .{ .ty_pl = .{
......@@ -13571,7 +13634,7 @@ fn zirShuffle(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air
1357113634 .a = a,
1357213635 .b = b,
1357313636 .mask = mask_index,
13574 .mask_len = @intCast(u32, mask_len),
13637 .mask_len = mask_len,
1357513638 }),
1357613639 } },
1357713640 });
test/behavior.zig+1-1
......@@ -151,6 +151,7 @@ test {
151151 _ = @import("behavior/bugs/2114.zig");
152152 _ = @import("behavior/bugs/3779.zig");
153153 _ = @import("behavior/bugs/10147.zig");
154 _ = @import("behavior/shuffle.zig");
154155 _ = @import("behavior/union_with_members.zig");
155156
156157 if (builtin.zig_backend == .stage1) {
......@@ -169,7 +170,6 @@ test {
169170 _ = @import("behavior/bugs/7027.zig");
170171 _ = @import("behavior/const_slice_child.zig");
171172 _ = @import("behavior/select.zig");
172 _ = @import("behavior/shuffle.zig");
173173 _ = @import("behavior/struct_contains_slice_of_itself.zig");
174174 _ = @import("behavior/typename.zig");
175175 _ = @import("behavior/vector.zig");
test/behavior/shuffle.zig+36-23
......@@ -4,12 +4,12 @@ const mem = std.mem;
44const expect = std.testing.expect;
55const Vector = std.meta.Vector;
66
7test "@shuffle" {
7test "@shuffle int" {
88 const S = struct {
99 fn doTheTest() !void {
1010 var v: Vector(4, i32) = [4]i32{ 2147483647, -2, 30, 40 };
1111 var x: Vector(4, i32) = [4]i32{ 1, 2147483647, 3, 4 };
12 const mask: Vector(4, i32) = [4]i32{ 0, ~@as(i32, 2), 3, ~@as(i32, 3) };
12 const mask = [4]i32{ 0, ~@as(i32, 2), 3, ~@as(i32, 3) };
1313 var res = @shuffle(i32, v, x, mask);
1414 try expect(mem.eql(i32, &@as([4]i32, res), &[4]i32{ 2147483647, 3, 40, 4 }));
1515
......@@ -18,40 +18,53 @@ test "@shuffle" {
1818 try expect(mem.eql(i32, &@as([4]i32, res), &[4]i32{ 2147483647, 3, 40, 4 }));
1919
2020 // Undefined
21 const mask2: Vector(4, i32) = [4]i32{ 3, 1, 2, 0 };
21 const mask2 = [4]i32{ 3, 1, 2, 0 };
2222 res = @shuffle(i32, v, undefined, mask2);
2323 try expect(mem.eql(i32, &@as([4]i32, res), &[4]i32{ 40, -2, 30, 2147483647 }));
2424
2525 // Upcasting of b
2626 var v2: Vector(2, i32) = [2]i32{ 2147483647, undefined };
27 const mask3: Vector(4, i32) = [4]i32{ ~@as(i32, 0), 2, ~@as(i32, 0), 3 };
27 const mask3 = [4]i32{ ~@as(i32, 0), 2, ~@as(i32, 0), 3 };
2828 res = @shuffle(i32, x, v2, mask3);
2929 try expect(mem.eql(i32, &@as([4]i32, res), &[4]i32{ 2147483647, 3, 2147483647, 4 }));
3030
3131 // Upcasting of a
3232 var v3: Vector(2, i32) = [2]i32{ 2147483647, -2 };
33 const mask4: Vector(4, i32) = [4]i32{ 0, ~@as(i32, 2), 1, ~@as(i32, 3) };
33 const mask4 = [4]i32{ 0, ~@as(i32, 2), 1, ~@as(i32, 3) };
3434 res = @shuffle(i32, v3, x, mask4);
3535 try expect(mem.eql(i32, &@as([4]i32, res), &[4]i32{ 2147483647, 3, -2, 4 }));
36 }
37 };
38 try S.doTheTest();
39 comptime try S.doTheTest();
40}
3641
37 // bool
38 {
39 var x2: Vector(4, bool) = [4]bool{ false, true, false, true };
40 var v4: Vector(2, bool) = [2]bool{ true, false };
41 const mask5: Vector(4, i32) = [4]i32{ 0, ~@as(i32, 1), 1, 2 };
42 var res2 = @shuffle(bool, x2, v4, mask5);
43 try expect(mem.eql(bool, &@as([4]bool, res2), &[4]bool{ false, false, true, false }));
44 }
45
46 // TODO re-enable when LLVM codegen is fixed
47 // https://github.com/ziglang/zig/issues/3246
48 if (false) {
49 var x2: Vector(3, bool) = [3]bool{ false, true, false };
50 var v4: Vector(2, bool) = [2]bool{ true, false };
51 const mask5: Vector(4, i32) = [4]i32{ 0, ~@as(i32, 1), 1, 2 };
52 var res2 = @shuffle(bool, x2, v4, mask5);
53 try expect(mem.eql(bool, &@as([4]bool, res2), &[4]bool{ false, false, true, false }));
54 }
42test "@shuffle bool" {
43 const S = struct {
44 fn doTheTest() !void {
45 var x: Vector(4, bool) = [4]bool{ false, true, false, true };
46 var v: Vector(2, bool) = [2]bool{ true, false };
47 const mask = [4]i32{ 0, ~@as(i32, 1), 1, 2 };
48 var res = @shuffle(bool, x, v, mask);
49 try expect(mem.eql(bool, &@as([4]bool, res), &[4]bool{ false, false, true, false }));
50 }
51 };
52 if (builtin.zig_backend == .stage1) try S.doTheTest();
53 comptime try S.doTheTest();
54}
55
56test "@shuffle bool" {
57 // TODO re-enable when LLVM codegen is fixed
58 // https://github.com/ziglang/zig/issues/3246
59 if (true) return error.SkipZigTest;
60
61 const S = struct {
62 fn doTheTest() !void {
63 var x: Vector(3, bool) = [3]bool{ false, true, false };
64 var v: Vector(2, bool) = [2]bool{ true, false };
65 const mask: Vector(4, i32) = [4]i32{ 0, ~@as(i32, 1), 1, 2 };
66 var res = @shuffle(bool, x, v, mask);
67 try expect(mem.eql(bool, &@as([4]bool, res), &[4]bool{ false, false, true, false }));
5568 }
5669 };
5770 try S.doTheTest();