Move the circuit library's entanglement logic to Rust (#12950)

* Move ``get_entangler_map`` to Rust

* add reno

* add docs, better arg order

* clippy

* add docs to make reviewers happier

* use Itertools::combinations

thanks @alexanderivrii for the hint!

* avoid usage of convert_idx

* implement reviews

* implement review comments

* rm unused import

crates/accelerate/src/circuit_library/entanglement.rs

@@ -0,0 +1,245 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2024
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use itertools::Itertools;
+use pyo3::prelude::*;
+use pyo3::{
+    types::{PyAnyMethods, PyInt, PyList, PyListMethods, PyString, PyTuple},
+    Bound, PyAny, PyResult,
+};
+
+use crate::QiskitError;
+
+/// Get all-to-all entanglement. For 4 qubits and block size 2 we have:
+/// [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
+fn full(num_qubits: u32, block_size: u32) -> impl Iterator<Item = Vec<u32>> {
+    (0..num_qubits).combinations(block_size as usize)
+}
+
+/// Get a linear entanglement structure. For ``n`` qubits and block size ``m`` we have:
+/// [(0..m-1), (1..m), (2..m+1), ..., (n-m..n-1)]
+fn linear(num_qubits: u32, block_size: u32) -> impl DoubleEndedIterator<Item = Vec<u32>> {
+    (0..num_qubits - block_size + 1)
+        .map(move |start_index| (start_index..start_index + block_size).collect())
+}
+
+/// Get a reversed linear entanglement. This is like linear entanglement but in reversed order:
+/// [(n-m..n-1), ..., (1..m), (0..m-1)]
+/// This is particularly interesting, as CX+"full" uses n(n-1)/2 gates, but operationally equals
+/// CX+"reverse_linear", which needs only n-1 gates.
+fn reverse_linear(num_qubits: u32, block_size: u32) -> impl Iterator<Item = Vec<u32>> {
+    linear(num_qubits, block_size).rev()
+}
+
+/// Return the qubit indices for circular entanglement. This is defined as tuples of length ``m``
+/// starting at each possible index ``(0..n)``. Historically, Qiskit starts with index ``n-m+1``.
+/// This is probably easiest understood for a concerete example of 4 qubits and block size 3:
+/// [(2,3,0), (3,0,1), (0,1,2), (1,2,3)]
+fn circular(num_qubits: u32, block_size: u32) -> Box<dyn Iterator<Item = Vec<u32>>> {
+    if block_size == 1 || num_qubits == block_size {
+        Box::new(linear(num_qubits, block_size))
+    } else {
+        let historic_offset = num_qubits - block_size + 1;
+        Box::new((0..num_qubits).map(move |start_index| {
+            (0..block_size)
+                .map(|i| (historic_offset + start_index + i) % num_qubits)
+                .collect()
+        }))
+    }
+}
+
+/// Get pairwise entanglement. This is typically used on 2 qubits and only has a depth of 2, as
+/// first all odd pairs, and then even pairs are entangled. For example on 6 qubits:
+/// [(0, 1), (2, 3), (4, 5), /* now the even pairs */ (1, 2), (3, 4)]
+fn pairwise(num_qubits: u32) -> impl Iterator<Item = Vec<u32>> {
+    // for Python-folks (like me): pairwise is equal to linear[::2] + linear[1::2]
+    linear(num_qubits, 2)
+        .step_by(2)
+        .chain(linear(num_qubits, 2).skip(1).step_by(2))
+}
+
+/// The shifted, circular, alternating (sca) entanglement is motivated from circuits 14/15 of
+/// https://arxiv.org/abs/1905.10876. It corresponds to circular entanglement, with the difference
+/// that in each layer (controlled by ``offset``) the entanglement gates are shifted by one, plus
+/// in each second layer, the entanglement gate is turned upside down.
+/// Offset 0 -> [(2,3,0), (3,0,1), (0,1,2), (1,2,3)]
+/// Offset 1 -> [(3,2,1), (0,3,2), (1,0,3), (2,1,0)]
+/// Offset 2 -> [(0,1,2), (1,2,3), (2,3,0), (3,0,1)]
+/// ...
+fn shift_circular_alternating(
+    num_qubits: u32,
+    block_size: u32,
+    offset: usize,
+) -> Box<dyn Iterator<Item = Vec<u32>>> {
+    // index at which we split the circular iterator -- we use Python-like indexing here,
+    // and define ``split`` as equivalent to a Python index of ``-offset``
+    let split = (num_qubits - (offset as u32 % num_qubits)) % num_qubits;
+    let shifted = circular(num_qubits, block_size)
+        .skip(split as usize)
+        .chain(circular(num_qubits, block_size).take(split as usize));
+    if offset % 2 == 0 {
+        Box::new(shifted)
+    } else {
+        // if the offset is odd, reverse the indices inside the qubit block (e.g. turn CX
+        // gates upside down)
+        Box::new(shifted.map(|indices| indices.into_iter().rev().collect()))
+    }
+}
+
+/// Get an entangler map for an arbitrary number of qubits.
+///
+/// Args:
+///     num_qubits: The number of qubits of the circuit.
+///     block_size: The number of qubits of the entangling block.
+///     entanglement: The entanglement strategy as string.
+///     offset: The block offset, can be used if the entanglements differ per block,
+///         for example used in the "sca" mode.
+///
+/// Returns:
+///     The entangler map using mode ``entanglement`` to scatter a block of ``block_size``
+///     qubits on ``num_qubits`` qubits.
+pub fn get_entanglement_from_str(
+    num_qubits: u32,
+    block_size: u32,
+    entanglement: &str,
+    offset: usize,
+) -> PyResult<Box<dyn Iterator<Item = Vec<u32>>>> {
+    if num_qubits == 0 || block_size == 0 {
+        return Ok(Box::new(std::iter::empty()));
+    }
+
+    if block_size > num_qubits {
+        return Err(QiskitError::new_err(format!(
+            "block_size ({}) cannot be larger than number of qubits ({})",
+            block_size, num_qubits
+        )));
+    }
+
+    match (entanglement, block_size) {
+        ("full", _) => Ok(Box::new(full(num_qubits, block_size))),
+        ("linear", _) => Ok(Box::new(linear(num_qubits, block_size))),
+        ("reverse_linear", _) => Ok(Box::new(reverse_linear(num_qubits, block_size))),
+        ("sca", _) => Ok(shift_circular_alternating(num_qubits, block_size, offset)),
+        ("circular", _) => Ok(circular(num_qubits, block_size)),
+        ("pairwise", 1) => Ok(Box::new(linear(num_qubits, 1))),
+        ("pairwise", 2) => Ok(Box::new(pairwise(num_qubits))),
+        ("pairwise", _) => Err(QiskitError::new_err(format!(
+            "block_size ({}) can be at most 2 for pairwise entanglement",
+            block_size
+        ))),
+        _ => Err(QiskitError::new_err(format!(
+            "Unsupported entanglement: {}",
+            entanglement
+        ))),
+    }
+}
+
+/// Get an entangler map for an arbitrary number of qubits.
+///
+/// Args:
+///     num_qubits: The number of qubits of the circuit.
+///     block_size: The number of qubits of the entangling block.
+///     entanglement: The entanglement strategy.
+///     offset: The block offset, can be used if the entanglements differ per block,
+///         for example used in the "sca" mode.
+///
+/// Returns:
+///     The entangler map using mode ``entanglement`` to scatter a block of ``block_size``
+///     qubits on ``num_qubits`` qubits.
+pub fn get_entanglement<'a>(
+    num_qubits: u32,
+    block_size: u32,
+    entanglement: &'a Bound<PyAny>,
+    offset: usize,
+) -> PyResult<Box<dyn Iterator<Item = PyResult<Vec<u32>>> + 'a>> {
+    // unwrap the callable, if it is one
+    let entanglement = if entanglement.is_callable() {
+        entanglement.call1((offset,))?
+    } else {
+        entanglement.to_owned()
+    };
+
+    if let Ok(strategy) = entanglement.downcast::<PyString>() {
+        let as_str = strategy.to_string();
+        return Ok(Box::new(
+            get_entanglement_from_str(num_qubits, block_size, as_str.as_str(), offset)?.map(Ok),
+        ));
+    } else if let Ok(list) = entanglement.downcast::<PyList>() {
+        let entanglement_iter = list.iter().map(move |el| {
+            let connections = el
+                .downcast::<PyTuple>()?
+                // .expect("Entanglement must be list of tuples") // clearer error message than `?`
+                .iter()
+                .map(|index| index.downcast::<PyInt>()?.extract())
+                .collect::<Result<Vec<u32>, _>>()?;
+
+            if connections.len() != block_size as usize {
+                return Err(QiskitError::new_err(format!(
+                    "Entanglement {:?} does not match block size {}",
+                    connections, block_size
+                )));
+            }
+            Ok(connections)
+        });
+        return Ok(Box::new(entanglement_iter));
+    }
+    Err(QiskitError::new_err(
+        "Entanglement must be a string or list of qubit indices.",
+    ))
+}
+
+/// Get the entanglement for given number of qubits and block size.
+///
+/// Args:
+///     num_qubits: The number of qubits to entangle.
+///     block_size: The entanglement block size (e.g. 2 for CX or 3 for CCX).
+///     entanglement: The entanglement strategy. This can be one of:
+///
+///         * string: Available options are ``"full"``, ``"linear"``, ``"pairwise"``
+///             ``"reverse_linear"``, ``"circular"``, or ``"sca"``.
+///         * list of tuples: A list of entanglements given as tuple, e.g. [(0, 1), (1, 2)].
+///         * callable: A callable that takes as input an offset as ``int`` (usually the layer
+///             in the variational circuit) and returns a string or list of tuples to use as
+///             entanglement in this layer.
+///
+///     offset: An offset used by certain entanglement strategies (e.g. ``"sca"``) or if the
+///         entanglement is given as callable. This is typically used to have different
+///         entanglement structures in different layers of variational quantum circuits.
+///
+/// Returns:
+///     The entanglement as list of tuples.
+///
+/// Raises:
+///     QiskitError: In case the entanglement is invalid.
+#[pyfunction]
+#[pyo3(signature = (num_qubits, block_size, entanglement, offset=0))]
+pub fn get_entangler_map<'py>(
+    py: Python<'py>,
+    num_qubits: u32,
+    block_size: u32,
+    entanglement: &Bound<PyAny>,
+    offset: usize,
+) -> PyResult<Vec<Bound<'py, PyTuple>>> {
+    // The entanglement is Result<impl Iterator<Item = Result<Vec<u32>>>>, so there's two
+    // levels of errors we must handle: the outer error is handled by the outer match statement,
+    // and the inner (Result<Vec<u32>>) is handled upon the PyTuple creation.
+    match get_entanglement(num_qubits, block_size, entanglement, offset) {
+        Ok(entanglement) => entanglement
+            .into_iter()
+            .map(|vec| match vec {
+                Ok(vec) => Ok(PyTuple::new_bound(py, vec)),
+                Err(e) => Err(e),
+            })
+            .collect::<Result<Vec<_>, _>>(),
+        Err(e) => Err(e),
+    }
+}

