diff --git a/examples/quimb_intro/benchmark_mpi_expectation.py b/examples/quimb_intro/benchmark_mpi_expectation.py new file mode 100644 index 00000000..ddbead10 --- /dev/null +++ b/examples/quimb_intro/benchmark_mpi_expectation.py @@ -0,0 +1,131 @@ +import sys +import time + +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..a81ad5a9 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", @@ -49,6 +50,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 +65,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 +74,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 +170,52 @@ 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 + + if self.rank == 0: + statevector = state.flatten() if return_array else state + else: + statevector = None + + 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 +286,28 @@ 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, + 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( @@ -297,23 +374,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. @@ -343,7 +409,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, } @@ -385,3 +451,197 @@ def __getattr__(name): return BACKENDS[name] except KeyError: 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 +): + """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``. + 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 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 = np.array([]) + + 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 + + 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=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 = _gather_dense_slice_futures(tree, fa, backend=backend) + + return amplitudes, rank + + +def exp_value_observable_symbolic_mpi_qu( + 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. + + 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``. + 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 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( + 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 + ) + + 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): + + ops = _string_to_quimb_operator(opstr) + coeff = coeff.real + + exp_val = circ_quimb.local_expectation( + ops, + where=sites, + backend=backend, + optimize=path_opts, + simplify_sequence="R", + ) + + expectation_value += coeff * exp_val + + if rank == 0: + return float(np.real(expectation_value)), rank + else: + return 0.0, rank 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..7450d4d3 --- /dev/null +++ b/tests/test_quimb_mpi_backend.py @@ -0,0 +1,102 @@ +# 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