Merge remote-tracking branch 'origin/bench' into bench

danlkv · Dec 22, 2023 · 576acfb · 576acfb
2 parents 27b9999 + b52d487
commit 576acfb
Show file tree

Hide file tree

Showing 13 changed files with 3,089 additions and 3,006 deletions.
diff --git a/bench/qc_simulation/src/simulators/qtensor.py b/bench/qc_simulation/src/simulators/qtensor.py
@@ -161,6 +161,7 @@ def simulate_tn(prep_data,
     import time
     from qtensor.contraction_algos import bucket_elimination
     from qtensor.compression.Compressor import CUSZCompressor, CUSZXCompressor, TorchCompressor, NEWSZCompressor
+    from qtensor.compression.Compressor import WriteToDiskCompressor
     import cupy
     cupy.cuda.profiler.start()
     peo, par_vars, tn = prep_data
@@ -184,6 +185,9 @@ def simulate_tn(prep_data,
             print(f"{r2r_error=} {r2r_threshold=}")
             compressor = NEWSZCompressor(r2r_error=r2r_error, r2r_threshold=r2r_threshold)
             compressor = qtensor.compression.ProfileCompressor(compressor)
+        elif compress == 'disk':
+            compressor = WriteToDiskCompressor(f'/grand/QTensor/compression/data/tensors_compressed_M{M}/')
+            compressor = qtensor.compression.ProfileCompressor(compressor)
         else:
             raise ValueError(f"Unknown compression algorithm: {compress}")
         backend = qtensor.contraction_backends.CompressionBackend(backend, compressor, M)