Skip to content

Commit

Permalink
Loop Invariant Code Motion
Browse files Browse the repository at this point in the history 
- More bug fixes.
- Add a second test case that passes but the optimization
  doesn't seem to be doing anything. This needs further
  investigation as to why nothing is being hoisted
  into the pre-headers.
  • Loading branch information
@kevinjoseph1995
kevinjoseph1995 committed Jan 19, 2025
1 parent b30eac1 commit 2402e50
Show file tree
Hide file tree
Showing 2 changed files with 218 additions and 60 deletions.
8 changes: 6 additions & 2 deletions common/src/cfg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ impl NodeEntry for BasicBlock {
pub struct Dominators<'a> {
cfg: &'a Cfg,
// The dominator set for each node
// If an element is present in set set_per_node[i] then it is a dominator of node i
// If an element x is present in set set_per_node[i] then it is a dominator of node i. That is x dominates i.
pub set_per_node: Vec<HashSet<NodeIndex>>,
}

Expand Down Expand Up @@ -237,6 +237,7 @@ impl Cfg {
index: NodeIndex,
mut node_data: BasicBlock,
new_node_name: String,
parent_nodes_to_skip: HashSet<NodeIndex>,
) {
let node_name = self.dag.get_node_name(index).to_string();
node_data.instruction_stream.push(Instruction::Effect {
Expand All @@ -254,7 +255,10 @@ impl Cfg {
predecessor_indices: SmallVec::new(),
});
let current_node_parents = self.dag.nodes[index].predecessor_indices.clone();
for parent_index in current_node_parents.iter() {
for parent_index in current_node_parents
.iter()
.filter(|p| !parent_nodes_to_skip.contains(p))
{
{
let parent = &mut self.dag.nodes[*parent_index];
let successor_index = parent
Expand Down
270 changes: 212 additions & 58 deletions optimizations/src/loop_invariant_code_motion.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use common::{
BasicBlock,
};
use dataflow_analysis::{Definition, LiveVariableAnalysis, ReachingDefinitions};
use std::collections::HashSet;
use std::collections::{HashMap, HashSet};

pub struct LoopInvariantCodeMotionPass {}

Expand Down Expand Up @@ -89,6 +89,7 @@ struct InstructionSite {
struct LoopMetadata {
loop_header: usize, // The index of the loop header.
loop_invariant_instructions: Vec<InstructionSite>, // The set of loop invariant instructions.
loop_nodes: HashSet<usize>, // The set of nodes in the loop.
}

fn find_loop_invariant_instructions(
Expand Down Expand Up @@ -118,6 +119,7 @@ fn find_loop_invariant_instructions(
op: ValueOps::Phi, ..
} => continue,
bril_rs::Instruction::Value { args, dest, .. } => (args, dest),
bril_rs::Instruction::Constant { dest, .. } => (&vec![], dest),
_ => continue,
};
if loop_invariant_instructions.contains(&InstructionSite {
Expand All @@ -128,43 +130,34 @@ fn find_loop_invariant_instructions(
// The instruction is already marked as loop invariant.
continue;
}
let defintion_sites = args
.iter()
.map(|arg| {
let definition_sites = reaching_definitions_in[*loop_node]
.iter()
.filter(|definition| definition.destination_variable == *arg)
.collect::<Vec<_>>();
if definition_sites.len() == 1 {
Some(definition_sites[0])
// A instruction is loop invariant when all of its arguments is either constant or variable "var" satisfying:
// 1. All reaching definitions of var are outside the loop.
// 2. There is exactly one reaching definition of var and the definition is loop-invariant.
let is_loop_invariant: bool = {
args.is_empty() || // This is a constant instruction as there are no arguments
args.iter().all(|argument| {
let incoming_definitions = &reaching_definitions_in[*loop_node];
let matching_definitions: Vec<&Definition> =
incoming_definitions.iter().filter(|definition| {
definition.destination_variable == argument
}).collect();
if matching_definitions.len() == 0 {
return false
} else if matching_definitions.len() > 1 {
return false;
} else if !loop_nodes.contains(&matching_definitions[0].basic_block_index) {
return true;
} else {
None
let definition = &matching_definitions[0];
// The argument is already marked as loop invariant
return loop_invariant_instructions.contains(&InstructionSite {
destination_variable: definition.destination_variable.to_string(),
basic_block_index: definition.basic_block_index,
instruction_index: definition.instruction_index
});
}
})
.collect::<Vec<_>>();
let is_loop_invariant = defintion_sites.iter().all(|definition| {
if let Some(definition) = definition {
if !loop_nodes.contains(&definition.basic_block_index) {
// The definition is not in the loop.
true
} else if loop_invariant_instructions.iter().any(
|site: &InstructionSite| {
site.destination_variable == definition.destination_variable
&& site.basic_block_index == definition.basic_block_index
&& site.instruction_index == definition.instruction_index
},
) {
// The definition is in the loop, but it is already marked as loop invariant.
true
} else {
// The definition is in the loop, but it is not marked as loop invariant.
false
}
} else {
// The definition is not found or there are multiple definitions.
false
}
});
};
if is_loop_invariant {
let site = InstructionSite {
destination_variable: dest.clone(),
Expand Down Expand Up @@ -201,10 +194,9 @@ fn filter_loop_invariant_instructions(
// Need to filter the loop invariant instructions.
// If the destination of an LII is d, it needs to satisfy the following condition:
// 1. The instruction should not have any side effects(Already satisfied as the candidate instructions are Value instructions). ✅
// 2 There is only one definition of d in the loop (Already satisfied by being SSA). ✅
// 3. The definition of d dominates all uses of d in the loop (Already satisfied by being SSA). ✅

// 4. d's block dominates all loop exits. -> Not satisfied. ⏳
// 2 There is only one definition of d in the loop. ⏳
// 3. d's block dominates all loop exits. ⏳
// 4. The definition of d dominates all uses of d in the loop (Already satisfied by being SSA). ⏳

// We need to find the loop exits.
let loop_exits = loop_nodes
Expand All @@ -216,28 +208,114 @@ fn filter_loop_invariant_instructions(
})
.cloned()
.collect::<HashSet<_>>();
let loop_defitntions = {
let mut loop_defitntions: HashMap<String, Vec<(usize, usize)>> = HashMap::new();
loop_nodes
.iter()
.map(|loop_node| {
cfg.get_basic_block(*loop_node)
.instruction_stream
.iter()
.filter(|instruction| match instruction {
Instruction::Constant { .. } => true,
Instruction::Value { .. } => true,
Instruction::Effect { .. } => false,
})
.enumerate()
.map(|(instruction_index, instruction)| match instruction {
Instruction::Constant { dest, .. } => {
(dest.clone(), instruction_index, *loop_node)
}
Instruction::Value { dest, .. } => {
(dest.clone(), instruction_index, *loop_node)
}
_ => {
panic!(
"Unexpected effect instruction found in loop node: {:?}",
instruction
);
}
})
})
.flatten()
.for_each(|(dest, instruction_index, loop_node)| {
loop_defitntions
.entry(dest)
.or_insert_with(Vec::new)
.push((loop_node, instruction_index));
});
loop_defitntions
};

let uses = {
let mut uses: HashMap<String, Vec<(usize, usize)>> = HashMap::new();
(0..cfg.dag.number_of_nodes())
.map(|node| {
cfg.get_basic_block(node)
.instruction_stream
.iter()
.filter(|instruction| match instruction {
Instruction::Constant { .. } => false,
Instruction::Value { .. } => true,
Instruction::Effect { .. } => false,
})
.enumerate()
.map(move |(instruction_index, instruction)| match instruction {
Instruction::Value { args, .. } => (args.iter(), instruction_index, node),
Instruction::Effect { args, .. } => (args.iter(), instruction_index, node),
_ => panic!(
"Unexpected constant instruction found in loop node: {:?}",
instruction
),
})
})
.flatten()
.for_each(|(args, instruction_index, loop_node)| {
args.for_each(|arg| {
uses.entry(arg.to_string())
.or_insert_with(Vec::new)
.push((loop_node, instruction_index));
});
});
uses
};

return instructions
.into_iter()
.filter(|instruction| {
// Check if the instruction satisfies the fourth condition.
let instruction_block = instruction.basic_block_index;
let instruction_index = instruction.instruction_index;
let destination_variable = instruction.destination_variable.as_str();
loop_exits.iter().all(|exit| {
if cfg
.dag
.get_successor_indices(*exit)
.iter()
.filter(|&successor| !loop_nodes.contains(successor))
.any(|successor| live_variables_in[*successor].contains(destination_variable))
{
// The destination variable is live at the exit.
dominators.set_per_node[*exit].contains(&instruction_block)
} else {
// The destination variable is not live at the exit.
true
loop_defitntions[destination_variable].len() == 1 // There is only one definition of d in the loop ✅
&& loop_exits.iter().all(|exit| { // d's block dominates all loop exits ✅
if cfg
.dag
.get_successor_indices(*exit)
.iter()
.filter(|&successor| !loop_nodes.contains(successor))
.any(|successor| {
live_variables_in[*successor].contains(destination_variable)
})
{
// The destination variable is live at the exit.
dominators.set_per_node[*exit].contains(&instruction_block)
} else {
// The destination variable is not live at the exit.
true
}
}) && {
if uses.contains_key(destination_variable) {
uses[destination_variable].iter().all(|(use_block, use_index)| {
// The definition of d dominates all uses of d in the loop ✅
dominators.set_per_node[*use_block].contains(&instruction_block)
|| (instruction_block == *use_block
&& instruction_index < *use_index) // The definition is before the use in the same block.
})
} else {
// The destination variable is not used. So, it is safe to move it.
true
}
}
})
})
.collect();
}
Expand Down Expand Up @@ -270,18 +348,29 @@ impl LoopInvariantCodeMotionPass {
LoopMetadata {
loop_header: *loop_header,
loop_invariant_instructions,
loop_nodes,
}
})
.collect();
for loop_metadata in loop_metadata_vector {
let mut instructions = Vec::<Instruction>::new();
for site in loop_metadata.loop_invariant_instructions {
for site in &loop_metadata.loop_invariant_instructions {
let instruction = cfg
.get_basic_block_mut(site.basic_block_index)
.instruction_stream
.remove(site.instruction_index);
.get_basic_block(site.basic_block_index)
.instruction_stream[site.instruction_index]
.clone();
instructions.push(instruction);
}
loop_metadata
.loop_invariant_instructions
.iter()
.for_each(|site| {
let basic_block = cfg.get_basic_block_mut(site.basic_block_index);
basic_block
.instruction_stream
.remove(site.instruction_index);
});

let new_node_name = format!(
"preheader_{}",
cfg.dag.get_node_name(loop_metadata.loop_header)
Expand All @@ -293,6 +382,7 @@ impl LoopInvariantCodeMotionPass {
instruction_stream: instructions,
},
new_node_name,
loop_metadata.loop_nodes,
);
}
cfg
Expand Down Expand Up @@ -351,4 +441,68 @@ mod tests {
let output = get_program_output(optimized_program.clone(), &[]);
assert_eq!(output_original, output, "{}", optimized_program);
}

#[test]
fn test_loop_invariant_code_motion2() {
let program = parse_program(indoc::indoc! {r#"
@main {
#--------------------------------
# Constants
#--------------------------------
outerLimit: int = const 3;
innerLimit: int = const 2;
one: int = const 1;
offset: int = const 10;
#--------------------------------
# Variables
#--------------------------------
i: int = const 0;
j: int = const 0;
sum: int = const 0;
#--------------------------------
# Outer loop
#--------------------------------
.outerLoop:
condOuter: bool = lt i outerLimit;
br condOuter .outerBody .outerDone;
.outerBody:
# Reset j for the inner loop
j: int = const 0;
#-----------------------------
# Inner loop
#-----------------------------
.innerLoop:
condInner: bool = lt j innerLimit;
br condInner .innerBody .innerDone;
.innerBody:
# Example partial-invariant expression for the inner loop
temp: int = add i offset; # i + offset is invariant wrt j
sum: int = add sum temp; # Accumulate it
sum: int = add sum j; # Also add j (changes each iteration of inner loop)
# Increment j
j: int = add j one;
jmp .innerLoop;
.innerDone:
# Increment i before repeating outer loop
i: int = add i one;
jmp .outerLoop;
.outerDone:
# Print sum and return
print sum;
ret;
}
"#});
let output_original = get_program_output(program.clone(), &[]);
let optimized_program = super::LoopInvariantCodeMotionPass::new().apply(program);
let output = get_program_output(optimized_program.clone(), &[]);
assert_eq!(output_original, output, "{}", optimized_program);
}
}

0 comments on commit 2402e50

Please sign in to comment.