Added comments

This commit is contained in:
tankya2
2023-10-04 11:18:05 +08:00
parent cec56f36d1
commit 15e90ebcc7

View File

@@ -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)