From 4f957172486328d628267ec63861014548fb8cbe Mon Sep 17 00:00:00 2001 From: vinitha-balachandran Date: Thu, 12 Mar 2026 10:09:16 +0800 Subject: [PATCH 1/5] MPI support for state and expectation calculations --- .../quimb_intro/benchmark_mpi_expectation.py | 127 +++++++++ src/qibotn/backends/quimb.py | 250 +++++++++++++++++- 2 files changed, 374 insertions(+), 3 deletions(-) create mode 100644 examples/quimb_intro/benchmark_mpi_expectation.py diff --git a/examples/quimb_intro/benchmark_mpi_expectation.py b/examples/quimb_intro/benchmark_mpi_expectation.py new file mode 100644 index 00000000..5bbcd62f --- /dev/null +++ b/examples/quimb_intro/benchmark_mpi_expectation.py @@ -0,0 +1,127 @@ +import time +import sys +import numpy as np +from qibo import Circuit, gates +from qibo.backends import construct_backend + +# Parse command line argument for number of processes expected +expected_procs = int(sys.argv[1]) if len(sys.argv) > 1 else 2 + +np.random.seed(42) + +def build_large_circuit(nqubits, nlayers): + """Build a larger circuit for benchmarking.""" + circ = Circuit(nqubits) + for _ in range(nlayers): + for q in range(nqubits): + circ.add(gates.RY(q=q, theta=np.random.random())) + circ.add(gates.RZ(q=q, theta=np.random.random())) + for q in range(nqubits): + circ.add(gates.CNOT(q % nqubits, (q + 1) % nqubits)) + return circ + +# Get MPI info +try: + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() +except ImportError: + print("ERROR: mpi4py not available") + exit(1) + +# Circuit parameters - larger for benchmarking +nqubits = 8 +nlayers = 3 + +if rank == 0: + print(f"=" * 60) + print(f"MPI EXPECTATION VALUE BENCHMARK") + print(f"=" * 60) + print(f"MPI processes: {size} (expected: {expected_procs})") + print(f"Circuit: {nqubits} qubits, {nlayers} layers") + print(f"State space: 2^{nqubits} = {2**nqubits} dimensions") + print(f"=" * 60) + +# Build circuit +circuit = build_large_circuit(nqubits, nlayers) + +# Define Hamiltonian with multiple terms +operators_list = ['z', 'x', 'y', 'zz', 'xx', 'yy', 'xyz'] +sites_list = [(0,), (1,), (2,), (3, 4), (5, 6), (1, 2), (0, 1, 2)] +coeffs_list = [1.0, 0.5, 0.3, 0.8, 0.6, 0.4, 0.2] + +if rank == 0: + print(f"\nHamiltonian: {len(operators_list)} terms") + for i, (ops, sites, coeff) in enumerate(zip(operators_list, sites_list, coeffs_list)): + print(f" Term {i+1}: {coeff} * {ops} on qubits {sites}") + +# Configure backend with MPI +backend = construct_backend(backend="qibotn", platform="quimb") +backend.configure_tn_simulation( + ansatz="mps", + max_bond_dimension=20, + MPI_enabled=True +) + +# Warm-up run +if rank == 0: + print("\nWarm-up run...") +circuit_warmup = build_large_circuit(4, 2) +operators_warmup = ['z', 'x'] +sites_warmup = [(0,), (1,)] +coeffs_warmup = [1.0, 1.0] +_ = backend.exp_value_observable_symbolic( + circuit_warmup, operators_warmup, sites_warmup, coeffs_warmup, 4 +) + +# Synchronize before timing +comm.Barrier() + +# Timed execution +if rank == 0: + print(f"\nStarting timed expectation value computation with {size} processes...") + +start_time = time.time() +exp_value = backend.exp_value_observable_symbolic( + circuit, operators_list, sites_list, coeffs_list, nqubits +) +end_time = time.time() + +execution_time = end_time - start_time + +# Gather timing from all ranks +all_times = comm.gather(execution_time, root=0) +all_values = comm.gather(exp_value, root=0) + +if rank == 0: + print(f"\n{'=' * 60}") + print(f"RESULTS") + print(f"{'=' * 60}") + print(f"\nTiming per rank:") + for i, t in enumerate(all_times): + print(f" Rank {i}: {t:.4f} seconds") + + avg_time = np.mean(all_times) + min_time = np.min(all_times) + max_time = np.max(all_times) + + print(f"\nTiming Statistics:") + print(f" Average: {avg_time:.4f} seconds") + print(f" Min: {min_time:.4f} seconds") + print(f" Max: {max_time:.4f} seconds") + print(f" Range: {max_time - min_time:.4f} seconds ({((max_time-min_time)/avg_time*100):.1f}%)") + + print(f"\nExpectation values per rank:") + for i, val in enumerate(all_values): + print(f" Rank {i}: {val:.10f}") + + print(f"\n{'=' * 60}") + print(f"Expectation Value: {exp_value:.10f}") + print(f"Computation Time: {execution_time:.4f} seconds") + print(f"{'=' * 60}") + print(f"✓ MPI expectation computation working with {size} processes") + print(f"{'=' * 60}") +else: + if rank == 0 or abs(exp_value) > 1e-10: + print(f"Rank {rank}: Completed in {execution_time:.4f} seconds") diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py index 3ee200d0..e5d1f1d8 100644 --- a/src/qibotn/backends/quimb.py +++ b/src/qibotn/backends/quimb.py @@ -49,6 +49,8 @@ def __init__(self, quimb_backend="numpy", contraction_optimizer="auto-hq"): self.max_bond_dimension = None self.svd_cutoff = None self.n_most_frequent_states = None + self.MPI_enabled = False + self.rank = None self.configure_tn_simulation() self.setup_backend_specifics( @@ -62,6 +64,7 @@ def configure_tn_simulation( max_bond_dimension: Optional[int] = None, svd_cutoff: Optional[float] = 1e-10, n_most_frequent_states: int = 100, + MPI_enabled: bool = False, ): """ Configure tensor network simulation. @@ -70,17 +73,25 @@ def configure_tn_simulation( ansatz : str, optional The tensor network ansatz to use. Default is `None` and, in this case, a generic Circuit Quimb class is used. - max_bond_dimension : int, optional + max_bond_dimension : int, optional The maximum bond dimension for the MPS ansatz. Default is 10. + svd_cutoff : float, optional + SVD cutoff value for MPS truncation. Default is 1e-10. + n_most_frequent_states : int, optional + Number of most frequent states to return. Default is 100. + MPI_enabled : bool, optional + Enable MPI-based multinode support. Default is False. Notes: - The ansatz determines the tensor network structure used for simulation. Currently, only "MPS" is supported. - The `max_bond_dimension` parameter controls the maximum allowed bond dimension for the MPS ansatz. + - MPI_enabled enables multinode support for large-scale simulations. """ self.ansatz = ansatz self.max_bond_dimension = max_bond_dimension self.svd_cutoff = svd_cutoff self.n_most_frequent_states = n_most_frequent_states + self.MPI_enabled = MPI_enabled @property @@ -158,8 +169,56 @@ def execute_circuit( - The ansatz determines the tensor network structure used for simulation. Currently, only "MPS" is supported. - If `initial_state` is provided, it must be compatible with the MPS ansatz. - The `nshots` parameter enables sampling from the circuit's output distribution. If not specified, the full statevector is computed. + - When MPI_enabled is True, multinode support is activated using dense_vector_tn_mpi_qu. """ - if initial_state is not None and self.ansatz == "MPS": + import numpy as np + + + if self.MPI_enabled: + if nshots is not None: + raise_error( + NotImplementedError, + "Sampling (nshots) is not supported with MPI-based execution." + ) + + + mps_opts = None + if self.ansatz == "mps": + mps_opts = { + "max_bond": self.max_bond_dimension, + "cutoff": self.svd_cutoff + } + + + state, self.rank = dense_vector_tn_mpi_qu( + qasm=circuit.to_qasm(), + nqubits=circuit.nqubits, + initial_state=initial_state, + mps_opts=mps_opts, + backend=self.backend + ) + + + if self.rank > 0: + state = np.array(0) + + + if return_array: + statevector = state.flatten() if self.rank == 0 else state + else: + statevector = state if self.rank == 0 else state + + return TensorNetworkResult( + nqubits=circuit.nqubits, + backend=self, + measures=None, + measured_probabilities=None, + prob_type=None, + statevector=statevector, + ) + + + if initial_state is not None and self.ansatz == "mps": initial_state = qtn.tensor_1d.MatrixProductState.from_dense( initial_state, 2 ) # 2 is the physical dimension @@ -230,7 +289,32 @@ def exp_value_observable_symbolic( float The real part of the expectation value of the Hamiltonian on the given circuit state. """ - # Validate that no term acts multiple times on the same qubit (no repeated indices in a sites tuple) + + if self.MPI_enabled: + + mps_opts = None + if self.ansatz == "mps": + mps_opts = { + "max_bond": self.max_bond_dimension, + "cutoff": self.svd_cutoff + } + + + expectation_value, self.rank = exp_value_observable_symbolic_mpi_qu( + qasm=circuit.to_qasm(), + nqubits=circuit.nqubits, + operators_list=operators_list, + sites_list=sites_list, + coeffs_list=coeffs_list, + mps_opts=mps_opts, + backend=self.backend, + contractions_optimizer=self.contractions_optimizer + ) + + return expectation_value + + # Standard (non-MPI) execution path + for sites in sites_list: if len(sites) != len(set(sites)): raise_error( @@ -385,3 +469,163 @@ def __getattr__(name): return BACKENDS[name] except KeyError: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + + +def dense_vector_tn_mpi_qu( + qasm: str, nqubits, initial_state, mps_opts, backend="numpy" +): + """Evaluate circuit in QASM format with Quimb using multi node multi cpu. + + Args: + qasm (str): QASM program. + nqubits (int): Number of qubits in the circuit + initial_state (list): Initial state in the dense vector form. If ``None`` the default ``|00...0>`` state is used. + mps_opts (dict): Parameters to tune the gate_opts for mps settings in ``class quimb.tensor.circuit.CircuitMPS``. + backend (str): Backend to perform the contraction with, e.g. ``numpy``, ``cupy``, ``jax``. Passed to ``opt_einsum``. + + Returns: + list: Amplitudes of final state after the simulation of the circuit. + """ + import numpy as np + import cotengra as ctg + from mpi4py import MPI + from mpi4py.futures import MPICommExecutor + from qibotn.eval_qu import init_state_tn + + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + target_size = int(2**nqubits / comm.size) + amplitudes = [] + + with MPICommExecutor() as pool: + + if pool is not None: + + if initial_state is not None: + initial_state = init_state_tn(nqubits, initial_state) + + circ_cls = qtn.circuit.CircuitMPS if mps_opts else qtn.circuit.Circuit + circ_quimb = circ_cls.from_openqasm2_str( + qasm, psi0=initial_state, gate_opts=mps_opts + ) + + # options to perform the slicing and finding contraction path usign Cotengra + + opt = ctg.ReusableHyperOptimizer( + parallel=pool, + # make sure we generate at least 1 slice per process + slicing_opts={"target_slices": comm.size}, + slicing_reconf_opts={"target_size": target_size}, + # uses basic greedy search algorithm to find optimal contraction path + methods=["greedy"], + # terminate search if contraction is cheap + max_time="rate:1e6", + # just uniformly sample the space + optlib="random", + # maximum number of trial contraction trees to generate + max_repeats=128, + # show the live progress of the best contraction found so far + progbar=False, + ) + + tensor_network = circ_quimb.psi + tree = tensor_network.contraction_tree(optimize=opt) + + arrays = [t.data for t in tensor_network] + + + fa = [ + pool.submit(tree.contract_slice, arrays, i) for i in range(tree.nslices) + ] + + amplitudes = [(c.result()).flatten() for c in fa] + + + return np.array(amplitudes), rank + + +def exp_value_observable_symbolic_mpi_qu( + qasm: str, nqubits, operators_list, sites_list, coeffs_list, mps_opts, + backend="numpy", contractions_optimizer="auto-hq" +): + """Evaluate expectation value of symbolic Hamiltonian with Quimb using multi node multi cpu. + + Args: + qasm (str): QASM program. + nqubits (int): Number of qubits in the circuit. + operators_list (list): List of operator strings representing the symbolic Hamiltonian terms. + sites_list (list): Tuples each specifying the qubits (sites) the corresponding operator acts on. + coeffs_list (list): The coefficients for each Hamiltonian term. + mps_opts (dict): Parameters to tune the gate_opts for mps settings in ``class quimb.tensor.circuit.CircuitMPS``. + backend (str): Backend to perform the contraction with, e.g. ``numpy``, ``cupy``, ``jax``. + contractions_optimizer (str): Contractions optimizer to use. + + Returns: + tuple: (expectation_value, rank) - The expectation value and MPI rank. + """ + import numpy as np + import cotengra as ctg + from mpi4py import MPI + from mpi4py.futures import MPICommExecutor + from qibo.config import raise_error + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + + for sites in sites_list: + if len(sites) != len(set(sites)): + raise_error( + ValueError, + f"Invalid Hamiltonian term sites {sites}: repeated qubit indices are not allowed " + "within a single term (e.g. (0,0,0) is invalid).", + ) + + expectation_value = 0.0 + + with MPICommExecutor() as pool: + + if pool is not None: + + circ_cls = qtn.circuit.CircuitMPS if mps_opts else qtn.circuit.Circuit + circ_quimb = circ_cls.from_openqasm2_str( + qasm, psi0=None, gate_opts=mps_opts + ) + + + for opstr, sites, coeff in zip(operators_list, sites_list, coeffs_list): + + op_str = opstr.lower() + ops = qu.pauli(op_str[0]) + for c in op_str[1:]: + ops = ops & qu.pauli(c) + + coeff = coeff.real + + target_size = int(2**nqubits / comm.size) + opt = ctg.ReusableHyperOptimizer( + parallel=pool, + slicing_opts={"target_slices": comm.size}, + slicing_reconf_opts={"target_size": target_size}, + methods=["greedy"], + max_time="rate:1e6", + optlib="random", + max_repeats=128, + progbar=False, + ) + + exp_val = circ_quimb.local_expectation( + ops, + where=sites, + backend=backend, + optimize=opt, + simplify_sequence="R", + ) + + expectation_value += coeff * exp_val + + if rank == 0: + return float(np.real(expectation_value)), rank + else: + return 0.0, rank From 7eb9fe207cef4c83ca9a166ce79ca2c8666ef4f8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Mar 2026 02:13:48 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../quimb_intro/benchmark_mpi_expectation.py | 32 ++--- src/qibotn/backends/quimb.py | 109 ++++++++---------- 2 files changed, 67 insertions(+), 74 deletions(-) diff --git a/examples/quimb_intro/benchmark_mpi_expectation.py b/examples/quimb_intro/benchmark_mpi_expectation.py index 5bbcd62f..ddbead10 100644 --- a/examples/quimb_intro/benchmark_mpi_expectation.py +++ b/examples/quimb_intro/benchmark_mpi_expectation.py @@ -1,5 +1,6 @@ -import time import sys +import time + import numpy as np from qibo import Circuit, gates from qibo.backends import construct_backend @@ -9,6 +10,7 @@ np.random.seed(42) + def build_large_circuit(nqubits, nlayers): """Build a larger circuit for benchmarking.""" circ = Circuit(nqubits) @@ -20,9 +22,11 @@ def build_large_circuit(nqubits, nlayers): circ.add(gates.CNOT(q % nqubits, (q + 1) % nqubits)) return circ + # Get MPI info try: from mpi4py import MPI + comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() @@ -47,28 +51,26 @@ def build_large_circuit(nqubits, nlayers): circuit = build_large_circuit(nqubits, nlayers) # Define Hamiltonian with multiple terms -operators_list = ['z', 'x', 'y', 'zz', 'xx', 'yy', 'xyz'] +operators_list = ["z", "x", "y", "zz", "xx", "yy", "xyz"] sites_list = [(0,), (1,), (2,), (3, 4), (5, 6), (1, 2), (0, 1, 2)] coeffs_list = [1.0, 0.5, 0.3, 0.8, 0.6, 0.4, 0.2] if rank == 0: print(f"\nHamiltonian: {len(operators_list)} terms") - for i, (ops, sites, coeff) in enumerate(zip(operators_list, sites_list, coeffs_list)): + for i, (ops, sites, coeff) in enumerate( + zip(operators_list, sites_list, coeffs_list) + ): print(f" Term {i+1}: {coeff} * {ops} on qubits {sites}") # Configure backend with MPI backend = construct_backend(backend="qibotn", platform="quimb") -backend.configure_tn_simulation( - ansatz="mps", - max_bond_dimension=20, - MPI_enabled=True -) +backend.configure_tn_simulation(ansatz="mps", max_bond_dimension=20, MPI_enabled=True) # Warm-up run if rank == 0: print("\nWarm-up run...") circuit_warmup = build_large_circuit(4, 2) -operators_warmup = ['z', 'x'] +operators_warmup = ["z", "x"] sites_warmup = [(0,), (1,)] coeffs_warmup = [1.0, 1.0] _ = backend.exp_value_observable_symbolic( @@ -101,21 +103,23 @@ def build_large_circuit(nqubits, nlayers): print(f"\nTiming per rank:") for i, t in enumerate(all_times): print(f" Rank {i}: {t:.4f} seconds") - + avg_time = np.mean(all_times) min_time = np.min(all_times) max_time = np.max(all_times) - + print(f"\nTiming Statistics:") print(f" Average: {avg_time:.4f} seconds") print(f" Min: {min_time:.4f} seconds") print(f" Max: {max_time:.4f} seconds") - print(f" Range: {max_time - min_time:.4f} seconds ({((max_time-min_time)/avg_time*100):.1f}%)") - + print( + f" Range: {max_time - min_time:.4f} seconds ({((max_time-min_time)/avg_time*100):.1f}%)" + ) + print(f"\nExpectation values per rank:") for i, val in enumerate(all_values): print(f" Rank {i}: {val:.10f}") - + print(f"\n{'=' * 60}") print(f"Expectation Value: {exp_value:.10f}") print(f"Computation Time: {execution_time:.4f} seconds") diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py index e5d1f1d8..bec85d73 100644 --- a/src/qibotn/backends/quimb.py +++ b/src/qibotn/backends/quimb.py @@ -172,42 +172,34 @@ def execute_circuit( - When MPI_enabled is True, multinode support is activated using dense_vector_tn_mpi_qu. """ import numpy as np - - + if self.MPI_enabled: if nshots is not None: raise_error( NotImplementedError, - "Sampling (nshots) is not supported with MPI-based execution." + "Sampling (nshots) is not supported with MPI-based execution.", ) - - + mps_opts = None if self.ansatz == "mps": - mps_opts = { - "max_bond": self.max_bond_dimension, - "cutoff": self.svd_cutoff - } - - + mps_opts = {"max_bond": self.max_bond_dimension, "cutoff": self.svd_cutoff} + state, self.rank = dense_vector_tn_mpi_qu( qasm=circuit.to_qasm(), nqubits=circuit.nqubits, initial_state=initial_state, mps_opts=mps_opts, - backend=self.backend + backend=self.backend, ) - - + if self.rank > 0: state = np.array(0) - - + if return_array: statevector = state.flatten() if self.rank == 0 else state else: statevector = state if self.rank == 0 else state - + return TensorNetworkResult( nqubits=circuit.nqubits, backend=self, @@ -216,8 +208,7 @@ def execute_circuit( prob_type=None, statevector=statevector, ) - - + if initial_state is not None and self.ansatz == "mps": initial_state = qtn.tensor_1d.MatrixProductState.from_dense( initial_state, 2 @@ -289,17 +280,13 @@ def exp_value_observable_symbolic( float The real part of the expectation value of the Hamiltonian on the given circuit state. """ - + if self.MPI_enabled: - + mps_opts = None if self.ansatz == "mps": - mps_opts = { - "max_bond": self.max_bond_dimension, - "cutoff": self.svd_cutoff - } - - + mps_opts = {"max_bond": self.max_bond_dimension, "cutoff": self.svd_cutoff} + expectation_value, self.rank = exp_value_observable_symbolic_mpi_qu( qasm=circuit.to_qasm(), nqubits=circuit.nqubits, @@ -308,13 +295,13 @@ def exp_value_observable_symbolic( coeffs_list=coeffs_list, mps_opts=mps_opts, backend=self.backend, - contractions_optimizer=self.contractions_optimizer + contractions_optimizer=self.contractions_optimizer, ) - + return expectation_value - + # Standard (non-MPI) execution path - + for sites in sites_list: if len(sites) != len(set(sites)): raise_error( @@ -469,7 +456,7 @@ def __getattr__(name): return BACKENDS[name] except KeyError: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None - + def dense_vector_tn_mpi_qu( qasm: str, nqubits, initial_state, mps_opts, backend="numpy" @@ -486,20 +473,20 @@ def dense_vector_tn_mpi_qu( Returns: list: Amplitudes of final state after the simulation of the circuit. """ - import numpy as np import cotengra as ctg + import numpy as np from mpi4py import MPI from mpi4py.futures import MPICommExecutor + from qibotn.eval_qu import init_state_tn - - + comm = MPI.COMM_WORLD rank = comm.Get_rank() target_size = int(2**nqubits / comm.size) amplitudes = [] with MPICommExecutor() as pool: - + if pool is not None: if initial_state is not None: @@ -528,26 +515,30 @@ def dense_vector_tn_mpi_qu( # show the live progress of the best contraction found so far progbar=False, ) - + tensor_network = circ_quimb.psi tree = tensor_network.contraction_tree(optimize=opt) - + arrays = [t.data for t in tensor_network] - fa = [ pool.submit(tree.contract_slice, arrays, i) for i in range(tree.nslices) ] - + amplitudes = [(c.result()).flatten() for c in fa] - - + return np.array(amplitudes), rank def exp_value_observable_symbolic_mpi_qu( - qasm: str, nqubits, operators_list, sites_list, coeffs_list, mps_opts, - backend="numpy", contractions_optimizer="auto-hq" + qasm: str, + nqubits, + operators_list, + sites_list, + coeffs_list, + mps_opts, + backend="numpy", + contractions_optimizer="auto-hq", ): """Evaluate expectation value of symbolic Hamiltonian with Quimb using multi node multi cpu. @@ -564,16 +555,15 @@ def exp_value_observable_symbolic_mpi_qu( Returns: tuple: (expectation_value, rank) - The expectation value and MPI rank. """ - import numpy as np import cotengra as ctg + import numpy as np from mpi4py import MPI from mpi4py.futures import MPICommExecutor from qibo.config import raise_error - + comm = MPI.COMM_WORLD rank = comm.Get_rank() - - + for sites in sites_list: if len(sites) != len(set(sites)): raise_error( @@ -581,28 +571,27 @@ def exp_value_observable_symbolic_mpi_qu( f"Invalid Hamiltonian term sites {sites}: repeated qubit indices are not allowed " "within a single term (e.g. (0,0,0) is invalid).", ) - + expectation_value = 0.0 - + with MPICommExecutor() as pool: - + if pool is not None: - + circ_cls = qtn.circuit.CircuitMPS if mps_opts else qtn.circuit.Circuit circ_quimb = circ_cls.from_openqasm2_str( qasm, psi0=None, gate_opts=mps_opts ) - - + for opstr, sites, coeff in zip(operators_list, sites_list, coeffs_list): - + op_str = opstr.lower() ops = qu.pauli(op_str[0]) for c in op_str[1:]: ops = ops & qu.pauli(c) - + coeff = coeff.real - + target_size = int(2**nqubits / comm.size) opt = ctg.ReusableHyperOptimizer( parallel=pool, @@ -614,7 +603,7 @@ def exp_value_observable_symbolic_mpi_qu( max_repeats=128, progbar=False, ) - + exp_val = circ_quimb.local_expectation( ops, where=sites, @@ -622,9 +611,9 @@ def exp_value_observable_symbolic_mpi_qu( optimize=opt, simplify_sequence="R", ) - + expectation_value += coeff * exp_val - + if rank == 0: return float(np.real(expectation_value)), rank else: From 96abc3a3b2df33730e6d3c860f6c3d04af02e447 Mon Sep 17 00:00:00 2001 From: Tankya2 Date: Fri, 10 Apr 2026 17:41:50 +0800 Subject: [PATCH 3/5] Add quimb gate support for benchmark circuits --- src/qibotn/backends/quimb.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py index bec85d73..5a0621d2 100644 --- a/src/qibotn/backends/quimb.py +++ b/src/qibotn/backends/quimb.py @@ -4,7 +4,6 @@ import quimb as qu import quimb.tensor as qtn from qibo.config import raise_error -from qibo.gates.abstract import ParametrizedGate from qibo.models import Circuit from qibotn.backends.abstract import QibotnBackend @@ -25,6 +24,8 @@ "cnot": "CNOT", "cy": "CY", "cz": "CZ", + "cu1": "CU1", + "rzz": "RZZ", "iswap": "ISWAP", "swap": "SWAP", "ccx": "CCX", @@ -368,19 +369,12 @@ def _qibo_circuit_to_quimb( params = getattr(gate, "parameters", ()) qubits = getattr(gate, "qubits", ()) - is_parametrized = isinstance(gate, ParametrizedGate) and getattr( - gate, "trainable", True + # Quimb's apply_gate does not accept the 'parametrized' kwarg in this path. + circ.apply_gate( + quimb_gate_name, + *params, + *qubits, ) - if is_parametrized: - circ.apply_gate( - quimb_gate_name, *params, *qubits, parametrized=is_parametrized - ) - else: - circ.apply_gate( - quimb_gate_name, - *params, - *qubits, - ) return circ From adef92578ef9a68a079a44e8ebf8f5f0d49f322c Mon Sep 17 00:00:00 2001 From: vinitha-balachandran Date: Tue, 14 Apr 2026 12:23:57 +0800 Subject: [PATCH 4/5] Refactored the code and added mpi pytest for new functions --- src/qibotn/backends/quimb.py | 224 ++++++++++++++++---------------- tests/conftest.py | 12 ++ tests/test_quimb_mpi_backend.py | 104 +++++++++++++++ 3 files changed, 230 insertions(+), 110 deletions(-) create mode 100644 tests/test_quimb_mpi_backend.py diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py index e5d1f1d8..d37e5c77 100644 --- a/src/qibotn/backends/quimb.py +++ b/src/qibotn/backends/quimb.py @@ -172,42 +172,39 @@ def execute_circuit( - When MPI_enabled is True, multinode support is activated using dense_vector_tn_mpi_qu. """ import numpy as np - - + if self.MPI_enabled: if nshots is not None: raise_error( NotImplementedError, - "Sampling (nshots) is not supported with MPI-based execution." + "Sampling (nshots) is not supported with MPI-based execution.", ) - - + mps_opts = None if self.ansatz == "mps": - mps_opts = { - "max_bond": self.max_bond_dimension, - "cutoff": self.svd_cutoff - } - - + mps_opts = {"max_bond": self.max_bond_dimension, "cutoff": self.svd_cutoff} + state, self.rank = dense_vector_tn_mpi_qu( qasm=circuit.to_qasm(), nqubits=circuit.nqubits, initial_state=initial_state, mps_opts=mps_opts, - backend=self.backend + backend=self.backend, ) - - + if self.rank > 0: state = np.array(0) - - + if return_array: statevector = state.flatten() if self.rank == 0 else state else: statevector = state if self.rank == 0 else state - + + if self.rank == 0: + statevector = state.flatten() if return_array else state + else: + statevector = None + return TensorNetworkResult( nqubits=circuit.nqubits, backend=self, @@ -216,8 +213,7 @@ def execute_circuit( prob_type=None, statevector=statevector, ) - - + if initial_state is not None and self.ansatz == "mps": initial_state = qtn.tensor_1d.MatrixProductState.from_dense( initial_state, 2 @@ -289,17 +285,13 @@ def exp_value_observable_symbolic( float The real part of the expectation value of the Hamiltonian on the given circuit state. """ - + if self.MPI_enabled: - + mps_opts = None if self.ansatz == "mps": - mps_opts = { - "max_bond": self.max_bond_dimension, - "cutoff": self.svd_cutoff - } - - + mps_opts = {"max_bond": self.max_bond_dimension, "cutoff": self.svd_cutoff} + expectation_value, self.rank = exp_value_observable_symbolic_mpi_qu( qasm=circuit.to_qasm(), nqubits=circuit.nqubits, @@ -308,13 +300,13 @@ def exp_value_observable_symbolic( coeffs_list=coeffs_list, mps_opts=mps_opts, backend=self.backend, - contractions_optimizer=self.contractions_optimizer + contraction_optimizer=self.contractions_optimizer, ) - + return expectation_value - + # Standard (non-MPI) execution path - + for sites in sites_list: if len(sites) != len(set(sites)): raise_error( @@ -381,23 +373,12 @@ def _qibo_circuit_to_quimb( params = getattr(gate, "parameters", ()) qubits = getattr(gate, "qubits", ()) - is_parametrized = isinstance(gate, ParametrizedGate) and getattr( - gate, "trainable", True - ) - if is_parametrized: - circ.apply_gate( - quimb_gate_name, *params, *qubits, parametrized=is_parametrized - ) - else: - circ.apply_gate( - quimb_gate_name, - *params, - *qubits, - ) + # Quimb's apply_gate does not accept the 'parametrized' kwarg in this path. + circ.apply_gate(quimb_gate_name, *params, *qubits) return circ -def _string_to_quimb_operator(self, op_str): +def _string_to_quimb_operator(op_str): """ Convert a Pauli string (e.g. 'xzy') to a Quimb operator using '&' chaining. @@ -427,7 +408,7 @@ def _string_to_quimb_operator(self, op_str): "execute_circuit": execute_circuit, "exp_value_observable_symbolic": exp_value_observable_symbolic, "_qibo_circuit_to_quimb": _qibo_circuit_to_quimb, - "_string_to_quimb_operator": _string_to_quimb_operator, + "_string_to_quimb_operator": staticmethod(_string_to_quimb_operator), "circuit_ansatz": circuit_ansatz, } @@ -469,10 +450,10 @@ def __getattr__(name): return BACKENDS[name] except KeyError: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None - + def dense_vector_tn_mpi_qu( - qasm: str, nqubits, initial_state, mps_opts, backend="numpy" + qasm: str, nqubits, initial_state, mps_opts, backend="numpy", path_opts=None ): """Evaluate circuit in QASM format with Quimb using multi node multi cpu. @@ -482,24 +463,26 @@ def dense_vector_tn_mpi_qu( initial_state (list): Initial state in the dense vector form. If ``None`` the default ``|00...0>`` state is used. mps_opts (dict): Parameters to tune the gate_opts for mps settings in ``class quimb.tensor.circuit.CircuitMPS``. backend (str): Backend to perform the contraction with, e.g. ``numpy``, ``cupy``, ``jax``. Passed to ``opt_einsum``. + path_opts (dict or object, optional): Contraction path options passed to Quimb/Cotengra. + If ``None``, a default ``ReusableHyperOptimizer`` is used. Returns: list: Amplitudes of final state after the simulation of the circuit. """ - import numpy as np import cotengra as ctg + import numpy as np from mpi4py import MPI from mpi4py.futures import MPICommExecutor + from qibotn.eval_qu import init_state_tn - - + comm = MPI.COMM_WORLD rank = comm.Get_rank() target_size = int(2**nqubits / comm.size) amplitudes = [] with MPICommExecutor() as pool: - + if pool is not None: if initial_state is not None: @@ -512,42 +495,50 @@ def dense_vector_tn_mpi_qu( # options to perform the slicing and finding contraction path usign Cotengra - opt = ctg.ReusableHyperOptimizer( - parallel=pool, - # make sure we generate at least 1 slice per process - slicing_opts={"target_slices": comm.size}, - slicing_reconf_opts={"target_size": target_size}, - # uses basic greedy search algorithm to find optimal contraction path - methods=["greedy"], - # terminate search if contraction is cheap - max_time="rate:1e6", - # just uniformly sample the space - optlib="random", - # maximum number of trial contraction trees to generate - max_repeats=128, - # show the live progress of the best contraction found so far - progbar=False, - ) - + if path_opts is None or isinstance(path_opts, dict): + path_kwargs = { + # make sure we generate at least 1 slice per process + "slicing_opts": {"target_slices": comm.size}, + "slicing_reconf_opts": {"target_size": target_size}, + # uses basic greedy search algorithm to find optimal contraction path + "methods": ["greedy"], + # terminate search if contraction is cheap + "max_time": "rate:1e6", + # just uniformly sample the space + "optlib": "random", + # maximum number of trial contraction trees to generate + "max_repeats": 128, + # show the live progress of the best contraction found so far + "progbar": False, + } + if isinstance(path_opts, dict): + path_kwargs.update(path_opts) + path_opts = ctg.ReusableHyperOptimizer(parallel=pool, **path_kwargs) + tensor_network = circ_quimb.psi - tree = tensor_network.contraction_tree(optimize=opt) - + tree = tensor_network.contraction_tree(optimize=path_opts) + arrays = [t.data for t in tensor_network] - fa = [ pool.submit(tree.contract_slice, arrays, i) for i in range(tree.nslices) ] - + amplitudes = [(c.result()).flatten() for c in fa] - - + return np.array(amplitudes), rank def exp_value_observable_symbolic_mpi_qu( - qasm: str, nqubits, operators_list, sites_list, coeffs_list, mps_opts, - backend="numpy", contractions_optimizer="auto-hq" + qasm: str, + nqubits, + operators_list, + sites_list, + coeffs_list, + mps_opts, + backend="numpy", + contraction_optimizer="auto-hq", + path_opts=None, ): """Evaluate expectation value of symbolic Hamiltonian with Quimb using multi node multi cpu. @@ -559,21 +550,24 @@ def exp_value_observable_symbolic_mpi_qu( coeffs_list (list): The coefficients for each Hamiltonian term. mps_opts (dict): Parameters to tune the gate_opts for mps settings in ``class quimb.tensor.circuit.CircuitMPS``. backend (str): Backend to perform the contraction with, e.g. ``numpy``, ``cupy``, ``jax``. - contractions_optimizer (str): Contractions optimizer to use. + contraction_optimizer (str, optional): The contractions_optimizer to use for the Quimb/Cotengra tensor network simulation. + If ``None``, defaults to "auto-hq". + path_opts (dict or object, optional): Contraction path options passed to Quimb/Cotengra. + If ``None``, defaults to a ``ReusableHyperOptimizer`` unless + ``contractions_optimizer`` is set to a non-default value. Returns: tuple: (expectation_value, rank) - The expectation value and MPI rank. """ - import numpy as np import cotengra as ctg + import numpy as np from mpi4py import MPI from mpi4py.futures import MPICommExecutor from qibo.config import raise_error - + comm = MPI.COMM_WORLD rank = comm.Get_rank() - - + for sites in sites_list: if len(sites) != len(set(sites)): raise_error( @@ -581,50 +575,60 @@ def exp_value_observable_symbolic_mpi_qu( f"Invalid Hamiltonian term sites {sites}: repeated qubit indices are not allowed " "within a single term (e.g. (0,0,0) is invalid).", ) - + expectation_value = 0.0 - + with MPICommExecutor() as pool: - + if pool is not None: - + circ_cls = qtn.circuit.CircuitMPS if mps_opts else qtn.circuit.Circuit circ_quimb = circ_cls.from_openqasm2_str( qasm, psi0=None, gate_opts=mps_opts ) - - + + target_size = int(2**nqubits / comm.size) + if path_opts is None: + if contraction_optimizer not in (None, "auto-hq"): + path_opts = contraction_optimizer + else: + path_opts = ctg.ReusableHyperOptimizer( + parallel=pool, + slicing_opts={"target_slices": comm.size}, + slicing_reconf_opts={"target_size": target_size}, + methods=["greedy"], + max_time="rate:1e6", + optlib="random", + max_repeats=128, + progbar=False, + ) + elif isinstance(path_opts, dict): + path_kwargs = { + "slicing_opts": {"target_slices": comm.size}, + "slicing_reconf_opts": {"target_size": target_size}, + "methods": ["greedy"], + "max_time": "rate:1e6", + "optlib": "random", + "max_repeats": 128, + "progbar": False, + } + path_kwargs.update(path_opts) + path_opts = ctg.ReusableHyperOptimizer(parallel=pool, **path_kwargs) for opstr, sites, coeff in zip(operators_list, sites_list, coeffs_list): - - op_str = opstr.lower() - ops = qu.pauli(op_str[0]) - for c in op_str[1:]: - ops = ops & qu.pauli(c) - + + ops = _string_to_quimb_operator(opstr) coeff = coeff.real - - target_size = int(2**nqubits / comm.size) - opt = ctg.ReusableHyperOptimizer( - parallel=pool, - slicing_opts={"target_slices": comm.size}, - slicing_reconf_opts={"target_size": target_size}, - methods=["greedy"], - max_time="rate:1e6", - optlib="random", - max_repeats=128, - progbar=False, - ) - + exp_val = circ_quimb.local_expectation( ops, where=sites, backend=backend, - optimize=opt, + optimize=path_opts, simplify_sequence="R", ) - + expectation_value += coeff * exp_val - + if rank == 0: return float(np.real(expectation_value)), rank else: diff --git a/tests/conftest.py b/tests/conftest.py index c5e9ed4b..6bf9e802 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ Pytest fixtures. """ +import os import sys import pytest @@ -57,6 +58,17 @@ def pytest_runtest_setup(item): def pytest_configure(config): config.addinivalue_line("markers", "linux: mark test to run only on linux") + if os.getenv("OMPI_COMM_WORLD_SIZE"): + if hasattr(config.option, "no_cov"): + config.option.no_cov = True + cov_plugin = config.pluginmanager.get_plugin("_cov") + if cov_plugin is not None: + config.pluginmanager.unregister(cov_plugin) + + +def pytest_addoption(parser): + # Keep pyproject's [tool.pytest.ini_options].env valid even when pytest-env is not installed. + parser.addini("env", type="linelist", help="Environment variables for tests.") def pytest_generate_tests(metafunc): diff --git a/tests/test_quimb_mpi_backend.py b/tests/test_quimb_mpi_backend.py new file mode 100644 index 00000000..cfd06d35 --- /dev/null +++ b/tests/test_quimb_mpi_backend.py @@ -0,0 +1,104 @@ +# mpirun -np 2 python -m pytest tests/test_quimb_mpi_backend.py -m mpi + +import math + +import numpy as np +import pytest +import qibo +from qibo import construct_backend, hamiltonians +from qibo.models import QFT +from qibo.symbols import X, Z + +pytest.importorskip("mpi4py") + +ABS_TOL = 1e-7 + + +def qibo_qft(nqubits, swaps): + circ_qibo = QFT(nqubits, swaps) + state_vec = circ_qibo().state(numpy=True) + return circ_qibo, state_vec + + +def build_observable(nqubits): + """Helper function to construct a target observable.""" + hamiltonian_form = 0 + for i in range(nqubits): + hamiltonian_form += 0.5 * X(i % nqubits) * Z((i + 1) % nqubits) + + hamiltonian = hamiltonians.SymbolicHamiltonian(form=hamiltonian_form) + return hamiltonian + + +def build_symbolic_lists(nqubits): + """Build operators/sites/coeffs accepted by exp_value_observable_symbolic.""" + operators_list = [] + sites_list = [] + coeffs_list = [] + + for i in range(nqubits): + operators_list.append("xz") + sites_list.append((i % nqubits, (i + 1) % nqubits)) + coeffs_list.append(0.5) + + return operators_list, sites_list, coeffs_list + + + +@pytest.mark.parametrize("nqubits", [2, 5, 7]) +def test_quimb_statevector_mpi(nqubits: int): + qibo.set_backend(backend="numpy") + qibo_circ, expected_sv = qibo_qft(nqubits, swaps=True) + + backend = construct_backend(backend="qibotn", platform="quimb") + backend.configure_tn_simulation( + ansatz="mps", + max_bond_dimension=None, + svd_cutoff=1e-12, + MPI_enabled=True, + ) + + outcome = backend.execute_circuit(circuit=qibo_circ, return_array=True) + + if backend.rank == 0: + got_sv = outcome.state().flatten() + assert np.allclose( + expected_sv, got_sv, atol=1e-7, rtol=1e-7 + ), "Resulting dense vectors do not match" + else: + assert outcome.state() is None + + + +@pytest.mark.parametrize("nqubits", [2, 5, 7]) +def test_quimb_expectation_mpi(nqubits: int): + qibo.set_backend(backend="numpy") + qibo_circ, _ = qibo_qft(nqubits, swaps=True) + + ham = build_observable(nqubits) + exact_expval = ham.expectation(qibo_circ) + + operators_list, sites_list, coeffs_list = build_symbolic_lists(nqubits) + + backend = construct_backend(backend="qibotn", platform="quimb") + backend.configure_tn_simulation( + ansatz="mps", + max_bond_dimension=None, + svd_cutoff=1e-12, + MPI_enabled=True, + ) + + result_tn = backend.exp_value_observable_symbolic( + qibo_circ, + operators_list, + sites_list, + coeffs_list, + nqubits, + ) + + if backend.rank == 0: + assert math.isclose( + float(exact_expval), float(result_tn), abs_tol=ABS_TOL + ), f"Rank {backend.rank}: mismatch, expected {exact_expval}, got {result_tn}" + else: + assert result_tn == 0.0 From 2f55e20b9ab36e1415158cc6ac6ea5e3e856a5fe Mon Sep 17 00:00:00 2001 From: Tankya2 Date: Tue, 28 Apr 2026 16:15:45 +0800 Subject: [PATCH 5/5] Fix dense vector TN MPI slice gathering --- src/qibotn/backends/quimb.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py index a84b3875..a81ad5a9 100644 --- a/src/qibotn/backends/quimb.py +++ b/src/qibotn/backends/quimb.py @@ -453,6 +453,17 @@ def __getattr__(name): raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None +def _gather_dense_slice_futures(tree, futures, backend): + """Assemble contracted slice futures into the final dense output. + + Cotengra's ``contract_slice`` returns one contribution per slice, which may + need summing and/or stacking depending on whether sliced indices overlap the + output. ``gather_slices`` performs that reconstruction for us. + """ + + return tree.gather_slices((future.result() for future in futures), backend=backend) + + def dense_vector_tn_mpi_qu( qasm: str, nqubits, initial_state, mps_opts, backend="numpy", path_opts=None ): @@ -480,7 +491,7 @@ def dense_vector_tn_mpi_qu( comm = MPI.COMM_WORLD rank = comm.Get_rank() target_size = int(2**nqubits / comm.size) - amplitudes = [] + amplitudes = np.array([]) with MPICommExecutor() as pool: @@ -525,9 +536,9 @@ def dense_vector_tn_mpi_qu( pool.submit(tree.contract_slice, arrays, i) for i in range(tree.nslices) ] - amplitudes = [(c.result()).flatten() for c in fa] + amplitudes = _gather_dense_slice_futures(tree, fa, backend=backend) - return np.array(amplitudes), rank + return amplitudes, rank def exp_value_observable_symbolic_mpi_qu(