Black formate
This commit is contained in:
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user