cutn integration draft

Signed-off-by: vedika-saravanan <vsaravanan@nvidia.com>

Author: vedika-saravanan

libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/nm_optimizer.py

@@ -283,7 +283,8 @@ class NMOptimizer(TensorNetworkDecoder):
         execute: Forward backend.  ``"codegen"`` (default) partial-evaluates
             the path into a flat Python function; ``"unrolled"`` keeps an
             interpretive einsum list; ``"opt_einsum"`` dispatches via
-            :func:`opt_einsum.contract_expression`.
+            :func:`opt_einsum.contract_expression`; ``"cutensornet"``
+            executes the reduced contraction with cuTensorNet on CUDA.
         compile_mode: Forwarded to :func:`torch.compile`; ignored when
             ``compile=False``.
         dynamic_syndromes: If ``True`` (default), syndromes are runtime
@@ -324,11 +325,12 @@ def __init__(
         device: str = "cuda",
         *,
         compile: bool = False,
-        execute: Literal["codegen", "unrolled", "opt_einsum"] = "codegen",
+        execute: Literal["codegen", "unrolled", "opt_einsum",
+                         "cutensornet"] = "codegen",
         compile_mode: str | None = None,
         dynamic_syndromes: bool = True,
     ) -> None:
-        if execute not in ("unrolled", "opt_einsum", "codegen"):
+        if execute not in ("unrolled", "opt_einsum", "codegen", "cutensornet"):
             raise ValueError(f"Invalid execute mode: {execute!r}")
         if dtype not in _SUPPORTED_DTYPES:
             raise ValueError(f"Invalid dtype {dtype!r}; expected one of "
@@ -347,6 +349,9 @@ def __init__(
             )
             target_device = "cpu"
 
+        if execute == "cutensornet" and "cuda" not in target_device:
+            raise ValueError("execute='cutensornet' requires a CUDA device.")
+
         # Build the topology directly so NMOptimizer never selects the
         # TensorNetworkDecoder cuTensorNet contractor during setup.
         qec.Decoder.__init__(self, H)
@@ -397,7 +402,9 @@ def __init__(
         self.slicing_batch = tuple()
         self.slicing_single = tuple()
 
-        self._set_contractor("oe_torch_compiled", target_device, "torch", dtype)
+        contractor = ("cutensornet"
+                      if execute == "cutensornet" else "oe_torch_compiled")
+        self._set_contractor(contractor, target_device, "torch", dtype)
         self.init_noise_model(
             factorized_noise_model(self.error_inds, noise_model),
             contract=False,
@@ -718,6 +725,7 @@ def _build_reduced_tn_state(self) -> None:
             })
 
         self._batched_einsum_groups = batched_groups
+        self._reduced_eq = reduced_eq
         self._reduced_static_positions = reduced_static
         self._reduced_syndrome_positions = reduced_syndrome
         self._reduced_recipe_positions = reduced_recipe_positions
@@ -735,6 +743,7 @@ def _build_predict_reduced(self):
             "opt_einsum": self._build_predict_reduced_opt_einsum,
             "unrolled": self._build_predict_reduced_unrolled,
             "codegen": self._build_predict_reduced_codegen,
+            "cutensornet": self._build_predict_reduced_cutensornet,
         }
         return builders[self._execute_mode]()
 
