Make CUDAQ objects pickleable

[zohimchandani]

Required prerequisites

• Search the issue tracker to check if your feature has already been mentioned or rejected in other issues.

Describe the feature

I am trying to manually hand-code the hamiltonian batching workflow where I distribute via MPI:

import cudaq
from mpi4py import MPI
import cupy as cp 
from cudaq import spin 

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
num_ranks = comm.Get_size()

cp.cuda.runtime.setDevice(rank)
print(f'rank {rank} running with GPU {cp.cuda.Device().pci_bus_id}')

cudaq.set_target("nvidia")


qubit_count = 4
@cudaq.kernel
def kernel(qubit_count: int):
    
    qvector = cudaq.qvector(qubit_count)
    
    h(qvector[0])
    for i in range(1, qubit_count):
        cx(qvector[0], qvector[i])


# we have a hamiltonian 
h = spin.x(0) + spin.y(1) + spin.z(2) + spin.x(3) 

# we split it into 4 batches for distribution across 4 gpus via mpi 
batched_h = [[spin.x(0)], [spin.y(1)], [spin.z(2)], [spin.x(3)]]

# we distribute each term across the gpus available 
comm.scatter(batched_h,  root = 0)

root@gorby:/home/cudaq/cudaq_work# mpirun -np 4 --allow-run-as-root python3 mpi.py
rank 2 running with GPU 0000:81:00.0
rank 3 running with GPU 0000:C2:00.0
rank 1 running with GPU 0000:47:00.0
rank 0 running with GPU 0000:01:00.0
TypeError: cannot pickle 'cudaq.mlir._mlir_libs._quakeDialects.cudaq_runtime.SpinOperator' object

Can we make CUDAQ types pickleable?

Thanks team.

[amccaskey]

You could probably get away with serializing / deserializing the spin operator

serialized = op.serialize() # list of floats
# distribute / scatter the data 
deserialized = cudaq.SpinOperator(serialized, numQubits)