-
Notifications
You must be signed in to change notification settings - Fork 2.4k
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
binary matrices utils in rust #12456
Conversation
One or more of the following people are relevant to this code:
|
Pull Request Test Coverage Report for Build 9697320760Warning: 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 |
I'm excited for this PR! Can you add a graph highlighting the runtime speedup of your Rust code compared to the old Python code? I don't think this graph would need to be comprehensive but it would be nice to get a clearer picture for some relevant procedures. :-) |
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 doing this, it's a great start. I left a couple inline comments for some suggestions. The larger comment shows how you can leverage rayon for multithreading part of the implementation. But we should benchmark it to make sure that multithreading actually speeds it up in this case. I could see the benefits from it not being huge here. In that suggestion you can just remove the .into_par_iter()
line and that will make the iterator chain serially execute instead of doing it in parallel.
Shelly and I were experimenting with this code trying to understand how to modify matrix elements in-place (rather than copying the matrix from Python space to Rust space and back), and this is what we currently have. The wrapper function looks as follows: #[pyfunction]
#[pyo3(signature = (mat, ncols=None, full_elim=false))]
fn _gauss_elimination(
_py: Python,
mat: &PyArray2<bool>,
ncols: Option<usize>,
full_elim: Option<bool>,
) {
unsafe {
let view = mat.as_array_mut();
gauss_elimination(view, ncols, full_elim);
}
} This seems to work, producing the expected results, yet there are a few points we are not sure about.
however replacing
but this does not:
giving the error
which is probably to be expected. |
If you want to make the rust function mutate the input numpy array this isn't too difficult. You need to change the types on things and that might require input normalization in the python side. I did a quick test locally and you would need to change the PR like: diff --git a/crates/accelerate/src/linear_matrix.rs b/crates/accelerate/src/linear_matrix.rs
index d708e71bd..197c13c3e 100644
--- a/crates/accelerate/src/linear_matrix.rs
+++ b/crates/accelerate/src/linear_matrix.rs
@@ -10,19 +10,19 @@
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.
-use ndarray::{s, Array2};
-use numpy::{AllowTypeChange, IntoPyArray, PyArray2, PyArrayLike2};
+use ndarray::{s, ArrayViewMut2};
+use numpy::PyReadwriteArray2;
use pyo3::prelude::*;
// Perform ROW operation on a matrix mat
-fn _row_op(mat: &mut Array2<i8>, ctrl: usize, trgt: usize) {
+fn _row_op(mat: &mut ArrayViewMut2<i8>, ctrl: usize, trgt: usize) {
let row0 = mat.row(ctrl).to_owned();
let mut row1 = mat.row_mut(trgt);
row1.zip_mut_with(&row0, |x, &y| *x ^= y);
}
// Perform COL operation on a matrix mat
-fn _col_op(mat: &mut Array2<i8>, ctrl: usize, trgt: usize) {
+fn _col_op(mat: &mut ArrayViewMut2<i8>, ctrl: usize, trgt: usize) {
let col0 = mat.column(ctrl).to_owned();
let mut col1 = mat.column_mut(trgt);
col1.zip_mut_with(&col0, |x, &y| *x ^= y);
@@ -31,11 +31,7 @@ fn _col_op(mat: &mut Array2<i8>, ctrl: usize, trgt: usize) {
// Gauss elimination of a matrix mat with m rows and n columns.
// If full_elim = True, it allows full elimination of mat[:, 0 : ncols]
// Returns the matrix mat.
-fn gauss_elimination(
- mut mat: Array2<i8>,
- ncols: Option<usize>,
- full_elim: Option<bool>,
-) -> Array2<i8> {
+fn gauss_elimination(mut mat: ArrayViewMut2<i8>, ncols: Option<usize>, full_elim: Option<bool>) {
let (m, mut n) = (mat.nrows(), mat.ncols()); // no. of rows and columns
if let Some(ncols_val) = ncols {
n = usize::min(n, ncols_val); // no. of active columns
@@ -62,7 +58,7 @@ fn gauss_elimination(
}
}
if !is_non_zero {
- return mat; // A is in the canonical form
+ return; // A is in the canonical form
}
if new_r != r {
@@ -87,21 +83,17 @@ fn gauss_elimination(
}
r += 1;
}
- mat
}
#[pyfunction]
#[pyo3(signature = (mat, ncols=None, full_elim=false))]
fn _gauss_elimination(
- py: Python,
- mat: PyArrayLike2<i8, AllowTypeChange>,
+ mut mat: PyReadwriteArray2<i8>,
ncols: Option<usize>,
full_elim: Option<bool>,
-) -> PyResult<Py<PyArray2<i8>>> {
- let view = mat.as_array().to_owned();
- Ok(gauss_elimination(view, ncols, full_elim)
- .into_pyarray_bound(py)
- .unbind())
+) {
+ let view = mat.as_array_mut();
+ gauss_elimination(view, ncols, full_elim);
}
#[pymodule] The tradeoff here is that you'll probably need to do a |
Thanks, @mtreinish, your code using Regarding |
It will make a copy if the dtype doesn't match the input array's type: https://numpy.org/doc/stable/reference/generated/numpy.asarray.html. The copy is needed to convert the elements to the specified dtype and would be unavoidable in this case. The previous |
A very quick experiment on my laptop:
Old version (python/numpy):
New code (ref. 0cf7502) but without using parallel iterators:
New code (ref. 0cf7502) with using parallel iterators:
|
The Rust code is compiled in release mode (without it it's significantly slower). |
Thank you! In terms of the preferred API, would it be better to be more general on the Rust side (i.e. restore |
I typically prefer doing the type normalization in Python directly. I think in this case that's preferable because you avoid a copy if the input is the correct dtype. While in the previous implementation you always needed to copy to get a mutable array. |
if full_elim == Some(true) { | ||
mat.axis_iter_mut(Axis(0)) | ||
.into_par_iter() | ||
.enumerate() | ||
.filter(|(i, row)| (*i < r) && row[new_k]) | ||
.for_each(|(_i, mut row)| { | ||
row.zip_mut_with(&row0, |x, &y| *x ^= y); | ||
}); | ||
} | ||
|
||
mat.axis_iter_mut(Axis(0)) | ||
.into_par_iter() | ||
.enumerate() | ||
.filter(|(i, row)| (*i > r) && (*i < m) && row[new_k]) | ||
.for_each(|(_i, mut row)| { | ||
row.zip_mut_with(&row0, |x, &y| *x ^= y); | ||
}); |
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 can be simplified to a single iterator with something like (double check that I got the filter condition correct:
if full_elim == Some(true) { | |
mat.axis_iter_mut(Axis(0)) | |
.into_par_iter() | |
.enumerate() | |
.filter(|(i, row)| (*i < r) && row[new_k]) | |
.for_each(|(_i, mut row)| { | |
row.zip_mut_with(&row0, |x, &y| *x ^= y); | |
}); | |
} | |
mat.axis_iter_mut(Axis(0)) | |
.into_par_iter() | |
.enumerate() | |
.filter(|(i, row)| (*i > r) && (*i < m) && row[new_k]) | |
.for_each(|(_i, mut row)| { | |
row.zip_mut_with(&row0, |x, &y| *x ^= y); | |
}); | |
mat.axis_iter_mut(Axis(0)) | |
.into_par_iter() | |
.enumerate() | |
.filter(|(i, row)| (full_elim == Some(true) && (*i < r) && row[new_k]) || (*i > r && *i < m)) | |
.for_each(|(_i, mut row)| { | |
row.zip_mut_with(&row0, |x, &y| *x ^= y); | |
}); |
You can also move the full_elim
condition into the for_each
block and just change the filter to be i < m
if that's easier.
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 other thing I was thinking about is you can probably get a slice view of the rows you care about and iterate over that slice instead of the full matrix. It's might be worth investigating that instead of doing a single big iterator like this. TBH, I'm not sure what is going to be the most efficient here we need to benchmark all the different possible ways to write this and figure out what ends up being the best performing.
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 only took a quick look at the numpy usage in this PR. I feel like there are a lot of extra copies happening which we can mostly avoid I think. Although I haven't reviewed the rust code in depth I was leaving that for @kevinhartman (as I think he volunteered to review it). As we're accelerating this code extra copies will have a real noticeable overhead so we should be intentional around when we copy.
@@ -53,7 +53,8 @@ def random_cnotdihedral(num_qubits, seed=None): | |||
random_invertible_binary_matrix, | |||
) | |||
|
|||
linear = random_invertible_binary_matrix(num_qubits, seed=rng) | |||
seed = rng.integers(100000, size=1, dtype=np.uint64)[0] | |||
linear = random_invertible_binary_matrix(num_qubits, seed=seed).astype(int) |
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 will always result in a copy. You could use .astype(int, copy=False)
but it still would need to copy. Also I think the dtype is incorrect shouldn't this be int8 because that's what cnot dihedral is using: https://github.com/qiskit/qiskit/blob/main/qiskit/quantum_info/operators/dihedral/dihedral.py#L129.
That being said wouldn't it just make more sense to change all the binary matrices in rust to use Array2<i8>
to match this? It'd be the same size as Array2<bool>
but then you could do a no-copy return to cnotdihedral and anywhere else using these functions.
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 original motivation for rust binary matrices to be of type Array2<bool>
is that python classes Clifford
and LinearFunction
internally store a matrix of booleans, and we wanted to avoid copying there. IMHO, the dihedral class is less important then either of the above (at least right now). However, I do agree that the python synthesize size needs to be cleaned up.
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 agree with @alexanderivrii . As follows form the function name, the aim is only to handle binary matrices (e.g. bool).
In the future, the plan is to port the synthesis code to rust, and so we can handle the other classes later.
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.
Overall this looks good, and I'm eager to merge it! Below are some questions and comments, my main concern being on whether we handle the functions as API-documented (meaning we keep the signature the same) or whether we can change them as if they were private 🙂
if verify { | ||
let mat2 = binary_matmul_inner(mat, (&invmat).into())?; | ||
let identity_matrix: Array2<bool> = Array2::from_shape_fn((n, n), |(i, j)| i == j); | ||
if mat2.ne(&identity_matrix) { | ||
return Err("The inverse matrix is not correct.".to_string()); | ||
} | ||
} |
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 understand that this was here before, but it seems very strange to essentially have a runtime test of our code here... if we provide the function then we should also be sure it is correct and not need to verify the result 😅 I'm fine keeping it because it existed before, but I would prefer removing this.
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.
There is a parameter "verify" in several parts of the synthesis code, which is mainly used for tests, since the algorithms are quite complicated (the default value is False)
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.
We can leave it for this PR 👍🏻 but a common subroutine like matrix inversion seems like something stable enough that we don't need to bake in a verification and instead just have some tests in the test suite (fwiw other packages like NumPy/SciPy also don't have any arguments to check that 🙂).
releasenotes/notes/linear-binary-matrix-utils-rust-c48b5577749c34ab.yaml
Outdated
Show resolved
Hide resolved
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.
LGTM thanks for the effort! 🦀
* gaussian elimination in rust * handle lint errors * replace python function by rust function for gauss elimination * change matrix elements type from bool to i8 * add parallelization in row operations * update matrices in place * change matrix type in rust code to bool * handle type in python code * update filter following review * remove parallelization using rayon * move _gauss_elimination_with_perm to rust * fix fmt error * simplify _gauss_elimination function * update _compute_rank_after_gauss_elim to rust * update _row_op and _col_op * transfer _row_op and _col_op from python to rust * fix code due to failing tests * minor update of types * move calc_inverse_matrix to rust, add _binary_matmul in rust * fix failing tests, by changing mat type from int to bool * update rust docstrings * add function _add_row_or_col to rust code * improve binary_matmul * proper error handling * unified format of function names * move compute_rank from python to rust, update errors * update type of mat in compute_rank * move random_invertible_binary_matrix and check_invertible_binary_matrix to rust * Updating HighLevelSynthesis tests that depend on the specific random number * Updating LinearSynthesis tests to pass seeds * Updating tests in test_linear_function * changing the matrix type in random dyhedral to be a matrix of ints rather than bools * updating cx_cz synthesis tests * updating clifford tests * remove unused imports * add option seed=None * enhance rust docs * add release notes * remove unnecessary copy in python * another copy fix * another copy fix * update rust docstrings * update release notes --------- Co-authored-by: AlexanderIvrii <alexi@il.ibm.com>
Summary
close #12330
Transfer to rust binary matrices utils from:
https://github.com/Qiskit/qiskit/blob/main/qiskit/synthesis/linear/linear_matrix_utils.py.
joint with @alexanderivrii
Details and comments