@@ -775,6 +784,34 @@ def _predict(noise_probs: torch.Tensor,
 
         return _predict
 
+    def _build_predict_reduced_cutensornet(self):
+        static_positions = self._reduced_static_positions
+        syndrome_positions = self._reduced_syndrome_positions
+        recipe_positions = self._reduced_recipe_positions
+        eq = self._reduced_eq
+        n = self._reduced_n_tensors
+        contractor = self.contractor_config.contractor
+        device_id = self.contractor_config.device_id
+
+        def _predict(noise_probs: torch.Tensor,
+                     syndrome_tuple: tuple[torch.Tensor, ...]) -> torch.Tensor:
+            arrays: list[torch.Tensor] = [None] * n  # type: ignore
+            for pos, arr in static_positions.items():
+                arrays[pos] = arr
+            for pos, syndrome_idx in syndrome_positions:
+                arrays[pos] = syndrome_tuple[syndrome_idx]
+            for pos, arr in zip(recipe_positions,
+                                self._precontract_reduced_noise(noise_probs)):
+                arrays[pos] = arr
+            out = contractor(eq,
+                             arrays,
+                             optimize=self.path_batch,
+                             slicing=self.slicing_batch,
+                             device_id=device_id)
+            return _normalize_prediction(out)
+
+        return _predict
+
     def _build_predict_reduced_unrolled(self):
         static_positions = self._reduced_static_positions
         syndrome_positions = self._reduced_syndrome_positions
@@ -852,7 +889,7 @@ def _maybe_torch_compile(self, fn, *, kind: str):
         On any compile failure, warn and fall back to eager.  ``kind``
         is included in the warning to disambiguate predict vs loss.
         """
-        if not self._use_torch_compile:
+        if not self._use_torch_compile or self._execute_mode == "cutensornet":
             return fn
         try:
             kwargs = self._torch_compile_kwargs()

libs/qec/python/tests/test_nm_optimizer.py

@@ -203,6 +203,23 @@ def test_invalid_dtype_rejected(device):
                   dtype="float16")
 
 
+def test_cutensornet_execute_requires_cuda():
+    H, logical, priors = _simple_repetition_code()
+    syn, flips = _sample_synthetic_dataset(H,
+                                           logical,
+                                           priors,
+                                           num_shots=4,
+                                           rng=np.random.default_rng(11))
+    with pytest.raises(ValueError, match="requires a CUDA device"):
+        _make_opt(H,
+                  logical,
+                  priors,
+                  syn,
+                  flips,
+                  device="cpu",
+                  execute="cutensornet")
+
+
 # -- forward pass / gradient -------------------------------------------------
 
 
@@ -429,6 +446,53 @@ def test_reduced_path_matches_full_network_reference_static_codegen(device):
                               rtol=1e-7)
 
 
+@pytest.mark.skipif(not _gpu_available(), reason="CUDA not available")
+def test_cutensornet_execute_matches_opt_einsum_and_backpropagates():
+    """cuTN can execute the reduced torch graph with autograd."""
+    pytest.importorskip("cuquantum")
+    rng = np.random.default_rng(31415)
+    H, logical = _nondegenerate_code()
+    init_priors = [0.2, 0.3, 0.4]
+    syn, flips = _sample_synthetic_dataset(H,
+                                           logical, [0.1, 0.15, 0.25],
+                                           num_shots=12,
+                                           rng=rng)
+    ref = _make_opt(H,
+                    logical,
+                    init_priors,
+                    syn,
+                    flips,
+                    device="cuda",
+                    dtype="float64",
+                    execute="opt_einsum")
+    cutn = _make_opt(H,
+                     logical,
+                     init_priors,
+                     syn,
+                     flips,
+                     device="cuda",
+                     dtype="float64",
+                     execute="cutensornet")
+
+    with torch.no_grad():
+        p_ref = ref.decoder_prediction()
+        p_cutn = cutn.decoder_prediction()
+        loss_ref = ref.cross_entropy_loss()
+        loss_cutn = cutn.cross_entropy_loss()
+
+    assert torch.allclose(p_cutn, p_ref, atol=1e-10, rtol=1e-10)
+    assert torch.allclose(loss_cutn, loss_ref, atol=1e-10, rtol=1e-10)
+    assert cutn.contractor_config.contractor_name == "cutensornet"
+
+    cutn.noise_params[0].grad = None
+    loss = cutn.cross_entropy_loss()
+    loss.backward()
+    grad = cutn.noise_params[0].grad
+    assert grad is not None
+    assert torch.isfinite(grad).all()
+    assert torch.any(grad != 0.0)
+
+
 # -- numerical guards --------------------------------------------------------
 
 

scripts/benchmark_nmoptimizer_d5r5_nico_dem.py

@@ -57,7 +57,7 @@
 
 from cudaq_qec import NMOptimizer, make_training_step
 
-EXECUTE_MODES = ("codegen", "unrolled", "opt_einsum")
+EXECUTE_MODES = ("codegen", "unrolled", "opt_einsum", "cutensornet")
 
 
 @dataclass
@@ -198,6 +198,22 @@ def path_opt_cost(info: Any) -> float:
         return float("nan")
 
 
+def einsum_rank(eq: str) -> int:
+    if "->" not in eq:
+        return 0
+    lhs, rhs = eq.split("->", 1)
+    terms = [term for term in lhs.split(",") if term]
+    terms.append(rhs)
+    return max((len(term) for term in terms), default=0)
+
+
+def path_max_einsum_rank(info: Any) -> int:
+    ranks = [
+        einsum_rank(step[2]) for step in getattr(info, "contraction_list", ())
+    ]
+    return max(ranks, default=0)
+
+
 def fmt_float(value: float) -> str:
     if math.isnan(value):
         return "n/a"
@@ -252,25 +268,32 @@ def construct_optimizer(problem: Problem, syn: np.ndarray, flips: np.ndarray,
                        compile=args.compile)
 
 
-def run_case(problem: Problem, batch_size: int, execute: str,
-             args: argparse.Namespace) -> dict[str, str]:
+def run_case(problem: Problem,
+             batch_size: int,
+             execute: str,
+             args: argparse.Namespace,
+             fixed_path: Any | None = None,
+             fixed_path_source: str = "") -> tuple[dict[str, str], Any | None]:
     label = f"batch={batch_size}, execute={execute}"
     row = {
         "batch": str(batch_size),
         "construct_batch": str(batch_size),
         "execute": execute,
+        "path_source": fixed_path_source if fixed_path is not None else "self",
         "status": "FAIL",
         "construct_s": "n/a",
         "largest_intermediate": "n/a",
         "intermediate_gib": "n/a",
         "opt_cost": "n/a",
+        "max_einsum_rank": "n/a",
         "step_s": "n/a",
         "peak_mem_mib": "n/a",
         "error": "",
     }
     init_priors = make_initial_priors(problem.priors, args.init)
     syn, flips = sample_problem(problem, batch_size, args.seed + batch_size)
 
+    selected_path = None
     try:
         clear_cuda(args.device)
         with timed(f"Construct NMOptimizer ({label})"):
@@ -310,6 +333,10 @@ def _progress(done: int, total: int, elapsed: float) -> None:
                 opt = construct_optimizer(problem, syn, flips, init_priors,
                                           args, execute)
                 step_fn = None
+            if fixed_path is not None:
+                print(f"Reusing reduced path from {fixed_path_source}",
+                      flush=True)
+                opt.optimize_path(fixed_path)
             row["construct_s"] = f"{time.perf_counter() - t0:.3f}"
 
         construct_batch = int(getattr(opt, "_batch_size", batch_size))
@@ -319,10 +346,12 @@ def _progress(done: int, total: int, elapsed: float) -> None:
             run_label += f", construct_batch={construct_batch}"
 
         info = getattr(opt, "_reduced_info", None)
+        selected_path = getattr(opt, "path_batch", None)
         largest = path_largest_intermediate(info)
         cost = path_opt_cost(info)
         row["largest_intermediate"] = fmt_float(largest)
         row["opt_cost"] = fmt_float(cost)
+        row["max_einsum_rank"] = str(path_max_einsum_rank(info))
         if not math.isnan(largest):
             gib = largest * element_size(args.dtype) / 2**30
             row["intermediate_gib"] = f"{gib:.3f}"
@@ -343,15 +372,15 @@ def _progress(done: int, total: int, elapsed: float) -> None:
 
         row["peak_mem_mib"] = peak_mem_mib(args.device)
         row["status"] = "OK"
-        return row
+        return row, selected_path
     except Exception as exc:  # intentionally broad: this is an OOM probe.
         row["error"] = repr(exc)
         with contextlib.suppress(Exception):
             row["peak_mem_mib"] = peak_mem_mib(args.device)
         print(f"FAILED CASE {label}: {exc!r}", flush=True)
         if args.stop_on_failure:
             raise
-        return row
+        return row, selected_path
     finally:
         with contextlib.suppress(Exception):
             del opt  # type: ignore[name-defined]
@@ -396,7 +425,7 @@ def make_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--execute",
         default="all",
-        help="all, codegen, unrolled, opt_einsum, or comma list")
+        help="all, codegen, unrolled, opt_einsum, cutensornet, or comma list")
     parser.add_argument("--device", default="cuda:0")
     parser.add_argument("--dtype",
                         choices=["float32", "float64"],
@@ -423,6 +452,11 @@ def make_parser() -> argparse.ArgumentParser:
                         type=int,
                         default=10,
                         help="Print batch-slice progress every N chunks.")
+    parser.add_argument(
+        "--reuse-first-path",
+        action="store_true",
+        help=("For each requested batch size, reuse the first execute mode's "
+              "reduced contraction path for later execute modes."))
     parser.add_argument("--lr",
                         type=float,
                         default=1e-2,
@@ -456,6 +490,7 @@ def main() -> int:
         f"d={args.distance}, r={args.rounds}, device={args.device}, "
         f"dtype={args.dtype}, execute={execute_modes}, "
         f"batch_sizes={args.batch_sizes}, run_step={args.run_step}, "
+        f"reuse_first_path={args.reuse_first_path}, "
         f"batch_slice_size={args.batch_slice_size}",
         flush=True)
 
@@ -477,14 +512,26 @@ def main() -> int:
 
     rows = []
     for batch_size in args.batch_sizes:
+        fixed_path = None
+        fixed_path_source = ""
         for execute in execute_modes:
-            rows.append(run_case(problem, batch_size, execute, args))
+            row, selected_path = run_case(problem,
+                                          batch_size,
+                                          execute,
+                                          args,
+                                          fixed_path=fixed_path,
+                                          fixed_path_source=fixed_path_source)
+            rows.append(row)
+            if (args.reuse_first_path and fixed_path is None and
+                    row["status"] == "OK" and selected_path is not None):
+                fixed_path = selected_path
+                fixed_path_source = f"execute={execute}"
 
     print("\n## NMOptimizer d=5/r=5 Large-DEM Memory Matrix")
     headers = [
-        "batch", "construct_batch", "execute", "status", "construct_s",
-        "largest_intermediate", "intermediate_gib", "opt_cost", "step_s",
-        "peak_mem_mib", "error"
+        "batch", "construct_batch", "execute", "path_source", "status",
+        "construct_s", "largest_intermediate", "intermediate_gib", "opt_cost",
+        "max_einsum_rank", "step_s", "peak_mem_mib", "error"
     ]
     print("| " + " | ".join(headers) + " |")
     print("| " + " | ".join(["---"] * len(headers)) + " |")