bell_state_on_gpu.py

import numpy as np
from numba import cuda

# Define Pauli matrices
sigma_x = np.array([[0, 1], [1, 0]], dtype=np.complex128)
sigma_z = np.array([[1, 0], [0, -1]], dtype=np.complex128)

# Define Hadamard gate
hadamard = (1 / np.sqrt(2)) * np.array([[1, 1], [1, -1]], dtype=np.complex128)

# Define CNOT gate
cnot = np.array([[1, 0, 0, 0],
                 [0, 1, 0, 0],
                 [0, 0, 0, 1],
                 [0, 0, 1, 0]], dtype=np.complex128)

@cuda.jit
def bell_state_kernel(qubits, out):
    tx = cuda.threadIdx.x
    ty = cuda.blockIdx.x
    bw = cuda.blockDim.x
    idx = tx + ty * bw
    if idx < out.size:
        # Apply Hadamard gate to first qubit
        qubits[idx][0] = hadamard[0, 0] * qubits[idx][0] + hadamard[0, 1] * qubits[idx][1]
        qubits[idx][1] = hadamard[1, 0] * qubits[idx][0] + hadamard[1, 1] * qubits[idx][1]

        # Apply CNOT gate
        qubits[idx][2] = qubits[idx][0] * cnot[0, 0] + qubits[idx][1] * cnot[0, 1] + qubits[idx][2] * cnot[0, 2] + qubits[idx][3] * cnot[0, 3]
        qubits[idx][3] = qubits[idx][0] * cnot[1, 0] + qubits[idx][1] * cnot[1, 1] + qubits[idx][2] * cnot[1, 2] + qubits[idx][3] * cnot[1, 3]

        # Reset qubits 0 and 1
        qubits[idx][0] = 0
        qubits[idx][1] = 0

        out[idx][0] = qubits[idx][2]
        out[idx][1] = qubits[idx][3]

# Main function to call the CUDA kernel
def apply_bell_state_on_gpu(qubits):
    # Allocate memory on the device
    d_qubits = cuda.to_device(qubits.astype(np.complex128))
    d_out = cuda.device_array_like(qubits)

    # Configure the kernel
    threads_per_block = 256
    blocks_per_grid = (qubits.shape[0] + (threads_per_block - 1)) // threads_per_block

    # Launch the kernel
    bell_state_kernel[blocks_per_grid, threads_per_block](d_qubits, d_out)

    # Copy the result back to the host
    return d_out.copy_to_host()

# Example usage
N = 2
qubits = np.array([[1, 0] for _ in range(N)], dtype=np.complex128)
print("Initial qubits:")
print(qubits)
result = apply_bell_state_on_gpu(qubits)
print(result)