Black formate

This commit is contained in:
tankya2
2023-08-30 17:26:07 +08:00
parent 1c9df26472
commit d12b8ab882

View File

@@ -12,50 +12,50 @@ def eval(qibo_circ, datatype):
def eval_tn_MPI(qibo_circ, datatype): def eval_tn_MPI(qibo_circ, datatype):
ncpu_threads = multiprocessing.cpu_count() // 2 ncpu_threads = multiprocessing.cpu_count() // 2
n_samples = 8 n_samples = 8
root = 0 root = 0
comm = MPI.COMM_WORLD comm = MPI.COMM_WORLD
rank = comm.Get_rank() rank = comm.Get_rank()
device_id = rank % getDeviceCount() device_id = rank % getDeviceCount()
cp.cuda.Device(device_id).use() cp.cuda.Device(device_id).use()
handle = cutn.create() handle = cutn.create()
cutn.distributed_reset_configuration(handle, *cutn.get_mpi_comm_pointer(comm)) cutn.distributed_reset_configuration(handle, *cutn.get_mpi_comm_pointer(comm))
network_opts = cutn.NetworkOptions(handle=handle, blocking="auto") network_opts = cutn.NetworkOptions(handle=handle, blocking="auto")
myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype) myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
operands_interleave = myconvertor.state_vector_operands() operands_interleave = myconvertor.state_vector_operands()
network = cutn.Network(*operands_interleave, options=network_opts) 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}
) # Calculate optimal path, returns path and info
result = network.contract() result = network.contract()
cutn.destroy(handle) cutn.destroy(handle)
if rank == root: if rank == root:
return result, rank return result, rank
if __name__ == "__main__": if __name__ == "__main__":
from qibo.models import QFT
from qibo.models import QFT
import cupy as cp import cupy as cp
import numpy as np import numpy as np
num_qubits = 10 num_qubits = 10
swaps = True swaps = True
circ_qibo = QFT(num_qubits, swaps) circ_qibo = QFT(num_qubits, swaps)
dtype="complex128" dtype = "complex128"
sv_mpi, rank = eval_tn_MPI(circ_qibo, dtype) sv_mpi, rank = eval_tn_MPI(circ_qibo, dtype)
if rank == 0: if rank == 0:
sv_reference = eval(circ_qibo, dtype) sv_reference = eval(circ_qibo, dtype)
state_vec = np.array(circ_qibo()) state_vec = np.array(circ_qibo())
print(f"State vector difference: {abs(sv_mpi-sv_reference).max():0.3e}") print(f"State vector difference: {abs(sv_mpi-sv_reference).max():0.3e}")
assert cp.allclose(sv_mpi, sv_reference) assert cp.allclose(sv_mpi, sv_reference)
assert cp.allclose(sv_mpi.flatten(), state_vec) assert cp.allclose(sv_mpi.flatten(), state_vec)