jaunatisblue/qibotn
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

完善脚本功能，添加计时估计

Browse Source
This commit is contained in:
jaunatisblue
2026-05-08 10:20:03 +08:00
parent 8b71ff96c8
commit cec0ba272a
4 changed files with 21 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
benchmark_contract_sliced.py
View File

@@ -31,13 +31,13 @@ else:
tree = None tree = None
tree = comm.bcast(tree, root=0) tree = comm.bcast(tree, root=0)
arrays = [torch.from_numpy(np.ascontiguousarray(t._data, dtype=np.complex128)) for t in tn.tensors] arrays = [torch.from_numpy(np.asarray(t._data)) for t in tn.tensors]
n_slices = tree.multiplicity n_slices = tree.multiplicity
if rank == 0: if rank == 0:
print(f"Slices: {n_slices}, Ranks: {size}, " print(f"Slices: {n_slices}, Ranks: {size}, "
f"Peak: {tree.max_size() * 16 / 1e9:.2f} GB, " f"Peak: {tree.max_size() * 16 / 1e9:.2f} GB, "
f"Threads/rank: {max(1, NCORES // size)}, Backend: torch") f"Threads/rank: {NCORES}, Backend: torch")
t0 = time.time() t0 = time.time()
result = None result = None

2
benchmark_slice.py
View File

@@ -8,7 +8,7 @@ with open(f"data/tree_q{NQUBITS}_l{NLAYERS}.pkl", 'rb') as f:
print(f"Original peak: {tree.max_size() * 16 / 1e9:.2f} GB") print(f"Original peak: {tree.max_size() * 16 / 1e9:.2f} GB")
tree_sliced = tree.slice_and_reconfigure(target_size=2**30) # 2^29 = 8 GB tree_sliced = tree.slice_and_reconfigure(target_size=2**28)
with open(f"data/tree_q{NQUBITS}_l{NLAYERS}_sliced.pkl", 'wb') as f: with open(f"data/tree_q{NQUBITS}_l{NLAYERS}_sliced.pkl", 'wb') as f:
pickle.dump(tree_sliced, f) pickle.dump(tree_sliced, f)

20
check_tree.py
View File

@@ -5,9 +5,21 @@ path = sys.argv[1] if len(sys.argv) > 1 else "data/tree_q25_l10.pkl"
with open(path, 'rb') as f: with open(path, 'rb') as f:
tree = pickle.load(f) tree = pickle.load(f)
# Intel 8558P: 96 cores, 2.1GHz, AVX-512 (16 FP64/cycle), FMA x2
# complex128 multiply-add = 6 real FLOPs
CORES = 96
FREQ = 2.1e9
AVX512_FP64 = 16
TFLOPS = CORES * FREQ * AVX512_FP64 * 2 / 1e12 # ~6.45 TFLOPS real FP64
COMPLEX_FLOPS = TFLOPS / 6 # complex128 effective
flops = tree.total_flops()
slices = tree.multiplicity
est_seconds = flops * slices / (COMPLEX_FLOPS * 1e12)
print(f"File: {path}") print(f"File: {path}")
print(f"Peak memory elements: {tree.max_size():.2e}") print(f"Peak memory (GB): {tree.max_size() * 16 / 1e9:.2f}")
print(f"Peak memory (GB): {tree.max_size() * 16 / 1e9:.2f}") # complex128 = 16 bytes print(f"Total FLOPs: {flops:.2e} x{slices} slices = {flops*slices:.2e}")
print(f"Total FLOPs: {tree.total_flops():.2e}")
print(f"Contraction width: {tree.contraction_width()}") print(f"Contraction width: {tree.contraction_width()}")
print(f"Multiplicity (slices): {tree.multiplicity}") print(f"Multiplicity (slices): {slices}")
print(f"Estimated time (96 cores): {est_seconds:.1f}s ({est_seconds/3600:.2f}h)")

2
src/qibotn/backends/quimb.py
View File

@@ -439,6 +439,7 @@ def _expectation_parallel(self, circuit, observable, method, opts):
mpi_contract = opts.get('mpi_contract', False) mpi_contract = opts.get('mpi_contract', False)
torch_threads = opts.get('torch_threads', None) torch_threads = opts.get('torch_threads', None)
slicing_opts = opts.get('slicing_opts', None) slicing_opts = opts.get('slicing_opts', None)
trial_timeout = opts.get('trial_timeout', None)
qc = self._qibo_circuit_to_quimb( qc = self._qibo_circuit_to_quimb(
circuit, circuit,
@@ -472,6 +473,7 @@ def _expectation_parallel(self, circuit, observable, method, opts):
max_time=max_time, max_time=max_time,
n_workers=search_workers, n_workers=search_workers,
slicing_opts=slicing_opts, slicing_opts=slicing_opts,
trial_timeout=trial_timeout,
) )
if tree is None: if tree is None: