stage2: implement @select

This commit is contained in:
John Schmidt 2022-03-23 21:41:35 +01:00
parent 1c33ea2c35
commit 12d5efcbe6
12 changed files with 223 additions and 21 deletions

View File

@ -553,6 +553,9 @@ pub const Inst = struct {
/// Constructs a vector by selecting elements from `a` and `b` based on `mask`.
/// Uses the `ty_pl` field with payload `Shuffle`.
shuffle,
/// Constructs a vector element-wise from `a` or `b` based on `pred`.
/// Uses the `ty_pl` field with payload `Select`.
select,
/// Given dest ptr, value, and len, set all elements at dest to value.
/// Result type is always void.
@ -785,6 +788,12 @@ pub const Shuffle = struct {
mask_len: u32,
};
pub const Select = struct {
pred: Inst.Ref,
a: Inst.Ref,
b: Inst.Ref,
};
pub const VectorCmp = struct {
lhs: Inst.Ref,
rhs: Inst.Ref,
@ -956,6 +965,7 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
.cmpxchg_weak,
.cmpxchg_strong,
.slice,
.select,
.shuffle,
.aggregate_init,
.union_init,

View File

@ -433,6 +433,10 @@ fn analyzeInst(
}
return extra_tombs.finish();
},
.select => {
const extra = a.air.extraData(Air.Select, inst_datas[inst].ty_pl.payload).data;
return trackOperands(a, new_set, inst, main_tomb, .{ extra.pred, extra.a, extra.b });
},
.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 });

View File

@ -14804,8 +14804,105 @@ fn analyzeShuffle(
fn zirSelect(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.zirSelect", .{});
const extra = sema.code.extraData(Zir.Inst.Select, inst_data.payload_index).data;
const elem_ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
const pred_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = inst_data.src_node };
const b_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);
const pred = sema.resolveInst(extra.pred);
const a = sema.resolveInst(extra.a);
const b = sema.resolveInst(extra.b);
const target = sema.mod.getTarget();
const pred_ty = sema.typeOf(pred);
switch (try pred_ty.zigTypeTagOrPoison()) {
.Vector => {
const scalar_ty = pred_ty.childType();
if (!scalar_ty.eql(Type.bool, target)) {
const bool_vec_ty = try Type.vector(sema.arena, pred_ty.vectorLen(), Type.bool);
return sema.fail(block, pred_src, "Expected '{}', found '{}'", .{ bool_vec_ty.fmt(target), pred_ty.fmt(target) });
}
},
else => return sema.fail(block, pred_src, "Expected vector type, found '{}'", .{pred_ty.fmt(target)}),
}
const vec_len = pred_ty.vectorLen();
const vec_ty = try Type.vector(sema.arena, vec_len, elem_ty);
const a_ty = sema.typeOf(a);
if (!a_ty.eql(vec_ty, target)) {
return sema.fail(block, a_src, "Expected '{}', found '{}'", .{ vec_ty.fmt(target), a_ty.fmt(target) });
}
const b_ty = sema.typeOf(b);
if (!b_ty.eql(vec_ty, target)) {
return sema.fail(block, b_src, "Expected '{}', found '{}'", .{ vec_ty.fmt(target), b_ty.fmt(target) });
}
const maybe_pred = try sema.resolveMaybeUndefVal(block, pred_src, pred);
const maybe_a = try sema.resolveMaybeUndefVal(block, a_src, a);
const maybe_b = try sema.resolveMaybeUndefVal(block, b_src, b);
const runtime_src = if (maybe_pred) |pred_val| rs: {
if (pred_val.isUndef()) return sema.addConstUndef(vec_ty);
if (maybe_a) |a_val| {
if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
if (maybe_b) |b_val| {
if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
var buf: Value.ElemValueBuffer = undefined;
const elems = try sema.gpa.alloc(Value, vec_len);
for (elems) |*elem, i| {
const pred_elem_val = pred_val.elemValueBuffer(i, &buf);
const should_choose_a = pred_elem_val.toBool();
if (should_choose_a) {
elem.* = a_val.elemValueBuffer(i, &buf);
} else {
elem.* = b_val.elemValueBuffer(i, &buf);
}
}
return sema.addConstant(
vec_ty,
try Value.Tag.aggregate.create(sema.arena, elems),
);
} else {
break :rs b_src;
}
} else {
if (maybe_b) |b_val| {
if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
}
break :rs a_src;
}
} else rs: {
if (maybe_a) |a_val| {
if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
}
if (maybe_b) |b_val| {
if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
}
break :rs pred_src;
};
try sema.requireRuntimeBlock(block, runtime_src);
return block.addInst(.{
.tag = .select,
.data = .{ .ty_pl = .{
.ty = try block.sema.addType(vec_ty),
.payload = try block.sema.addExtra(Air.Select{
.pred = pred,
.a = a,
.b = b,
}),
} },
});
}
fn zirAtomicLoad(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {

View File

@ -640,6 +640,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),
.select => try self.airSelect(inst),
.shuffle => try self.airShuffle(inst),
.reduce => try self.airReduce(inst),
.aggregate_init => try self.airAggregateInit(inst),
@ -3746,6 +3747,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
}
fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for {}", .{self.target.cpu.arch});
return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
}
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});