crates/accelerate/src/circuit_library/mod.rs

@@ -0,0 +1,21 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2024
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use pyo3::prelude::*;
+
+mod entanglement;
+
+#[pymodule]
+pub fn circuit_library(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(entanglement::get_entangler_map))?;
+    Ok(())
+}

crates/accelerate/src/lib.rs

@@ -14,6 +14,7 @@ use std::env;
 
 use pyo3::import_exception;
 
+pub mod circuit_library;
 pub mod convert_2q_block_matrix;
 pub mod dense_layout;
 pub mod edge_collections;

crates/pyext/src/lib.rs

@@ -13,14 +13,14 @@
 use pyo3::prelude::*;
 
 use qiskit_accelerate::{
-    convert_2q_block_matrix::convert_2q_block_matrix, dense_layout::dense_layout,
-    error_map::error_map, euler_one_qubit_decomposer::euler_one_qubit_decomposer,
-    isometry::isometry, nlayout::nlayout, optimize_1q_gates::optimize_1q_gates,
-    pauli_exp_val::pauli_expval, results::results, sabre::sabre, sampled_exp_val::sampled_exp_val,
-    sparse_pauli_op::sparse_pauli_op, star_prerouting::star_prerouting,
-    stochastic_swap::stochastic_swap, synthesis::synthesis, target_transpiler::target,
-    two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate, utils::utils,
-    vf2_layout::vf2_layout,
+    circuit_library::circuit_library, convert_2q_block_matrix::convert_2q_block_matrix,
+    dense_layout::dense_layout, error_map::error_map,
+    euler_one_qubit_decomposer::euler_one_qubit_decomposer, isometry::isometry, nlayout::nlayout,
+    optimize_1q_gates::optimize_1q_gates, pauli_exp_val::pauli_expval, results::results,
+    sabre::sabre, sampled_exp_val::sampled_exp_val, sparse_pauli_op::sparse_pauli_op,
+    star_prerouting::star_prerouting, stochastic_swap::stochastic_swap, synthesis::synthesis,
+    target_transpiler::target, two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate,
+    utils::utils, vf2_layout::vf2_layout,
 };
 
 #[inline(always)]
