-
Notifications
You must be signed in to change notification settings - Fork 2.6k
Fully port ALAPScheduleAnalysis to Rust #14664
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.
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
Conversation
Thank you for opening a new pull request. Before your PR can be merged it will first need to pass continuous integration tests and be reviewed. Sometimes the review process can be slow, so please be patient. While you're waiting, please feel free to review other open PRs. While only a subset of people are authorized to approve pull requests for merging, everyone is encouraged to review open pull requests. Doing reviews helps reduce the burden on the core team and helps make the project's code better for everyone. One or more of the following people are relevant to this code:
|
Pull Request Test Coverage Report for Build 16786305064Warning: This coverage report may be inaccurate.This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.
Details
💛 - Coveralls |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks pretty good for a start, I have a couple comments.
clbit_write_latency: u64, | ||
node_durations: &Bound<PyDict>, | ||
) -> PyResult<Py<PyDict>> { | ||
if dag.qregs().len() != 1 || !dag.qregs().iter().any(|reg| reg.name() == "q") { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
To avoid iterating over all the registers in the dag, you can use dag.qregs_data()
to get an immutable reference to the inner RegisterData
within the dag. This may also allow you to use RegisterData::contains_key
which accepts a register name for querying.
let mut node_start_time: HashMap<NodeIndex, f64> = HashMap::new(); | ||
let mut idle_before: HashMap<Wire, f64> = HashMap::new(); | ||
|
||
for (i, _) in dag.qubits().objects().iter().enumerate() { | ||
idle_before.insert(Wire::Qubit(Qubit::new(i)), 0.0); | ||
} | ||
|
||
for (i, _) in dag.clbits().objects().iter().enumerate() { | ||
idle_before.insert(Wire::Clbit(Clbit::new(i)), 0.0); | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For this you may not have to iterate over ObjectRegistry::objects()
. Since the indices are always contiguous in cases like these, you may want to just do something of the sorts of:
for index in 0..dag.qubits().len() {
idle_before.insert(Wire::Qubit(Qubit::new(i)), 0.0);
}
for node_index in dag | ||
.topological_op_nodes()? | ||
.collect::<Vec<_>>() | ||
.into_iter() | ||
.rev() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is there a reason you need to collect this iterator before reversing it?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The reason is that for ALAP Scheduling, nodes are scheduled in reverse topological ordering (unlike ASAP) and are packed from the end of the circuit. I will move the comment from lines 84-92 above this so that it becomes more readable.
py: Python, | ||
dag: &DAGCircuit, | ||
clbit_write_latency: u64, | ||
node_durations: &Bound<PyDict>, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If node_durations
is a dictionary of DAGNode
and f64
. Maybe we should try to find a way of displaying this with a Rust native struct, or maybe just bringing InstructionDurations
to Rust might solve this.
|
||
let op = match dag.dag().node_weight(node_index) { | ||
Some(NodeType::Operation(op)) => op, | ||
_ => continue, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
More of a nitpick, but you could use panic!()
here. Though it is a bit of a drastic measure considering that panic will stop the process and cannot be caught in Python, there's no reason why topological_op_nodes()
should give you anything that isn't an instance DAGOpNode
.
OperationRef::Gate(_) => (true, false), | ||
OperationRef::StandardGate(_) => (true, false), | ||
_ => (false, op_name == "delay"), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
OperationRef::Gate(_) => (true, false), | |
OperationRef::StandardGate(_) => (true, false), | |
_ => (false, op_name == "delay"), | |
OperationRef::Gate(_) => (true, false), | |
OperationRef::StandardGate(_) => (true, false), | |
OperationRef::StandardInstruction(StandardInstruction::Delay(_)) => (false, true), | |
_ => (false, op_name == "delay"), |
You might want to add this condition for completion's sake.
// Get operation type | ||
let op_name = op.op.name(); | ||
let op_view = op.op.view(); | ||
let (is_gate, is_delay) = match op_view { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nitpick: This could be simplified to be one boolean value instead of two.
let node_start_time_final: HashMap<NodeIndex, f64> = node_start_time | ||
.into_iter() | ||
.map(|(n, t1)| (n, circuit_duration - t1)) | ||
.collect(); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You could avoid this step if you're going to convert to a PyDict
anyway. But if that's the case you should probably consider having a Rust native version of this function that returns a HashMap
, and then have the Python version perform the conversion to Python types.
test/benchmarks/passes.py
Outdated
class ALAPScheduleAnalysisPassBenchmarks: | ||
params = ([5, 14, 20], [1024]) | ||
|
||
param_names = ["n_qubits", "depth"] | ||
timeout = 300 | ||
|
||
def setup(self, n_qubits, depth): | ||
seed = 42 | ||
self.circuit = random_circuit( | ||
n_qubits, depth, measure=True, conditional=True, reset=True, seed=seed | ||
) | ||
self.dag = circuit_to_dag(self.circuit) | ||
gate_names = {inst[0].name for inst in self.circuit.data} | ||
self.durations = InstructionDurations([(name, None, 1000) for name in gate_names], dt=1e-7) | ||
|
||
def time_alap_schedule_analysis(self, _, __): | ||
ALAPScheduleAnalysis(durations=self.durations).run(self.dag) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think there are tests that perform something similar to this, check tests/benchmarks/scheduling_passes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh, I think I may have missed that. I have removed this class from the file.
This reverts commit e240bae.
…prove docstrings for clarity - The run method now constructs node_durations as a Python dict mapping node indices to durations, which is passed to Rust as a native HashMap<usize, f64>. - Enhanced docstrings for the ALAPScheduleAnalysis class and run method to clarify scheduling behavior
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is progressing quite well. I still think you could tweak it around to work a bit better, considering the rust side of things. Remember, we should be able to rust this pass without necessarily calling Python.
py: Python, | ||
dag: &DAGCircuit, | ||
clbit_write_latency: u64, | ||
node_durations: HashMap<usize, f64>, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmmm... I'm not quite sure using this type should be the best way of going around this. It should be best to use the defined NodeIndex
type. Though I get you won't be able to directly extract it from Python, but there should be a way around that too...
pub fn run_alap_schedule_analysis( | ||
py: Python, | ||
dag: &DAGCircuit, | ||
clbit_write_latency: u64, | ||
node_durations: HashMap<usize, f64>, | ||
) -> PyResult<Py<PyDict>> { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If we're trying to keep everything rust native, is there a way you could return a HashMap for this too? We want to be able to run this pass standalone (that meaning running the pass entirely from rust without it depending on Python for anything).
You might still be able to receive a PyDict of type { node_type: duration_type }
while keeping a path open for rust users to use the respective rust native types. Which would allow you to remove the python types altogether for a rust side without changing anything from what you've done on the Python side.
// Note that ALAP pass is inversely scheduled, thus | ||
// t0 is computed by subtracting t1 from the entire circuit duration. | ||
let py_dict = PyDict::new(py); | ||
for (node_idx, t1) in node_start_time { | ||
let node = dag.get_node(py, node_idx)?; | ||
let time = circuit_duration - t1; | ||
if time.fract() == 0.0 { | ||
py_dict.set_item(node, time as u64)?; | ||
} else { | ||
py_dict.set_item(node, time)?; | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This step shouldn't be in a Rust native function. Though it is still necessary to achieve what you want. There should be a way of not making this conversion to a PyDict required.
node_durations = { | ||
node._node_id: self._get_node_duration(node, dag) for node in dag.topological_op_nodes() | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm... I'm not quite sure you should you should be accessing node._node_id here, even though it is a python expose method. Perhaps you could keep the accessing the node_id in Rust instead of here. Afterall DAGNode is still a PyO3 struct and it stores an optional instance of NodeIndex
, the rust native indexing type for DAG.
for node_index in dag | ||
.topological_op_nodes()? | ||
.collect::<Vec<_>>() | ||
.into_iter() | ||
.rev() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why are we going back to the topological op_nodes of the dag here? I believe you already did this in Python and what we receive here is a dict of {DAGNode
: duration }, if that's the case, why are you doing this iteration again? I think you might benefit from just using the data you obtained from Python.
…ap_schedule_analysis, remove the extra node iteration.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is almost ready. However, there is still one big pain point that you had brought up before: Dealing with durations where the unit is either seconds or dt
. I left a suggestion on how to better handle it in one of my comments.
.extract::<DAGNode>()? | ||
.node | ||
.expect("Node index not found."); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nitpick: For the extraction here, an extra step you could add is to use downcast
to cast the node into a DAGOpNode
and then extract the DAGNode
just to make sure we are receiving DAGOpNodes
.
It could look somewhat like:
let node_idx = py_node
.downcast_into::<DAGOpNode>()?
.extract::<DAGNode>()?
.node
.expect("Node index not found.");
let op = match dag.dag().node_weight(node_index) { | ||
Some(NodeType::Operation(op)) => op, | ||
_ => panic!("topological_op_nodes() should only return instances of DagOpNode."), | ||
}; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think that now you should be able to access the DAG
via indexing to get the node's weight (NodeType
). And you also have the unwrap_operation()
to get the PackedInstruction.
let op = match dag.dag().node_weight(node_index) { | |
Some(NodeType::Operation(op)) => op, | |
_ => panic!("topological_op_nodes() should only return instances of DagOpNode."), | |
}; | |
let op = dag[node_index].unwrap_operation(); |
if t1.fract() == 0.0 { | ||
py_dict.set_item(node, t1 as u64)?; | ||
} else { | ||
py_dict.set_item(node, t1)?; | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Although this looks like a good way of handling things, it may not be the case. Since you'd want to preserve the integer cases (for dt), maybe you'd have to do something different. Perhaps finding a way of receiving either a float (seconds) or int (dt) from Python in Rust depending on whether the duration's unit. Since we only ever use either two for this pass and all durations will always be of the same type for each one circuit, you could create something that allows you to handle each unit. Perhaps with an enum.
Thank you, I have addressed the comments and created an enum |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A couple more comments about the new enumeration.
fn from_f64(value: f64, is_int_type: bool) -> TimeValue { | ||
if is_int_type { | ||
TimeValue::Int(value as u64) | ||
} else { | ||
TimeValue::Float(value) | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You should implement this using the From<T>
trait. With this you could automatically transform integers into TimeValue::Int
and f64
is correctly cast into TimeValue::Float
.
fn as_f64(&self) -> f64 { | ||
match self { | ||
TimeValue::Float(f) => *f, | ||
TimeValue::Int(i) => *i as f64, | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't believe this might be the right usage. You should maybe focus on handling any operation in a case to case basis, just to handle things more cleanly.
I comment on this mainly because casting from float
back to int
is a bit wonky.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think you're getting closer but there are some things that need ironing out.
I believe that, for any duration provided, the unit will end up being the same based on what the BaseScheduler
pass does. Based on whether the provided durations there are in terms of dt
or not (s
), the _get_node_duration()
method will always provide units of the same type. This should influence your approach when handling the two different types here.
} else if let Ok(val) = py_duration.extract::<f64>() { | ||
val.into() | ||
} else { | ||
return Err(TranspilerError::new_err("Duration must be u64 or f64")); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
return Err(TranspilerError::new_err("Duration must be u64 or f64")); | |
return Err(TranspilerError::new_err("Duration must be int or float")); |
Nitpick: Since this error message is only exposed to Python, we can use int
or float
instead of the rust types.
let mut node_start_time: HashMap<NodeIndex, f64> = HashMap::new(); | ||
let mut idle_before: HashMap<Wire, f64> = HashMap::new(); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What's the guarantee that these types are always float
and they don't match the same unit as the duration?
let mut result: HashMap<NodeIndex, TimeValue> = HashMap::new(); | ||
for (node_idx, t1) in node_start_time { | ||
let final_time = match node_durations.get(&node_idx) { | ||
Some(TimeValue::Int(_)) => TimeValue::from((circuit_duration - t1).round() as u64), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not sure we should be rounding things for the calculation. If we're dealing with TimeValue::Int
from the beginning, shouldn't t1 and circuit_duration
also be of TimeValue::Int
too??
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done in 8cc9893! All variables and HashMaps now use the TimeValue
enum, so the need of typecasting/ rounding values anywhere is eliminated.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks much better as this approach reduces the lines of code needed to achieve the same result. I had a couple of comments and suggestions.
pub trait TimeOps: Copy + PartialOrd + Add<Output = Self> + Sub<Output = Self> { | ||
fn zero() -> Self; | ||
fn max<'a>(a: &'a Self, b: &'a Self) -> &'a Self; | ||
fn from_u64(val: u64) -> Self; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why is the from_u64
needed here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I had included this function mainly for clbit_write_latency
. Since it is always provided as u64
type, It would need to be converted to T
to make the data types consistent for the addition/subtraction operations, but I realized the type conversion for it could be done at the function calling stage itself. So I have removed the from_u64
method completely.
let mut op_durations_int: Option<HashMap<NodeIndex, u64>> = None; | ||
let mut op_durations_float: Option<HashMap<NodeIndex, f64>> = None; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is there a reason why you're having these co-exist? I believe you can deal with each case individually within their respective if clauses, and do all of the python processing within those as well.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done in the latest commit!
…_durations only once.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the quick update
* Implement ALAP scheduling analysis in Rust and integrate with Python module * Fix lint issues * Add benchmark for ALAPScheduleAnalysis to passes.py * Add copyright and licensing information to alap_schedule_analysis.rs * Refactor ALAP schedule analysis based on suggestions * Remove the benchmark added to passes.py * Fix lint issues * Refactor comments in ALAP schedule analysis * Fix more lint issues * Simplify node duration handling and update function signature * Revert "Simplify node duration handling and update function signature" This reverts commit e240bae. * Update run function to pass node_durations as native Rust HashMap; improve docstrings for clarity - The run method now constructs node_durations as a Python dict mapping node indices to durations, which is passed to Rust as a native HashMap<usize, f64>. - Enhanced docstrings for the ALAPScheduleAnalysis class and run method to clarify scheduling behavior * Make documentation consistent * Fix lint * Split the pass into Rust-native function and Python-facing wrapper * Fix lint * Fix failing tests by move the type conversion logic back to py_run_alap_schedule_analysis, remove the extra node iteration. * Add enum for type handling of duration values * Rename enum to TimeValue * Implement From trait, ensure conversion to integer on case to case basis per node * Modify TimeValue to support arithmetic and max ops; update scheduling to use TimeValue throughout. * Replace enum with TimeOps trait for improved type flexibility * Remove from_u64 method from TimeOps trait; Refactor code to create op_durations only once. * Add back documentation example; Update node_duration argument in docstring
Summary
This PR migrates the implementation of the
ALAPScheduleAnalysis
pass from Python to Rust, maintaining identical scheduling logic and user-facing behavior.Details and comments
Fixes #12258