View File

@ -630,6 +630,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),
.select => try self.airSelect(inst),
.shuffle => try self.airShuffle(inst),
.reduce => try self.airReduce(inst),
.aggregate_init => try self.airAggregateInit(inst),
@ -4323,6 +4324,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
}
fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for arm", .{});
return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
}
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", .{});

View File

@ -600,6 +600,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),
.select => try self.airSelect(inst),
.shuffle => try self.airShuffle(inst),
.reduce => try self.airReduce(inst),
.aggregate_init => try self.airAggregateInit(inst),
@ -2396,6 +2397,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
}
fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for riscv64", .{});
return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
}
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", .{});

View File

@ -1371,6 +1371,7 @@ fn genInst(self: *Self, inst: Air.Inst.Index) !WValue {
.ret_ptr => self.airRetPtr(inst),
.ret_load => self.airRetLoad(inst),
.splat => self.airSplat(inst),
.select => self.airSelect(inst),
.shuffle => self.airShuffle(inst),
.reduce => self.airReduce(inst),
.aggregate_init => self.airAggregateInit(inst),
@ -3265,6 +3266,16 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
return self.fail("TODO: Implement wasm airSplat", .{});
}
fn airSelect(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
const ty = try self.resolveInst(ty_pl.ty);
_ = ty;
return self.fail("TODO: Implement wasm airSelect", .{});
}
fn airShuffle(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
if (self.liveness.isUnused(inst)) return WValue{ .none = {} };

View File

@ -721,6 +721,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),
.select => try self.airSelect(inst),
.shuffle => try self.airShuffle(inst),
.reduce => try self.airReduce(inst),
.aggregate_init => try self.airAggregateInit(inst),
@ -5678,6 +5679,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
}
fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
const extra = self.air.extraData(Air.Select, ty_pl.payload).data;
const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for x86_64", .{});
return self.finishAir(inst, result, .{ extra.pred, extra.a, extra.b });
}
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", .{});

View File

