diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py index 3d42eb7..5267bc0 100644 --- a/src/qibotn/cutn.py +++ b/src/qibotn/cutn.py @@ -1,7 +1,6 @@ from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum from cuquantum import contract from cuquantum import cutensornet as cutn -from mpi4py import MPI # this line initializes MPI import multiprocessing from cupy.cuda.runtime import getDeviceCount import cupy as cp @@ -12,9 +11,16 @@ def eval(qibo_circ, datatype): return contract(*myconvertor.state_vector_operands()) -def eval_tn_MPI(qibo_circ, datatype): +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() afterwhich 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 - n_samples = 8 comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -25,14 +31,15 @@ def eval_tn_MPI(qibo_circ, datatype): 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} - ) # Calculate optimal path, returns path and info + 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)