improve performance

direct write data to GPU memory if already allocated

Author: derwind

qmlant/neural_networks/estimator_tn.py

@@ -31,8 +31,7 @@ def prepare_circuit(
         expr: str,
         operands: list[cp.ndarray],
     ) -> tuple[str, list[cp.ndarray]]:
-        """prepare a circuit for forward process setting batch and parameters to the circuit.
-        """
+        """prepare a circuit for forward process setting batch and parameters to the circuit."""
 
         pname2theta_list = {
             f"x[{i}]": batch[:, i].flatten().tolist() for i in range(batch.shape[1])

qmlant/neural_networks/neural_network.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 
-def Ry(theta: float, xp=cp):
+def Ry(theta: float, xp=cp) -> cp.ndarray:
     cos = xp.cos(theta / 2)
     sin = xp.sin(theta / 2)
     return xp.array(
@@ -13,11 +13,19 @@ def Ry(theta: float, xp=cp):
     )
 
 
-def Ry_Rydag(theta: float, xp=cp):
+def Ry_Rydag(theta: float, xp=cp) -> tuple[np.ndarray | cp.ndarray, np.ndarray | cp.ndarray]:
     cos = xp.cos(theta / 2)
     sin = xp.sin(theta / 2)
     ry_rydag = xp.array(
         [[[cos, -sin], [sin, cos]], [[cos, sin], [-sin, cos]]],
         dtype=complex,
     )
     return ry_rydag[0], ry_rydag[1]
+
+
+def Ry_Rydag_direct(theta: float, mat: cp.ndarray, mat_dag: cp.ndarray) -> None:
+    cos = cp.cos(theta / 2)
+    sin = cp.sin(theta / 2)
+    mat[0][0] = mat[1][1] = mat_dag[0][0] = mat_dag[1][1] = cos
+    mat[0][1] = mat_dag[1][0] = -sin
+    mat[1][0] = mat_dag[0][1] = sin

qmlant/neural_networks/utils/utils.py

@@ -9,7 +9,7 @@
 from cuquantum import CircuitToEinsum
 from qiskit import QuantumCircuit
 
-from ..neural_network import Ry_Rydag
+from ..neural_network import Ry_Rydag, Ry_Rydag_direct
 
 
 @overload
@@ -101,10 +101,8 @@ def replace_ry(
     pname2locs: dict[str, tuple[int, int]],
 ) -> list[cp.ndarray]:
     for pname, theta in pname2theta.items():  # e.g. pname[0] = "θ[0]"
-        ry, ry_dag = Ry_Rydag(theta)
         loc, dag_loc = pname2locs[pname]
-        operands[loc] = ry
-        operands[dag_loc] = ry_dag
+        Ry_Rydag_direct(theta, operands[loc], operands[dag_loc])
 
     return operands
 
@@ -119,10 +117,8 @@ def replace_ry_phase_shift(
     # θ[0]: [π/2, -π/2], θ[1]: [π/2, -π/2], ...
     for pname, theta in pname2theta.items():  # e.g. pname[0] = "θ[0]"
         for phase_shift in phase_shift_list:
-            ry, ry_dag = Ry_Rydag(theta + phase_shift)
             loc, dag_loc = pname2locs[pname]
-            operands[loc][i, :] = ry
-            operands[dag_loc][i, :] = ry_dag
+            Ry_Rydag_direct(theta + phase_shift, operands[loc][i], operands[dag_loc][i])
             i += 1
 
     return operands