diff --git a/src/qibotn/QiboCircuitConvertor.py b/src/qibotn/QiboCircuitConvertor.py index 9212f70..d72a09c 100644 --- a/src/qibotn/QiboCircuitConvertor.py +++ b/src/qibotn/QiboCircuitConvertor.py @@ -94,7 +94,8 @@ class QiboCircuitToEinsum: required_shape = self.op_shape_from_qubits(len(gate_qubits)) self.gate_tensors.append( ( - cp.asarray(gate.matrix).reshape(required_shape), + cp.asarray(gate.matrix(), dtype=self.dtype).reshape( + required_shape), gate_qubits, ) ) diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py index ae36497..eb0e0d4 100644 --- a/src/qibotn/cutn.py +++ b/src/qibotn/cutn.py @@ -1,5 +1,9 @@ from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum from cuquantum import contract +from cuquantum import cutensornet as cutn +import multiprocessing +from cupy.cuda.runtime import getDeviceCount +import cupy as cp from qibotn.QiboCircuitToMPS import QiboCircuitToMPS from qibotn.mps_contraction_helper import MPSContractionHelper @@ -10,6 +14,42 @@ def eval(qibo_circ, datatype): return contract(*myconvertor.state_vector_operands()) +def eval_tn_MPI(qibo_circ, datatype, n_samples=8): + """Convert qibo circuit to tensornet (TN) format and perform contraction using multi node and multi GPU through MPI. + The conversion is performed by QiboCircuitToEinsum(), after which it goes through 2 steps: pathfinder and execution. + The pathfinder looks at user defined number of samples (n_samples) iteratively to select the least costly contraction path. This is sped up with multi thread. + After pathfinding the optimal path is used in the actual contraction to give a dense vector representation of the TN. + """ + + from mpi4py import MPI # this line initializes MPI + + ncpu_threads = multiprocessing.cpu_count() // 2 + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + device_id = rank % getDeviceCount() + cp.cuda.Device(device_id).use() + + handle = cutn.create() + cutn.distributed_reset_configuration(handle, *cutn.get_mpi_comm_pointer(comm)) + network_opts = cutn.NetworkOptions(handle=handle, blocking="auto") + + # Perform circuit conversion + myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype) + operands_interleave = myconvertor.state_vector_operands() + + # Pathfinder: To search for the optimal path. Optimal path are assigned to path and info attribute of the network object. + network = cutn.Network(*operands_interleave, options=network_opts) + network.contract_path(optimize={"samples": n_samples, "threads": ncpu_threads}) + + # Execution: To execute the contraction using the optimal path found previously + result = network.contract() + + cutn.destroy(handle) + + return result, rank + + def eval_mps(qibo_circ, gate_algo, datatype): myconvertor = QiboCircuitToMPS(qibo_circ, gate_algo, dtype=datatype) mps_helper = MPSContractionHelper(myconvertor.num_qubits)