Skip to content

Commit

Permalink
Reduce allocations random_unitaries
Browse files Browse the repository at this point in the history
The previous implementation had 4 heap allocations for each random
unitary constructed, this commit uses some fixed sized stack allocated
arrays and reduces that to two allocations one for q and r from the
factorization. We'll always need at least one for the `Array2` that gets
stored in each `UnitaryGate` as a numpy array. But to reduce to just
this we'll need a method of computing the QR factorization without an
allocation for the result space, nalgebtra might be a path for doing
that. While this currently isn't a bottleneck as the `UnitaryGate`
python object creation is the largest source of runtime, but assuming
that's fixed in the future this might have a larger impact.
  • Loading branch information
mtreinish committed Sep 30, 2024
1 parent 69d4cfa commit 32b4c1f
Showing 1 changed file with 43 additions and 11 deletions.
54 changes: 43 additions & 11 deletions crates/accelerate/src/circuit_library/quantum_volume.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use ndarray::prelude::*;
use num_complex::Complex64;
use numpy::IntoPyArray;
use rand::prelude::*;
use rand_distr::Normal;
use rand_distr::StandardNormal;
use rand_pcg::Pcg64Mcg;
use rayon::prelude::*;

Expand All @@ -33,29 +33,61 @@ use qiskit_circuit::packed_instruction::PackedOperation;
use qiskit_circuit::Qubit;
use smallvec::smallvec;

#[inline(always)]
fn random_complex(rng: &mut Pcg64Mcg) -> Complex64 {
Complex64::new(rng.sample(StandardNormal), rng.sample(StandardNormal))
* std::f64::consts::FRAC_1_SQRT_2
}

// This function's implementation was modeled off of the algorithm used in the
// `scipy.stats.unitary_group.rvs()` function defined here:
//
// https://github.com/scipy/scipy/blob/v1.14.1/scipy/stats/_multivariate.py#L4224-L4256
#[inline]
fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Array2<Complex64>> {
let mut rng = Pcg64Mcg::seed_from_u64(seed);
let dist = Normal::new(0., 1.0).unwrap();

(0..size).map(move |_| {
let mut z: Array2<Complex64> = Array2::from_shape_simple_fn((4, 4), || {
Complex64::new(dist.sample(&mut rng), dist.sample(&mut rng))
});
z.mapv_inplace(|x| x * std::f64::consts::FRAC_1_SQRT_2);
let qr = z.view().into_faer_complex().qr();
let r = qr.compute_r().as_ref().into_ndarray_complex().to_owned();
let mut d = r.into_diag();
d.mapv_inplace(|d| d / d.norm());
let raw_numbers: [[Complex64; 4]; 4] = [
[
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
],
[
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
],
[
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
],
[
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
random_complex(&mut rng),
],
];

let qr = aview2(&raw_numbers).into_faer_complex().qr();
let r = qr.compute_r();
let diag: [Complex64; 4] = [
r[(0, 0)].to_num_complex() / r[(0, 0)].abs(),
r[(1, 1)].to_num_complex() / r[(1, 1)].abs(),
r[(2, 2)].to_num_complex() / r[(2, 2)].abs(),
r[(3, 3)].to_num_complex() / r[(3, 3)].abs(),
];
let mut q = qr.compute_q().as_ref().into_ndarray_complex().to_owned();
q.axis_iter_mut(Axis(0)).for_each(|mut row| {
row.iter_mut()
.enumerate()
.for_each(|(index, val)| *val *= d.diag()[index])
.for_each(|(index, val)| *val *= diag[index])
});
q
})
Expand Down

0 comments on commit 32b4c1f

Please sign in to comment.