@ -1825,6 +1825,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),
.select => try airSelect(f, inst),
.shuffle => try airShuffle(f, inst),
.reduce => try airReduce(f, inst),
.aggregate_init => try airAggregateInit(f, inst),
@ -3794,6 +3795,21 @@ fn airSplat(f: *Function, inst: Air.Inst.Index) !CValue {
return f.fail("TODO: C backend: implement airSplat", .{});
}
fn airSelect(f: *Function, inst: Air.Inst.Index) !CValue {
if (f.liveness.isUnused(inst)) return CValue.none;
const inst_ty = f.air.typeOfIndex(inst);
const ty_pl = f.air.instructions.items(.data)[inst].ty_pl;
const writer = f.object.writer();
const local = try f.allocLocal(inst_ty, .Const);
try writer.writeAll(" = ");
_ = local;
_ = ty_pl;
return f.fail("TODO: C backend: implement airSelect", .{});
}
fn airShuffle(f: *Function, inst: Air.Inst.Index) !CValue {
if (f.liveness.isUnused(inst)) return CValue.none;

View File

@ -3444,6 +3444,7 @@ pub const FuncGen = struct {
.tag_name => try self.airTagName(inst),
.error_name => try self.airErrorName(inst),
.splat => try self.airSplat(inst),
.select => try self.airSelect(inst),
.shuffle => try self.airShuffle(inst),
.reduce => try self.airReduce(inst),
.aggregate_init => try self.airAggregateInit(inst),
@ -6355,6 +6356,18 @@ pub const FuncGen = struct {
return self.builder.buildShuffleVector(op_vector, undef_vector, mask_llvm_ty.constNull(), "");
}
fn airSelect(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.Select, ty_pl.payload).data;
const pred = try self.resolveInst(extra.pred);
const a = try self.resolveInst(extra.a);
const b = try self.resolveInst(extra.b);
return self.builder.buildSelect(pred, a, b, "");
}
fn airShuffle(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
if (self.liveness.isUnused(inst)) return null;

View File

@ -264,6 +264,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),
.select => try w.writeSelect(s, inst),
.shuffle => try w.writeShuffle(s, inst),
.reduce => try w.writeReduce(s, inst),
.cmp_vector => try w.writeCmpVector(s, inst),
@ -396,6 +397,18 @@ const Writer = struct {
try s.print(", mask {d}, len {d}", .{ extra.mask, extra.mask_len });
}
fn writeSelect(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
const ty_pl = w.air.instructions.items(.data)[inst].ty_pl;
const extra = w.air.extraData(Air.Select, ty_pl.payload).data;
try s.print("{}, ", .{w.air.getRefType(ty_pl.ty).fmtDebug()});
try w.writeOperand(s, inst, 0, extra.pred);
try s.writeAll(", ");
try w.writeOperand(s, inst, 1, extra.a);
try s.writeAll(", ");
try w.writeOperand(s, inst, 2, extra.b);
}
fn writeReduce(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
const reduce = w.air.instructions.items(.data)[inst].reduce;

View File

@ -4,23 +4,29 @@ const mem = std.mem;
const expect = std.testing.expect;
test "@select" {
if (@import("builtin").zig_backend != .stage1) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_c) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_wasm) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
const S = struct {
fn doTheTest() !void {
var a: @Vector(4, bool) = [4]bool{ true, false, true, false };
var b: @Vector(4, i32) = [4]i32{ -1, 4, 999, -31 };
var c: @Vector(4, i32) = [4]i32{ -5, 1, 0, 1234 };
var abc = @select(i32, a, b, c);
try expect(mem.eql(i32, &@as([4]i32, abc), &[4]i32{ -1, 1, 999, 1234 }));
var x: @Vector(4, bool) = [4]bool{ false, false, false, true };
var y: @Vector(4, f32) = [4]f32{ 0.001, 33.4, 836, -3381.233 };
var z: @Vector(4, f32) = [4]f32{ 0.0, 312.1, -145.9, 9993.55 };
var xyz = @select(f32, x, y, z);
try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
}
};
try S.doTheTest();
comptime try S.doTheTest();
try doTheTest();
comptime try doTheTest();
}
fn doTheTest() !void {
var a = @Vector(4, bool){ true, false, true, false };
var b = @Vector(4, i32){ -1, 4, 999, -31 };
var c = @Vector(4, i32){ -5, 1, 0, 1234 };
var abc = @select(i32, a, b, c);
try expect(abc[0] == -1);
try expect(abc[1] == 1);
try expect(abc[2] == 999);
try expect(abc[3] == 1234);
var x = @Vector(4, bool){ false, false, false, true };
var y = @Vector(4, f32){ 0.001, 33.4, 836, -3381.233 };
var z = @Vector(4, f32){ 0.0, 312.1, -145.9, 9993.55 };
var xyz = @select(f32, x, y, z);
try expect(mem.eql(f32, &@as([4]f32, xyz), &[4]f32{ 0.0, 312.1, -145.9, -3381.233 }));
}