Merge pull request #298 from ghanem-nv/cuquantum-dev

cuTN integration

Author: zkysfls

examples/cuquantum/qaoa.py

@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+
+import math
+import argparse
+
+import torch
+from torch import nn
+from torchquantum.plugin.cuquantum import *
+from torchquantum.operator.standard_gates import *
+
+
+
+
+
+class MAXCUT(nn.Module):
+    def __init__(self, n_wires, input_graph, n_layers):
+        super().__init__()
+        self.n_wires = n_wires
+        self.input_graph = input_graph
+        self.n_layers = n_layers
+
+        self.circuit = ParameterizedQuantumCircuit(n_wires=n_wires, n_input_params=0, n_trainable_params=2 * n_layers)
+        self.circuit.set_trainable_params(torch.randn(2 * n_layers))
+
+        for wire in range(self.n_wires):
+            self.circuit.append_gate(Hadamard, wires=wire)
+
+        for l in range(self.n_layers):
+            # mixer layer
+            for i in range(self.n_wires):
+                self.circuit.append_gate(RX, wires=i, trainable_idx=l)
+
+            # entangler layer
+            for edge in self.input_graph:
+                self.circuit.append_gate(CNOT, wires=[edge[0], edge[1]])
+                self.circuit.append_gate(RZ, wires=edge[1], trainable_idx=n_layers + l)
+                self.circuit.append_gate(CNOT, wires=[edge[0], edge[1]])
+
+
+        hamiltonian = {}
+        for edge in self.input_graph:
+            pauli_string = ""
+            for wire in range(self.n_wires):
+                if wire in edge:
+                    pauli_string += "Z"
+                else:
+                    pauli_string += "I"
+            hamiltonian[pauli_string] = 0.5
+
+        backend = CuTensorNetworkBackend(TNConfig(num_hyper_samples=10))
+        self.energy = QuantumExpectation(self.circuit, backend, [hamiltonian])
+        self.sampling = QuantumSampling(self.circuit, backend, 100)
+
+    def forward(self):
+        start_time = torch.cuda.Event(enable_timing=True)
+        end_time = torch.cuda.Event(enable_timing=True)
+
+        start_time.record()
+        output = self.energy() - len(self.input_graph) / 2
+        end_time.record()
+
+        torch.cuda.synchronize()
+        elapsed_time = start_time.elapsed_time(end_time)
+        print(f"Forward pass took {elapsed_time:.2f} ms")
+
+        return output
+
+
+def optimize(model, n_steps=100, lr=0.1):
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    print(f"The initial parameters are:\n{next(model.parameters()).data.tolist()}")
+    print("")
+    for step in range(n_steps):
+        optimizer.zero_grad()
+        loss = model()
+        start_time = torch.cuda.Event(enable_timing=True)
+        end_time = torch.cuda.Event(enable_timing=True)
+
+        start_time.record()
+        loss.backward()
+        end_time.record()
+
+        torch.cuda.synchronize()
+        elapsed_time = start_time.elapsed_time(end_time)
+        print(f"Backward pass took {elapsed_time:.2f} ms")
+
+        optimizer.step()
+
+        print(f"Step: {step}, Cost Objective: {loss.item()}")
+
+    print("")
+    print(f"The optimal parameters are:\n{next(model.parameters()).data.tolist()}")
+    print("")
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n_wires", type=int, default=4, help="number of wires")
+    parser.add_argument("--n_layers", type=int, default=4, help="number of layers")
+    parser.add_argument("--steps", type=int, default=100, help="number of steps")
+    parser.add_argument("--lr", type=float, default=0.01, help="learning rate")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    args = parser.parse_args()
+
+    torch.manual_seed(args.seed)
+
+    # create a fully connected graph
+    input_graph = []
+    for i in range(args.n_wires):
+        for j in range(i):
+            input_graph.append((i, j))
+
+    print(f"Cost Objective Minimum (Analytic Reference Result): {math.floor(args.n_wires**2 // 4)}")
+
+    model = MAXCUT(n_wires=args.n_wires, input_graph=input_graph, n_layers=args.n_layers)
+    optimize(model, n_steps=args.steps, lr=args.lr)
+    samples = model.sampling()
+
+    print(f"Sampling Results: {samples}")
+
+
+if __name__ == "__main__":
+    main()

torchquantum/plugin/cuquantum/LICENSE

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

torchquantum/plugin/cuquantum/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+
+from .circuit import ParameterizedQuantumCircuit
+from .cutn import CuTensorNetworkBackend, TNConfig, MPSConfig
+from .expectation import QuantumExpectation
+from .sampling import QuantumSampling
+from .amplitude import QuantumAmplitude
+
+__all__ = [
+    "ParameterizedQuantumCircuit",
+    "CuTensorNetworkBackend",
+    "TNConfig",
+    "MPSConfig",
+    "QuantumExpectation",
+    "QuantumSampling",
+    "QuantumAmplitude",
+]

