Skip to content

Commit

Permalink
riscv: implement switch_dispatch & loop_switch_br
Browse files Browse the repository at this point in the history
  • Loading branch information
@Rexicon226 @mlugg
Rexicon226 authored and mlugg committed Sep 1, 2024
1 parent 97ed239 commit 0d295d7
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 15 deletions.
129 changes: 114 additions & 15 deletions src/arch/riscv64/CodeGen.zig
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ frame_allocs: std.MultiArrayList(FrameAlloc) = .{},
free_frame_indices: std.AutoArrayHashMapUnmanaged(FrameIndex, void) = .{},
frame_locs: std.MultiArrayList(Mir.FrameLoc) = .{},

loop_repeat_info: std.AutoHashMapUnmanaged(Air.Inst.Index, struct {
loops: std.AutoHashMapUnmanaged(Air.Inst.Index, struct {
/// The state to restore before branching.
state: State,
/// The branch target.
Expand Down Expand Up @@ -232,11 +232,12 @@ const MCValue = union(enum) {
.register,
.register_pair,
.register_offset,
.load_frame,
.load_symbol,
.load_tlv,
.indirect,
=> true,

.load_frame => |frame_addr| !frame_addr.index.isNamed(),
};
}

Expand Down Expand Up @@ -804,7 +805,7 @@ pub fn generate(
function.frame_allocs.deinit(gpa);
function.free_frame_indices.deinit(gpa);
function.frame_locs.deinit(gpa);
function.loop_repeat_info.deinit(gpa);
function.loops.deinit(gpa);
var block_it = function.blocks.valueIterator();
while (block_it.next()) |block| block.deinit(gpa);
function.blocks.deinit(gpa);
Expand Down Expand Up @@ -1588,7 +1589,7 @@ fn genBody(func: *Func, body: []const Air.Inst.Index) InnerError!void {
.block => try func.airBlock(inst),
.br => try func.airBr(inst),
.repeat => try func.airRepeat(inst),
.switch_dispatch => return func.fail("TODO implement `switch_dispatch`", .{}),
.switch_dispatch => try func.airSwitchDispatch(inst),
.trap => try func.airTrap(),
.breakpoint => try func.airBreakpoint(),
.ret_addr => try func.airRetAddr(inst),
Expand Down Expand Up @@ -1678,7 +1679,7 @@ fn genBody(func: *Func, body: []const Air.Inst.Index) InnerError!void {
.field_parent_ptr => try func.airFieldParentPtr(inst),

.switch_br => try func.airSwitchBr(inst),
.loop_switch_br => return func.fail("TODO implement `loop_switch_br`", .{}),
.loop_switch_br => try func.airLoopSwitchBr(inst),

.ptr_slice_len_ptr => try func.airPtrSliceLenPtr(inst),
.ptr_slice_ptr_ptr => try func.airPtrSlicePtrPtr(inst),
Expand Down Expand Up @@ -5610,11 +5611,11 @@ fn airLoop(func: *Func, inst: Air.Inst.Index) !void {
func.scope_generation += 1;
const state = try func.saveState();

try func.loop_repeat_info.putNoClobber(func.gpa, inst, .{
try func.loops.putNoClobber(func.gpa, inst, .{
.state = state,
.jmp_target = @intCast(func.mir_instructions.len),
});
defer assert(func.loop_repeat_info.remove(inst));
defer assert(func.loops.remove(inst));

try func.genBody(body);

Expand Down Expand Up @@ -5671,12 +5672,7 @@ fn lowerBlock(func: *Func, inst: Air.Inst.Index, body: []const Air.Inst.Index) !

fn airSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
const switch_br = func.air.unwrapSwitch(inst);

const liveness = try func.liveness.getSwitchBr(func.gpa, inst, switch_br.cases_len + 1);
defer func.gpa.free(liveness.deaths);

const condition = try func.resolveInst(switch_br.operand);
const condition_ty = func.typeOf(switch_br.operand);

// If the condition dies here in this switch instruction, process
// that death now instead of later as this has an effect on
Expand All @@ -5685,6 +5681,22 @@ fn airSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
if (switch_br.operand.toIndex()) |op_inst| try func.processDeath(op_inst);
}

try func.lowerSwitchBr(inst, switch_br, condition);

// We already took care of pl_op.operand earlier, so there's nothing left to do
func.finishAirBookkeeping();
}

fn lowerSwitchBr(
func: *Func,
inst: Air.Inst.Index,
switch_br: Air.UnwrappedSwitch,
condition: MCValue,
) !void {
const condition_ty = func.typeOf(switch_br.operand);
const liveness = try func.liveness.getSwitchBr(func.gpa, inst, switch_br.cases_len + 1);
defer func.gpa.free(liveness.deaths);

func.scope_generation += 1;
const state = try func.saveState();

Expand Down Expand Up @@ -5785,8 +5797,92 @@ fn airSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
.close_scope = true,
});
}
}

fn airLoopSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
const switch_br = func.air.unwrapSwitch(inst);
const condition = try func.resolveInst(switch_br.operand);

const mat_cond = if (condition.isMutable() and
func.reuseOperand(inst, switch_br.operand, 0, condition))
condition
else mat_cond: {
const ty = func.typeOf(switch_br.operand);
const mat_cond = try func.allocRegOrMem(ty, inst, true);
try func.genCopy(ty, mat_cond, condition);
break :mat_cond mat_cond;
};
func.inst_tracking.putAssumeCapacityNoClobber(inst, InstTracking.init(mat_cond));

// If the condition dies here in this switch instruction, process
// that death now instead of later as this has an effect on
// whether it needs to be spilled in the branches
if (func.liveness.operandDies(inst, 0)) {
if (switch_br.operand.toIndex()) |op_inst| try func.processDeath(op_inst);
}

func.scope_generation += 1;
const state = try func.saveState();

try func.loops.putNoClobber(func.gpa, inst, .{
.state = state,
.jmp_target = @intCast(func.mir_instructions.len),
});
defer assert(func.loops.remove(inst));

// Stop tracking block result without forgetting tracking info
try func.freeValue(mat_cond);

try func.lowerSwitchBr(inst, switch_br, mat_cond);

try func.processDeath(inst);
func.finishAirBookkeeping();
}

fn airSwitchDispatch(func: *Func, inst: Air.Inst.Index) !void {
const br = func.air.instructions.items(.data)[@intFromEnum(inst)].br;

const block_ty = func.typeOfIndex(br.block_inst);
const block_tracking = func.inst_tracking.getPtr(br.block_inst).?;
const loop_data = func.loops.getPtr(br.block_inst).?;
done: {
try func.getValue(block_tracking.short, null);
const src_mcv = try func.resolveInst(br.operand);

if (func.reuseOperandAdvanced(inst, br.operand, 0, src_mcv, br.block_inst)) {
try func.getValue(block_tracking.short, br.block_inst);
// .long = .none to avoid merging operand and block result stack frames.
const current_tracking: InstTracking = .{ .long = .none, .short = src_mcv };
try current_tracking.materializeUnsafe(func, br.block_inst, block_tracking.*);
for (current_tracking.getRegs()) |src_reg| func.register_manager.freeReg(src_reg);
break :done;
}

try func.getValue(block_tracking.short, br.block_inst);
const dst_mcv = block_tracking.short;
try func.genCopy(block_ty, dst_mcv, try func.resolveInst(br.operand));
break :done;
}

// Process operand death so that it is properly accounted for in the State below.
if (func.liveness.operandDies(inst, 0)) {
if (br.operand.toIndex()) |op_inst| try func.processDeath(op_inst);
}

try func.restoreState(loop_data.state, &.{}, .{
.emit_instructions = true,
.update_tracking = false,
.resurrect = false,
.close_scope = false,
});

// Emit a jump with a relocation. It will be patched up after the block ends.
// Leave the jump offset undefined
_ = try func.jump(loop_data.jmp_target);

// Stop tracking block result without forgetting tracking info
try func.freeValue(block_tracking.short);

// We already took care of pl_op.operand earlier, so there's nothing left to do
func.finishAirBookkeeping();
}

Expand Down Expand Up @@ -5867,7 +5963,7 @@ fn airBr(func: *Func, inst: Air.Inst.Index) !void {

fn airRepeat(func: *Func, inst: Air.Inst.Index) !void {
const loop_inst = func.air.instructions.items(.data)[@intFromEnum(inst)].repeat.loop_inst;
const repeat_info = func.loop_repeat_info.get(loop_inst).?;
const repeat_info = func.loops.get(loop_inst).?;
try func.restoreState(repeat_info.state, &.{}, .{
.emit_instructions = true,
.update_tracking = false,
Expand Down Expand Up @@ -8298,7 +8394,10 @@ fn typeOf(func: *Func, inst: Air.Inst.Ref) Type {

fn typeOfIndex(func: *Func, inst: Air.Inst.Index) Type {
const zcu = func.pt.zcu;
return func.air.typeOfIndex(inst, &zcu.intern_pool);
return switch (func.air.instructions.items(.tag)[@intFromEnum(inst)]) {
.loop_switch_br => func.typeOf(func.air.unwrapSwitch(inst).operand),
else => func.air.typeOfIndex(inst, &zcu.intern_pool),
};
}

fn hasFeature(func: *Func, feature: Target.riscv.Feature) bool {
Expand Down
1 change: 1 addition & 0 deletions test/behavior/switch_loop.zig
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ test "switch loop on tagged union" {
if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_sparc64) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_spirv64) return error.SkipZigTest; // TODO
if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;

const S = struct {
const U = union(enum) {
Expand Down

0 comments on commit 0d295d7

Please sign in to comment.