more options for benchmark cli

Author: refraction-ray

benchmarks/scripts_v2/benchmark_cli.py

@@ -3,10 +3,10 @@
 import json
 import argparse
 import time
-import numpy as np
-import cpuinfo
 import sys
 import datetime
+import numpy as np
+import cpuinfo
 import tensorcircuit as tc
 
 from benchmark_core import benchmark_mega_function
@@ -131,6 +131,22 @@ def arg():
         help="Contractor setting (e.g., cotengra-16-128)",
         default=[None],
     )
+    parser.add_argument(
+        "-bond_dim",
+        dest="bond_dim",
+        type=int,
+        nargs=1,
+        help="Bond dimension for MPS circuits",
+        default=[16],
+    )
+    parser.add_argument(
+        "-jit_compile",
+        dest="jit_compile",
+        type=int,
+        nargs=1,
+        help="Whether to use JIT compilation (0 or 1)",
+        default=[1],
+    )
     args = parser.parse_args()
     return [
         args.n[0],
@@ -152,6 +168,8 @@ def arg():
         args.backend[0],
         args.dtype[0],
         args.contractor[0],
+        args.bond_dim[0],
+        args.jit_compile[0],
     ]
 
 
@@ -204,6 +222,8 @@ def benchmark_cli(
     backend,
     dtype,
     contractor,
+    bond_dim,
+    jit_compile,
     path,
 ):
     meta = {}
@@ -244,6 +264,8 @@ def benchmark_cli(
         "backend": backend,
         "dtype": dtype,
         "contractor": contractor,
+        "bond_dim": bond_dim,
+        "jit_compile": jit_compile,
     }
     meta["UUID"] = uuid
     meta["Benchmark Time"] = (
@@ -258,13 +280,15 @@ def benchmark_cli(
         lx=lx,
         ly=ly,
         circuit_type=circuit_type,
+        bond_dim=bond_dim,
         layout_type=layout_type,
         operation=operation,
         noisy=bool(noisy),
         noisy_type=noisy_type,
         use_grad=bool(use_grad),
         use_vmap=bool(use_vmap),
         contractor=contractor,
+        jit_compile=bool(jit_compile),
     )
 
     # Create parameters for testing
@@ -313,6 +337,8 @@ def benchmark_cli(
         backend,
         dtype,
         contractor,
+        bond_dim,
+        jit_compile,
     ) = arg()
 
     results = benchmark_cli(
@@ -335,6 +361,8 @@ def benchmark_cli(
         backend,
         dtype,
         contractor,
+        bond_dim,
+        jit_compile,
         path,
     )
     save(results, _uuid, path)

benchmarks/scripts_v2/benchmark_core.py

@@ -36,7 +36,6 @@ def generate_2d_circuit(c, lx, ly, params, nqubits, nlayers):
 
 def generate_noisy_circuit(c, status, type="depolarizing"):
     noise_conf = noisemodel.NoiseConf()
-    # print(type)
     if type == "depolarizing":
         error1 = channels.depolarizingchannel(0.1, 0.1, 0.1)  # px, py, pz probabilities
     elif type == "amplitudedamping":
@@ -78,6 +77,7 @@ def benchmark_mega_function(
     use_grad=False,  # True, False
     use_vmap=False,  # True, False
     contractor=None,  # contractor setting like "cotengra-16-128"
+    jit_compile=True,  # True, False
 ):
     """
     Mega benchmark function that can control all parameters via arguments.
@@ -88,12 +88,13 @@ def benchmark_mega_function(
         lx: Lattice size x (for 2D)
         ly: Lattice size y (for 2D)
         circuit_type: Type of circuit ("circuit", "dmcircuit", "mpscircuit")
+        bond_dim: Bond dimension for MPS circuits
         layout_type: Circuit layout ("1d", "2d")
         operation: Operation to perform ("state", "sample", "exps")
         noisy: Whether to add noise (only for "circuit" and "dmcircuit")
+        noisy_type: Type of noise channel ("depolarizing", "amplitudedamping")
         use_grad: Whether to compute gradient (AD)
         use_vmap: Whether to use vectorized operations
-        batch_size: Batch size for vmap operations
         contractor: Contractor setting like "cotengra-16-128"
 
     Returns:
@@ -136,17 +137,17 @@ def circuit_func(params):
     # Handle gradient computation
     if use_grad and not use_vmap:
         grad_func = tc.backend.grad(circuit_func)
-        return tc.backend.jit(grad_func, jit_compile=True)
+        return tc.backend.jit(grad_func, jit_compile=jit_compile)
 
     # Handle vmap computation
     if use_vmap and not use_grad:
-        return tc.backend.jit(tc.backend.vmap(circuit_func), jit_compile=True)
+        return tc.backend.jit(tc.backend.vmap(circuit_func), jit_compile=jit_compile)
 
     # Handle both grad and vmap
     if use_grad and use_vmap:
         vvag_func = tc.backend.vvag(circuit_func)
-        return tc.backend.jit(vvag_func, jit_compile=True)
+        return tc.backend.jit(vvag_func, jit_compile=jit_compile)
 
     # Regular operation (no grad, no vmap)
     # Always JIT the returned function
-    return tc.backend.jit(circuit_func, jit_compile=True)
+    return tc.backend.jit(circuit_func, jit_compile=jit_compile)