torchquantum/plugin/cuquantum/amplitude.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+
+from typing import List
+
+import torch.nn as nn
+
+from .utils import check_input_params
+from .backend import QuantumBackend
+from .circuit import ParameterizedQuantumCircuit
+
+
+class QuantumAmplitude(nn.Module):
+    """A PyTorch module for computing quantum state amplitudes.
+
+    This module computes the amplitudes of specified bitstrings in the quantum state prepared by a given quantum circuit.
+
+    Args:
+        circuit: The quantum circuit that prepares the state.
+        backend: The quantum backend to use for computation.
+        bitstrings: List of bitstrings whose amplitudes to compute.
+    """
+
+    def __init__(self, circuit: ParameterizedQuantumCircuit, backend: QuantumBackend, bitstrings: List[str]):
+        super().__init__()
+        self._circuit = circuit.copy()
+        self._bitstrings = bitstrings.copy()
+        self._backend = backend
+        self._amplitude_module = self.backend._create_amplitude_module(circuit, bitstrings)
+
+    def forward(self, input_params=None):
+        """Compute the amplitudes for the bitstrings specified in the constructor.
+
+        Args:
+            input_params: 2D Tensor of input parameters for the circuit. Shape should be (batch_size, n_input_params). If
+                only one batch is being processed, the tensor can be instead a 1D tensor with shape (n_input_params,). If
+                the circuit has no input parameters, this argument can be omitted (i.e. None).
+
+        Returns:
+            2D Tensor of amplitudes for each bitstring in each batch. The shape is (batch_size, len(bitstrings)).
+        """
+        input_params = check_input_params(input_params, self._circuit.n_input_params)
+        return self._amplitude_module(input_params)
+
+    @property
+    def bitstrings(self):
+        """Get the list of bitstrings whose amplitudes are being computed."""
+        return self._bitstrings.copy()
+
+    @property
+    def circuit(self):
+        """Get the quantum circuit used for state preparation."""
+        return self._circuit.copy()
+
+    @property
+    def backend(self):
+        """Get the quantum backend being used for computation."""
+        return self._backend

torchquantum/plugin/cuquantum/backend.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+
+from abc import ABC, abstractmethod
+from typing import List, Union, Dict, Optional
+
+import torch.nn as nn
+
+from .circuit import ParameterizedQuantumCircuit
+
+
+class QuantumBackend(ABC):
+    """Abstract base class for quantum backends.
+
+    This class defines the interface that all quantum backends must implement. Each backend must provide methods for 
+    creating PyTorch modules that compute:
+    - Expectation values of Pauli operators.
+    - State amplitudes for given bitstrings.
+    - Sampling from the quantum state.
+    """
+
+    @abstractmethod
+    def _create_expectation_module(
+        self, circuit: ParameterizedQuantumCircuit, pauli_ops: Union[List[str], Dict[str, float]]
+    ) -> nn.Module:
+        """Create a module for computing expectation values of Pauli operators.
+
+        Args:
+            circuit: The quantum circuit that prepares the state
+            pauli_ops: List of Pauli operators to compute expectations for. Each Pauli operator can be either:
+                - A single Pauli string specifying the pauli operator for each qubit ("I", "X", "Y", or "Z").
+                - A linear combination of Pauli strings specified as a dictionary mapping each single Pauli string to its 
+                  corresponding coefficient.
+
+        Returns:
+            A PyTorch module that computes the expectation values.
+        """
+        pass
+
+    @abstractmethod
+    def _create_amplitude_module(self, circuit: ParameterizedQuantumCircuit, bitstrings: List[str]) -> nn.Module:
+        """Create a module for computing state amplitudes.
+
+        Args:
+            circuit: The quantum circuit that prepares the state.
+            bitstrings: List of bitstrings whose amplitudes to compute.
+
+        Returns:
+            A PyTorch module that computes the amplitudes.
+        """
+        pass
+
+    @abstractmethod
+    def _create_sampling_module(
+        self, circuit: ParameterizedQuantumCircuit, n_samples: int, wires: Optional[List[int]] = None
+    ) -> nn.Module:
+        """Create a module for sampling from the quantum state.
+
+        Args:
+            circuit: The quantum circuit that prepares the state.
+            n_samples: Number of samples to generate.
+            wires: Optional list of wires/qubits to sample from. If not provided, all wires/qubits are sampled from.
+
+        Returns:
+            A PyTorch module that generates samples from the quantum state.
+        """
+        pass