Skip to content

chore: Cleanup unrolling pass #6743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TomAFrench merged 2 commits into from
Dec 9, 2024
Merged

chore: Cleanup unrolling pass #6743

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e61c388
Unrolling cleanup
jfecher Dec 9, 2024
b3e080e
Update test
jfecher Dec 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 26 additions & 43 deletions compiler/noirc_evaluator/src/ssa/opt/unrolling.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -279,10 +279,10 @@ impl Loop {
&self,
function: &Function,
cfg: &ControlFlowGraph,
) -> Result<Option<FieldElement>, CallStack> {
let pre_header = self.get_pre_header(function, cfg)?;
let jump_value = get_induction_variable(function, pre_header)?;
Ok(function.dfg.get_numeric_constant(jump_value))
) -> Option<FieldElement> {
let pre_header = self.get_pre_header(function, cfg).ok()?;
let jump_value = get_induction_variable(function, pre_header).ok()?;
function.dfg.get_numeric_constant(jump_value)
}

/// Find the upper bound of the loop in the loop header and return it
Expand Down Expand Up @@ -327,14 +327,10 @@ impl Loop {
&self,
function: &Function,
cfg: &ControlFlowGraph,
) -> Result<Option<(FieldElement, FieldElement)>, CallStack> {
let Some(lower) = self.get_const_lower_bound(function, cfg)? else {
return Ok(None);
};
let Some(upper) = self.get_const_upper_bound(function) else {
return Ok(None);
};
Ok(Some((lower, upper)))
) -> Option<(FieldElement, FieldElement)> {
let lower = self.get_const_lower_bound(function, cfg)?;
let upper = self.get_const_upper_bound(function)?;
Some((lower, upper))
}

/// Unroll a single loop in the function.
Expand Down Expand Up @@ -547,32 +543,29 @@ impl Loop {
&self,
function: &Function,
cfg: &ControlFlowGraph,
) -> Result<HashSet<ValueId>, CallStack> {
) -> Option<HashSet<ValueId>> {
// We need to traverse blocks from the pre-header up to the block entry point.
let pre_header = self.get_pre_header(function, cfg)?;
let pre_header = self.get_pre_header(function, cfg).ok()?;
let function_entry = function.entry_block();

// The algorithm in `find_blocks_in_loop` expects to collect the blocks between the header and the back-edge of the loop,
// but technically works the same if we go from the pre-header up to the function entry as well.
let blocks = Self::find_blocks_in_loop(function_entry, pre_header, cfg).blocks;

// Collect allocations in all blocks above the header.
let allocations = blocks.iter().flat_map(|b| {
function.dfg[*b]
.instructions()
.iter()
let allocations = blocks.iter().flat_map(|block| {
let instructions = function.dfg[*block].instructions().iter();
instructions
.filter(|i| matches!(&function.dfg[**i], Instruction::Allocate))
.map(|i| {
// Get the value into which the allocation was stored.
function.dfg.instruction_results(*i)[0]
})
// Get the value into which the allocation was stored.
.map(|i| function.dfg.instruction_results(*i)[0])
});

// Collect reference parameters of the function itself.
let params =
function.parameters().iter().filter(|p| function.dfg.value_is_reference(**p)).copied();

Ok(params.chain(allocations).collect())
Some(params.chain(allocations).collect())
}

/// Count the number of load and store instructions of specific variables in the loop.
Expand Down Expand Up @@ -603,13 +596,11 @@ impl Loop {

/// Count the number of instructions in the loop, including the terminating jumps.
fn count_all_instructions(&self, function: &Function) -> usize {
self.blocks
.iter()
.map(|block| {
let block = &function.dfg[*block];
block.instructions().len() + block.terminator().map(|_| 1).unwrap_or_default()
})
.sum()
let iter = self.blocks.iter().map(|block| {
let block = &function.dfg[*block];
block.instructions().len() + block.terminator().is_some() as usize
});
iter.sum()
}

/// Count the number of increments to the induction variable.
Expand Down Expand Up @@ -640,18 +631,11 @@ impl Loop {
function: &Function,
cfg: &ControlFlowGraph,
) -> Option<BoilerplateStats> {
let Ok(Some((lower, upper))) = self.get_const_bounds(function, cfg) else {
return None;
};
let Some(lower) = lower.try_to_u64() else {
return None;
};
let Some(upper) = upper.try_to_u64() else {
return None;
};
let Ok(refs) = self.find_pre_header_reference_values(function, cfg) else {
return None;
};
let (lower, upper) = self.get_const_bounds(function, cfg)?;
let lower = lower.try_to_u64()?;
let upper = upper.try_to_u64()?;
let refs = self.find_pre_header_reference_values(function, cfg)?;

let (loads, stores) = self.count_loads_and_stores(function, &refs);
let increments = self.count_induction_increments(function);
let all_instructions = self.count_all_instructions(function);
Expand Down Expand Up @@ -1142,7 +1126,6 @@ mod tests {

let (lower, upper) = loops.yet_to_unroll[0]
.get_const_bounds(function, &loops.cfg)
.expect("should find bounds")
.expect("bounds are numeric const");

assert_eq!(lower, FieldElement::from(0u32));
Expand Down