Skip to content

Commit

Permalink
YJIT: Delete otherwise-empty defer_compilation() blocks
Browse files Browse the repository at this point in the history 
Calls to defer_compilation() leave behind a stub and a `struct Block`
that we retain. If the block is empty, it's only reason for existing is
to hold the `struct Branch` that the stub needs.

This patch transplants the branch out of the empty block into the newly
generated block when the defer_compilation() stub is hit, and deletes
the empty block to save memory.

To assist the transplantation, `Block::outgoing` is now a
`MutableBranchList`, and `Branch::Block` now in a `Cell`. These types
don't incur a size cost.

On the `lobsters` benchmark, `yjit_alloc_size` is roughly 97.7% of what
it was before the change.

Co-authored-by: Kevin Menard <[email protected]>
Co-authored-by: Randy Stauner <[email protected]>
Co-authored-by: Maxime Chevalier-Boisvert <[email protected]>
  • Loading branch information
@XrXr @nirvdrum
@rwstauner @maximecb
4 people committed Jun 11, 2024
1 parent 792e9c4 commit ff8529c
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 27 deletions.
1 change: 1 addition & 0 deletions yjit.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,7 @@ def _print_stats(out: $stderr) # :nodoc:
out.puts "compiled_iseq_count: " + format_number(13, stats[:compiled_iseq_count])
out.puts "compiled_blockid_count:" + format_number(13, stats[:compiled_blockid_count])
out.puts "compiled_block_count: " + format_number(13, stats[:compiled_block_count])
out.puts "deleted_defer_block_count:" + format_number_pct(10, stats[:deleted_defer_block_count], stats[:compiled_block_count])
if stats[:compiled_blockid_count] != 0
out.puts "versions_per_block: " + format_number(13, "%4.3f" % (stats[:compiled_block_count].fdiv(stats[:compiled_blockid_count])))
end
Expand Down
109 changes: 82 additions & 27 deletions yjit/src/core.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -693,7 +693,7 @@ struct BranchStub {
/// Note: care must be taken to minimize the size of branch objects
pub struct Branch {
// Block this is attached to
block: BlockRef,
block: Cell<BlockRef>,

// Positions where the generated code starts and ends
start_addr: CodePtr,
Expand Down Expand Up @@ -830,7 +830,7 @@ impl PendingBranch {
fn into_branch(mut self, uninit_block: BlockRef) -> BranchRef {
// Make the branch
let branch = Branch {
block: uninit_block,
block: Cell::new(uninit_block),
start_addr: self.start_addr.get().unwrap(),
end_addr: Cell::new(self.end_addr.get().unwrap()),
targets: self.targets,
Expand Down Expand Up @@ -919,14 +919,11 @@ pub struct Block {
end_addr: Cell<CodePtr>,

// List of incoming branches (from predecessors)
// These are reference counted (ownership shared between predecessor and successors)
incoming: MutableBranchList,

// NOTE: we might actually be able to store the branches here without refcounting
// however, using a RefCell makes it easy to get a pointer to Branch objects
//
// List of outgoing branches (to successors)
outgoing: Box<[BranchRef]>,
// Infrequently mutated for control flow graph edits for saving memory.
outgoing: MutableBranchList,

// FIXME: should these be code pointers instead?
// Offsets for GC managed objects in the mainline code block
Expand Down Expand Up @@ -993,6 +990,26 @@ impl MutableBranchList {
current_list.push(branch);
self.0.set(current_list.into_boxed_slice());
}

/// Iterate through branches in the list by moving out of the cell
/// and then putting it back when done. Modifications to this cell
/// during iteration will be discarded.
///
/// Assumes panic=abort since panic=unwind during iteration would
/// leave the cell empty.
fn for_each(&self, mut f: impl FnMut(BranchRef)) {
let list = self.0.take();
for branch in list.iter() {
f(*branch);
}
self.0.set(list);
}

/// Length of the list.
fn len(&self) -> usize {
// SAFETY: No cell mutation inside unsafe.
unsafe { self.0.ref_unchecked().len() }
}
}

impl fmt::Debug for MutableBranchList {
Expand Down Expand Up @@ -1227,7 +1244,7 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
}

// Mark outgoing branch entries
for branch in block.outgoing.iter() {
block.outgoing.for_each(|branch| {
let branch = unsafe { branch.as_ref() };
for target in branch.targets.iter() {
// SAFETY: no mutation inside unsafe
Expand All @@ -1247,7 +1264,7 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
unsafe { rb_gc_mark_movable(target_iseq.into()) };
}
}
}
});

// Mark references to objects in generated code.
// Skip for dead blocks since they shouldn't run.
Expand Down Expand Up @@ -1328,7 +1345,7 @@ pub extern "C" fn rb_yjit_iseq_update_references(iseq: IseqPtr) {
}

// Update outgoing branch entries
for branch in block.outgoing.iter() {
block.outgoing.for_each(|branch| {
let branch = unsafe { branch.as_ref() };
for target in branch.targets.iter() {
// SAFETY: no mutation inside unsafe
Expand All @@ -1352,7 +1369,7 @@ pub extern "C" fn rb_yjit_iseq_update_references(iseq: IseqPtr) {
unsafe { target.ref_unchecked().as_ref().unwrap().set_iseq(updated_iseq) };
}
}
}
});

// Update references to objects in generated code.
// Skip for dead blocks since they shouldn't run and
Expand Down Expand Up @@ -1632,11 +1649,11 @@ impl JITState {
entry_exit: self.get_block_entry_exit(),
cme_dependencies: self.method_lookup_assumptions.into_iter().map(Cell::new).collect(),
// Pending branches => actual branches
outgoing: self.pending_outgoing.into_iter().map(|pending_out| {
outgoing: MutableBranchList(Cell::new(self.pending_outgoing.into_iter().map(|pending_out| {
let pending_out = Rc::try_unwrap(pending_out)
.ok().expect("all PendingBranchRefs should be unique when ready to construct a Block");
pending_out.into_branch(NonNull::new(blockref as *mut Block).expect("no null from Box"))
}).collect()
}).collect()))
});
// Initialize it on the heap
// SAFETY: allocated with Box above
Expand Down Expand Up @@ -1681,10 +1698,10 @@ impl Block {

pub fn get_ctx_count(&self) -> usize {
let mut count = 1; // block.ctx
for branch in self.outgoing.iter() {
self.outgoing.for_each(|branch| {
// SAFETY: &self implies it's initialized
count += unsafe { branch.as_ref() }.get_stub_count();
}
});
count
}

Expand Down Expand Up @@ -2353,9 +2370,10 @@ fn gen_block_series_body(
let mut last_blockref = first_block;
loop {
// Get the last outgoing branch from the previous block.
let last_branchref = {
let last_block = unsafe { last_blockref.as_ref() };
match last_block.outgoing.last() {
// SAFETY: No cell mutation inside unsafe. Copying out a BranchRef.
let last_branchref: BranchRef = unsafe {
let last_block = last_blockref.as_ref();
match last_block.outgoing.0.ref_unchecked().last() {
Some(branch) => *branch,
None => {
break;
Expand Down Expand Up @@ -2645,7 +2663,7 @@ fn regenerate_branch(cb: &mut CodeBlock, branch: &Branch) {
cb.remove_comments(branch.start_addr, branch.end_addr.get());

// SAFETY: having a &Branch implies branch.block is initialized.
let block = unsafe { branch.block.as_ref() };
let block = unsafe { branch.block.get().as_ref() };

let branch_terminates_block = branch.end_addr.get() == block.get_end_addr();

Expand Down Expand Up @@ -2747,7 +2765,6 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
// code calls this function.
let branch = unsafe { branch_ref.as_ref() };
let branch_size_on_entry = branch.code_size();
let housing_block = unsafe { branch.block.as_ref() };

let target_idx: usize = target_idx.as_usize();
let target_branch_shape = match target_idx {
Expand Down Expand Up @@ -2820,7 +2837,7 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
// If the new block can be generated right after the branch (at cb->write_pos)
if cb.get_write_ptr() == branch.end_addr.get() {
// This branch should be terminating its block
assert!(branch.end_addr == housing_block.end_addr);
assert!(branch.end_addr == unsafe { branch.block.get().as_ref() }.end_addr);

// Change the branch shape to indicate the target block will be placed next
branch.gen_fn.set_shape(target_branch_shape);
Expand Down Expand Up @@ -2854,6 +2871,44 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
// Branch shape should reflect layout
assert!(!(branch.gen_fn.get_shape() == target_branch_shape && new_block.start_addr != branch.end_addr.get()));

// If we've hit a deferred branch, and the housing block is solely the branch,
// rewire incoming branches to the new block and delete the housing block.
{
// This &Block should be unique, relying on the VM lock
let housing_block: &Block = unsafe { branch.block.get().as_ref() };
if target_ctx.is_deferred() &&
target_blockid == housing_block.get_blockid() &&
housing_block.outgoing.len() == 1
{
// Divert incoming branches of housing_block to the new block
housing_block.incoming.for_each(|incoming| {
let incoming = unsafe { incoming.as_ref() };
for target in 0..incoming.targets.len() {
// SAFETY: No cell mutation; copying out a BlockRef.
if Some(BlockRef::from(housing_block)) == unsafe {
incoming.targets[target]
.ref_unchecked()
.as_ref()
.and_then(|target| target.get_block())
} {
incoming.targets[target].set(Some(Box::new(BranchTarget::Block(new_block.into()))));
}
}
new_block.push_incoming(incoming.into());
});

// Transplant the branch we've just hit to the new block
mem::drop(housing_block.outgoing.0.take());
new_block.outgoing.push(branch_ref);
let housing_block: BlockRef = branch.block.replace(new_block.into());
// Free the old housing block. There should now be no live &Block.
remove_block_version(&housing_block);
unsafe { free_block(housing_block, false) };

incr_counter!(deleted_defer_block_count);
}
}

// Add this branch to the list of incoming branches for the target
new_block.push_incoming(branch_ref);

Expand Down Expand Up @@ -3142,7 +3197,7 @@ pub fn defer_compilation(
idx: jit.get_insn_idx(),
};

// Likely a stub due to the increased chain depth
// Likely a stub since the context is marked as deferred().
let target0_address = branch.set_target(0, blockid, &next_ctx, ocb);

// Pad the block if it has the potential to be invalidated. This must be
Expand Down Expand Up @@ -3197,7 +3252,7 @@ unsafe fn remove_from_graph(blockref: BlockRef) {
}

// For each outgoing branch
for out_branchref in block.outgoing.iter() {
block.outgoing.for_each(|out_branchref| {
let out_branch = unsafe { out_branchref.as_ref() };
// For each successor block
for out_target in out_branch.targets.iter() {
Expand All @@ -3213,11 +3268,11 @@ unsafe fn remove_from_graph(blockref: BlockRef) {
// Temporarily move out of succ_block.incoming.
let succ_incoming = succ_block.incoming.0.take();
let mut succ_incoming = succ_incoming.into_vec();
succ_incoming.retain(|branch| branch != out_branchref);
succ_incoming.retain(|branch| *branch != out_branchref);
succ_block.incoming.0.set(succ_incoming.into_boxed_slice()); // allocs. Rely on oom=abort
}
}
}
});
}

/// Tear down a block and deallocate it.
Expand Down Expand Up @@ -3251,7 +3306,7 @@ pub unsafe fn free_block(blockref: BlockRef, graph_intact: bool) {
/// Caller must ensure that we have unique ownership for the referent block
unsafe fn dealloc_block(blockref: BlockRef) {
unsafe {
for outgoing in blockref.as_ref().outgoing.iter() {
for outgoing in blockref.as_ref().outgoing.0.take().iter() {
// this Box::from_raw matches the Box::into_raw from PendingBranch::into_branch
mem::drop(Box::from_raw(outgoing.as_ptr()));
}
Expand Down Expand Up @@ -3688,7 +3743,7 @@ mod tests {
// we're always working with &Branch (a shared reference to a Branch).
let branch: &Branch = &Branch {
gen_fn: BranchGenFn::JZToTarget0,
block,
block: Cell::new(block),
start_addr: dumm_addr,
end_addr: Cell::new(dumm_addr),
targets: [Cell::new(None), Cell::new(Some(Box::new(BranchTarget::Stub(Box::new(BranchStub {
Expand Down
1 change: 1 addition & 0 deletions yjit/src/stats.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -554,6 +554,7 @@ make_counters! {
block_next_count,
defer_count,
defer_empty_count,
deleted_defer_block_count,
branch_insn_count,
branch_known_count,
max_inline_versions,
Expand Down

0 comments on commit ff8529c

Please sign in to comment.