@@ -39,6 +39,7 @@ fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
     add_submodule(m, qiskit_circuit::circuit, "circuit")?;
     add_submodule(m, qiskit_qasm2::qasm2, "qasm2")?;
     add_submodule(m, qiskit_qasm3::qasm3, "qasm3")?;
+    add_submodule(m, circuit_library, "circuit_library")?;
     add_submodule(m, convert_2q_block_matrix, "convert_2q_block_matrix")?;
     add_submodule(m, dense_layout, "dense_layout")?;
     add_submodule(m, error_map, "error_map")?;

qiskit/__init__.py

@@ -60,6 +60,7 @@
 # We manually define them on import so people can directly import qiskit._accelerate.* submodules
 # and not have to rely on attribute access.  No action needed for top-level extension packages.
 sys.modules["qiskit._accelerate.circuit"] = _accelerate.circuit
+sys.modules["qiskit._accelerate.circuit_library"] = _accelerate.circuit_library
 sys.modules["qiskit._accelerate.convert_2q_block_matrix"] = _accelerate.convert_2q_block_matrix
 sys.modules["qiskit._accelerate.dense_layout"] = _accelerate.dense_layout
 sys.modules["qiskit._accelerate.error_map"] = _accelerate.error_map

qiskit/circuit/library/n_local/n_local.py

