Skip to content
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

Jump to bottom

Fully port Depth and Size passes to Rust #13040

Merged
merged 4 commits into from
Aug 28, 2024
Merged

Fully port Depth and Size passes to Rust #13040

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5b41efa
Avoid isinstance check on every node in DAG during
kevinhartman Aug 26, 2024
80ff154
Avoid unnecessary downcast.
kevinhartman Aug 27, 2024
af244c2
Check if there is control flow before doing loops.
kevinhartman Aug 27, 2024
3c70630
Avoid py_op_nodes in DAGCircuit::count_ops.
kevinhartman Aug 27, 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
66 changes: 33 additions & 33 deletions crates/circuit/src/dag_circuit.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2168,14 +2168,15 @@ def _format(operand):
}

let circuit_to_dag = imports::CIRCUIT_TO_DAG.get_bound(py);
for node_index in
self.op_nodes_by_py_type(imports::CONTROL_FLOW_OP.get_bound(py).downcast()?, true)
{
let NodeType::Operation(node) = &self.dag[node_index] else {
return Err(DAGCircuitError::new_err("unknown control-flow type"));
for node in self.dag.node_weights() {
let NodeType::Operation(node) = node else {
continue;
};
if !node.op.control_flow() {
continue;
}
let OperationRef::Instruction(inst) = node.op.view() else {
unreachable!("Control Flow operations must be a PyInstruction");
panic!("control flow op must be an instruction");
};
let inst_bound = inst.instruction.bind(py);
if inst_bound.is_instance(imports::FOR_LOOP_OP.get_bound(py))? {
Expand Down Expand Up @@ -2239,34 +2240,33 @@ def _format(operand):
Ok(if recurse {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we want an && CONTROL_FLOW_OP_NAMES.iter().any(|x| self.op_names.contains_key(&x.to_string())) on this condition so we don't bother with the for loop if we know there are no control flow instructions. The same question applies to size() as well.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a good idea, and allowed me to make the code a bit cleaner across both methods. I added a helper as well, which is private to DAGCircuit for now.

Done in af244c2.

let circuit_to_dag = imports::CIRCUIT_TO_DAG.get_bound(py);
let mut node_lookup: HashMap<NodeIndex, usize> = HashMap::new();

for node_index in
self.op_nodes_by_py_type(imports::CONTROL_FLOW_OP.get_bound(py).downcast()?, true)
{
if let NodeType::Operation(node) = &self.dag[node_index] {
if let OperationRef::Instruction(inst) = node.op.view() {
let inst_bound = inst.instruction.bind(py);
let weight =
if inst_bound.is_instance(imports::FOR_LOOP_OP.get_bound(py))? {
node.params_view().len()
} else {
1
};
if weight == 0 {
node_lookup.insert(node_index, 0);
} else {
let raw_blocks = inst_bound.getattr("blocks")?;
let blocks = raw_blocks.downcast::<PyTuple>()?;
let mut block_weights: Vec<usize> = Vec::with_capacity(blocks.len());
for block in blocks.iter() {
let inner_dag: &DAGCircuit =
&circuit_to_dag.call1((block,))?.extract()?;
block_weights.push(inner_dag.depth(py, true)?);
}
node_lookup
.insert(node_index, weight * block_weights.iter().max().unwrap());
}
for (node_index, node) in self.dag.node_references() {
let NodeType::Operation(node) = node else {
continue;
};
if !node.op.control_flow() {
continue;
}
let OperationRef::Instruction(inst) = node.op.view() else {
panic!("control flow op must be an instruction")
};
let inst_bound = inst.instruction.bind(py);
let weight = if inst_bound.is_instance(imports::FOR_LOOP_OP.get_bound(py))? {
node.params_view().len()
} else {
1
};
if weight == 0 {
node_lookup.insert(node_index, 0);
} else {
let raw_blocks = inst_bound.getattr("blocks")?;
let blocks = raw_blocks.downcast::<PyTuple>()?;
let mut block_weights: Vec<usize> = Vec::with_capacity(blocks.len());
for block in blocks.iter() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need the PyTuple downcast here? Should we be able to just call len() and iter() on a Bound<PyAny>. I think the only difference is their fallible in PyAnyMethods but not in PyTupleMethods.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good idea. Done in 80ff154.

let inner_dag: &DAGCircuit = &circuit_to_dag.call1((block,))?.extract()?;
block_weights.push(inner_dag.depth(py, true)?);
}
node_lookup.insert(node_index, weight * block_weights.iter().max().unwrap());
}
}

Expand Down
Loading