Merge branch 'add_collapse_lgpu' into add_LGPUSIM

PennyLaneAI · Oct 10, 2024 · f6b5202 · f6b5202
2 parents 7d82d15 + 0e9f1db
commit f6b5202
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 12 deletions.
diff --git a/mpitests/test_native_mcm.py b/mpitests/test_native_mcm.py
@@ -18,9 +18,6 @@
 from conftest import LightningDevice, device_name
 from mpi4py import MPI
 
-if device_name not in ("lightning.qubit", "lightning.kokkos", "lightning.gpu"):
-    pytest.skip("Native MCM not supported. Skipping.", allow_module_level=True)
-
 if not LightningDevice._CPP_BINARY_AVAILABLE:  # pylint: disable=protected-access
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)
 
@@ -41,7 +38,6 @@ def func(x, y):
     comm.Barrier()
 
     with pytest.raises(
-        qml.DeviceError,
-        match=f"LightningGPU-MPI does not support Mid-circuit measurements.",
+        qml.DeviceError, match="LightningGPU-MPI does not support Mid-circuit measurements."
     ):
         func(*params)
diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.39.0-dev42"
+__version__ = "0.39.0-dev43"
diff --git a/pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaMPI.hpp b/pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaMPI.hpp
@@ -355,7 +355,7 @@ class StateVectorCudaMPI final
             collapse_local_(wireInt, branch);
         } else {
             // global wire
-            int local_wire = 0;
+            constexpr int local_wire = 0;
             std::vector<int2> wirePairs{make_int2(wireInt, local_wire)};
             applyMPI_Dispatcher(wirePairs, &StateVectorCudaMPI::collapse_local_,
                                 local_wire, branch);
@@ -1680,12 +1680,12 @@ class StateVectorCudaMPI final
             /* const int32_t * */ basisBits.data(),
             /* const uint32_t nBasisBits */ basisBits.size()));
 
-        auto abs2sum0 = mpi_manager_.allreduce(abs2sum0_local, "sum");
-        auto abs2sum1 = mpi_manager_.allreduce(abs2sum1_local, "sum");
+        auto abs2sum0 = mpi_manager_.allreduce<double>(abs2sum0_local, "sum");
+        auto abs2sum1 = mpi_manager_.allreduce<double>(abs2sum1_local, "sum");
 
         double norm = (branch == 0) ? abs2sum0 : abs2sum1;
 
-        int parity = branch;
+        const int parity = static_cast<int>(branch);
 
         PL_CUSTATEVEC_IS_SUCCESS(custatevecCollapseOnZBasis(
             /* custatevecHandle_t */ handle_.get(),

diff --git a/pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaManaged.hpp b/pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaManaged.hpp
@@ -522,7 +522,7 @@ class StateVectorCudaManaged
 
         double norm = (branch == 0) ? abs2sum0 : abs2sum1;
 
-        int parity = branch;
+        const int parity = static_cast<int>(branch);
 
         PL_CUSTATEVEC_IS_SUCCESS(custatevecCollapseOnZBasis(
             /* custatevecHandle_t */ handle_.get(),

diff --git a/pennylane_lightning/lightning_gpu/lightning_gpu.py b/pennylane_lightning/lightning_gpu/lightning_gpu.py
@@ -516,7 +516,7 @@ def simulate(
                 raise qml.DeviceError("LightningGPU-MPI does not support Mid-circuit measurements.")
 
             results = []
-            aux_circ = qml.tape.QuantumScript(
+            aux_circ = QuantumScript(
                 circuit.operations,
                 circuit.measurements,
                 shots=[1],