@@ -31,9 +31,11 @@
 )
 from qiskit.exceptions import QiskitError
 from qiskit.circuit.library.standard_gates import get_standard_gate_name_mapping
+from qiskit._accelerate.circuit_library import get_entangler_map as fast_entangler_map
 
 from ..blueprintcircuit import BlueprintCircuit
 
+
 if typing.TYPE_CHECKING:
     import qiskit  # pylint: disable=cyclic-import
 
@@ -1037,51 +1039,11 @@ def get_entangler_map(
     Raises:
         ValueError: If the entanglement mode ist not supported.
     """
-    n, m = num_circuit_qubits, num_block_qubits
-    if m > n:
-        raise ValueError(
-            "The number of block qubits must be smaller or equal to the number of "
-            "qubits in the circuit."
-        )
-
-    if entanglement == "pairwise" and num_block_qubits > 2:
-        raise ValueError("Pairwise entanglement is not defined for blocks with more than 2 qubits.")
-
-    if entanglement == "full":
-        return list(itertools.combinations(list(range(n)), m))
-    elif entanglement == "reverse_linear":
-        # reverse linear connectivity. In the case of m=2 and the entanglement_block='cx'
-        # then it's equivalent to 'full' entanglement
-        reverse = [tuple(range(n - i - m, n - i)) for i in range(n - m + 1)]
-        return reverse
-    elif entanglement in ["linear", "circular", "sca", "pairwise"]:
-        linear = [tuple(range(i, i + m)) for i in range(n - m + 1)]
-        # if the number of block qubits is 1, we don't have to add the 'circular' part
-        if entanglement == "linear" or m == 1:
-            return linear
-
-        if entanglement == "pairwise":
-            return linear[::2] + linear[1::2]
-
-        # circular equals linear plus top-bottom entanglement (if there's space for it)
-        if n > m:
-            circular = [tuple(range(n - m + 1, n)) + (0,)] + linear
-        else:
-            circular = linear
-        if entanglement == "circular":
-            return circular
-
-        # sca is circular plus shift and reverse
-        shifted = circular[-offset:] + circular[:-offset]
-        if offset % 2 == 1:  # if odd, reverse the qubit indices
-            sca = [ind[::-1] for ind in shifted]
-        else:
-            sca = shifted
-
-        return sca
-
-    else:
-        raise ValueError(f"Unsupported entanglement type: {entanglement}")
+    try:
+        return fast_entangler_map(num_circuit_qubits, num_block_qubits, entanglement, offset)
+    except Exception as exc:
+        # need this as Rust is now raising a QiskitError, where this function was raising ValueError
+        raise ValueError("Something went wrong in Rust space, here's the error:") from exc
 
 
 _StdlibGateResult = collections.namedtuple("_StdlibGateResult", ("gate", "num_params"))