From f48af0116c775a440a8dc46c50f86d4f730b02df Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Fri, 29 Mar 2024 15:04:36 -0400 Subject: [PATCH 1/2] Expose internal rust interface to DenseLayout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit makes a small change to the rust code for DenseLayout that enables calling it more easily from rust. The primary obstacle was the pyfunction used PyReadonlyArray2 inputs which precludes calling it with rust constructed Array2Views. This adds a new inner public function which takes the array view directly and then the pyfunction's only job is to convert the inputs and outputs to Python. The python side of the function is still building a sparse matrix and then runs reverse Cuthill–McKee to get a permutation of the densest subgraph so any rust consumers will want to keep that in mind (and maybe use sprs to do the same). At the same time it corrects an oversight in the original implementation where the returned numpy arrays of the densest subgraph are copied instead of being returned as references. This should slightly improve performance by eliminating 3 array copies that weren't needed. --- crates/accelerate/src/dense_layout.rs | 35 +++++++++++++++++++++------ 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/crates/accelerate/src/dense_layout.rs b/crates/accelerate/src/dense_layout.rs index 4570eafc0bf1..e3e63f33c053 100644 --- a/crates/accelerate/src/dense_layout.rs +++ b/crates/accelerate/src/dense_layout.rs @@ -16,7 +16,7 @@ use hashbrown::{HashMap, HashSet}; use indexmap::IndexSet; use ndarray::prelude::*; use numpy::PyReadonlyArray2; -use numpy::ToPyArray; +use numpy::IntoPyArray; use rayon::prelude::*; use pyo3::prelude::*; @@ -110,8 +110,33 @@ pub fn best_subset( error_matrix: PyReadonlyArray2, ) -> PyResult<(PyObject, PyObject, PyObject)> { let coupling_adj_mat = coupling_adjacency.as_array(); - let coupling_shape = coupling_adj_mat.shape(); let err = error_matrix.as_array(); + let (rows, cols, best_map) = best_subset_inner( + num_qubits, + coupling_adj_mat, + num_meas, + num_cx, + use_error, + symmetric_coupling_map, + err, + )?; + Ok(( + rows.into_pyarray(py).into(), + cols.into_pyarray(py).into(), + best_map.into_pyarray(py).into(), + )) +} + +pub fn best_subset_inner( + num_qubits: usize, + coupling_adj_mat: ArrayView2, + num_meas: usize, + num_cx: usize, + use_error: bool, + symmetric_coupling_map: bool, + err: ArrayView2, +) -> PyResult<(Vec, Vec, Vec)> { + let coupling_shape = coupling_adj_mat.shape(); let avg_meas_err = err.diag().mean().unwrap(); let map_fn = |k| -> SubsetResult { @@ -216,11 +241,7 @@ pub fn best_subset( let rows: Vec = new_cmap.iter().map(|edge| edge[0]).collect(); let cols: Vec = new_cmap.iter().map(|edge| edge[1]).collect(); - Ok(( - rows.to_pyarray(py).into(), - cols.to_pyarray(py).into(), - best_map.to_pyarray(py).into(), - )) + Ok((rows, cols, best_map)) } #[pymodule] From 33493a0e8858a7ccad3952caaba071fcf02e3333 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Wed, 22 May 2024 07:30:58 -0400 Subject: [PATCH 2/2] Remove PyResult --- crates/accelerate/src/dense_layout.rs | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/crates/accelerate/src/dense_layout.rs b/crates/accelerate/src/dense_layout.rs index 4805def78291..901a906d9c81 100644 --- a/crates/accelerate/src/dense_layout.rs +++ b/crates/accelerate/src/dense_layout.rs @@ -15,8 +15,8 @@ use ahash::RandomState; use hashbrown::{HashMap, HashSet}; use indexmap::IndexSet; use ndarray::prelude::*; -use numpy::PyReadonlyArray2; use numpy::IntoPyArray; +use numpy::PyReadonlyArray2; use rayon::prelude::*; use pyo3::prelude::*; @@ -108,10 +108,10 @@ pub fn best_subset( use_error: bool, symmetric_coupling_map: bool, error_matrix: PyReadonlyArray2, -) -> PyResult<(PyObject, PyObject, PyObject)> { +) -> (PyObject, PyObject, PyObject) { let coupling_adj_mat = coupling_adjacency.as_array(); let err = error_matrix.as_array(); - let (rows, cols, best_map) = best_subset_inner( + let [rows, cols, best_map] = best_subset_inner( num_qubits, coupling_adj_mat, num_meas, @@ -119,12 +119,12 @@ pub fn best_subset( use_error, symmetric_coupling_map, err, - )?; - Ok(( - rows.to_pyarray_bound(py).into(), - cols.to_pyarray_bound(py).into(), - best_map.to_pyarray_bound(py).into(), - )) + ); + ( + rows.into_pyarray_bound(py).into(), + cols.into_pyarray_bound(py).into(), + best_map.into_pyarray_bound(py).into(), + ) } pub fn best_subset_inner( @@ -135,7 +135,7 @@ pub fn best_subset_inner( use_error: bool, symmetric_coupling_map: bool, err: ArrayView2, -) -> PyResult<(Vec, Vec, Vec)> { +) -> [Vec; 3] { let coupling_shape = coupling_adj_mat.shape(); let avg_meas_err = err.diag().mean().unwrap(); @@ -241,7 +241,7 @@ pub fn best_subset_inner( let rows: Vec = new_cmap.iter().map(|edge| edge[0]).collect(); let cols: Vec = new_cmap.iter().map(|edge| edge[1]).collect(); - Ok((rows, cols, best_map)) + [rows, cols, best_map] } #[pymodule]