✨ Extend QCO DD construction/simulation beyond static unitaries

CHANGELOG.md

@@ -24,8 +24,11 @@ releases may include breaking changes.
   [#1755], [#1787], [#1815], [#1823], [#1830], [#1886], [#1933], [#1978],
   [#1979]) ([**@burgholzer**], [**@denialhaag**], [**@simon1hofmann**],
   [**@li-mingbao**], [**@DRovara**], [**@MatthiasReumann**])
-- ✨ Add decision diagram-based construction and simulation of static unitary
-  QCO functions ([#1915]) ([**@simon1hofmann**])
+- ✨ Add decision diagram-based construction and simulation of QCO functions,
+  including static unitaries, mid-circuit `measure`/`reset`, concrete
+  `if`/`index_switch`/`scf.for`/`func.call`, richer classical SSA, dense `k>3`
+  wire embedding, multi-shot `sample` / `sampleWithClassics`, `memref` classical
+  registers, and Python wrappers ([#1915], [#1973]) ([**@simon1hofmann**])
 - ✨ Add a `fuse-two-qubit-unitary-runs` pass for fusing compile-time two-qubit
   unitary windows via Weyl/KAK resynthesis ([#1865], [#1961])
   ([**@simon1hofmann**], [**@burgholzer**])
@@ -701,6 +704,7 @@ changelogs._
 [#1978]: https://github.com/munich-quantum-toolkit/core/pull/1978
 [#1976]: https://github.com/munich-quantum-toolkit/core/pull/1976
 [#1975]: https://github.com/munich-quantum-toolkit/core/pull/1975
+[#1973]: https://github.com/munich-quantum-toolkit/core/pull/1973
 [#1972]: https://github.com/munich-quantum-toolkit/core/pull/1972
 [#1967]: https://github.com/munich-quantum-toolkit/core/pull/1967
 [#1965]: https://github.com/munich-quantum-toolkit/core/pull/1965

bindings/mlir/CMakeLists.txt

@@ -23,7 +23,8 @@ if(NOT TARGET ${TARGET_NAME})
     .
     LINK_LIBS
     MQTCompilerPipeline
-    MQT::CoreIR)
+    MQT::CoreIR
+    MLIRQCODDFunctionality)
 
   # install the Python stub file in editable mode for better IDE support
   if(SKBUILD_STATE STREQUAL "editable")

bindings/mlir/register_mlir.cpp

@@ -8,11 +8,20 @@
  * Licensed under the MIT License
  */
 
+#include "dd/Package.hpp"
 #include "ir/QuantumComputation.hpp"
 #include "mlir/Compiler/Programs.h"
+#include "mlir/Dialect/QCO/Utils/DDFunctionality.h"
 
+#include <llvm/Support/raw_ostream.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/IR/Diagnostics.h>
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/Support/LogicalResult.h>
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/filesystem.h>  // NOLINT(misc-include-cleaner)
+#include <nanobind/stl/map.h>         // NOLINT(misc-include-cleaner)
+#include <nanobind/stl/optional.h>    // NOLINT(misc-include-cleaner)
 #include <nanobind/stl/pair.h>        // NOLINT(misc-include-cleaner)
 #include <nanobind/stl/string.h>      // NOLINT(misc-include-cleaner)
 #include <nanobind/stl/string_view.h> // NOLINT(misc-include-cleaner)
@@ -21,8 +30,10 @@
 
 #include <cctype>
 #include <cstddef>
+#include <cstdint>
 #include <filesystem>
 #include <optional>
+#include <random>
 #include <span>
 #include <stdexcept>
 #include <string>
@@ -95,6 +106,58 @@ void requireValid(const mlir::Program& program) {
   }
 }
 
+[[nodiscard]] mlir::func::FuncOp entryFunc(const mlir::QCOProgram& program) {
+  requireValid(program);
+  auto func = program.entryFunc();
+  if (!func) {
+    throw nb::value_error("QCO program has no func.func entry point");
+  }
+  return *func;
+}
+
+[[nodiscard]] std::mt19937_64 makeRng(const std::uint64_t seed) {
+  if (seed == 0) {
+    std::random_device rd;
+    return std::mt19937_64(rd());
+  }
+  return std::mt19937_64(seed);
+}
+
+[[nodiscard]] std::mt19937_64
+makeRng(const std::optional<std::uint64_t>& seed) {
+  if (!seed.has_value()) {
+    std::random_device rd;
+    return std::mt19937_64(rd());
+  }
+  return makeRng(*seed);
+}
+
+/// Run @p fn under a diagnostic handler and raise `ValueError` on failure,
+/// appending any emitted MLIR diagnostics to @p message.
+template <typename Fn>
+[[nodiscard]] auto takeFailureOr(mlir::MLIRContext* context,
+                                 const char* message, Fn&& fn) {
+  std::string diagnostics;
+  const mlir::ScopedDiagnosticHandler handler(
+      context, [&](mlir::Diagnostic& diag) {
+        if (!diagnostics.empty()) {
+          diagnostics.push_back('\n');
+        }
+        llvm::raw_string_ostream os(diagnostics);
+        os << diag;
+        return mlir::success();
+      });
+  auto result = std::forward<Fn>(fn)();
+  if (mlir::failed(result)) {
+    std::string full = message;
+    if (!diagnostics.empty()) {
+      full.append(": ").append(diagnostics);
+    }
+    throw nb::value_error(full.c_str());
+  }
+  return *std::move(result);
+}
+
 template <class ProgramType>
 [[nodiscard]] ProgramType copiedOrConsumed(ProgramType& program,
                                            const bool copy) {
@@ -513,6 +576,132 @@ LLVM bitcode.)pb");
            &BooleanMemberAdapter<&mlir::QIRProgram::writeBitcode>::call,
            "path"_a, "Write this program as LLVM bitcode.");
 
+  nb::module_::import_("mqt.core.dd");
+
+  nb::class_<mlir::qco::SampleResult>(m, "SampleResult",
+                                      "Histograms from QCO DD sampling.")
+      .def_ro("shots", &mlir::qco::SampleResult::shots,
+              "Final computational-basis outcome histogram.")
+      .def_ro("classical", &mlir::qco::SampleResult::classical,
+              "Mid-circuit measure-bit histogram (encounter order).");
+
+  m.def(
+      "build_functionality",
+      [](const mlir::QCOProgram& program, dd::Package& ddPackage) {
+        auto func = entryFunc(program);
+        return takeFailureOr(
+            func.getContext(),
+            "cannot build DD functionality for this QCO program",
+            [&] { return mlir::qco::buildFunctionality(func, ddPackage); });
+      },
+      "program"_a, "dd_package"_a,
+      // Keep the DD package alive while the returned matrix DD is alive
+      // (arg index 2; free-function equivalent of method keep_alive<0, 1>).
+      nb::keep_alive<0, 2>(),
+      R"pb(Build a matrix DD for a static unitary QCO program.
+
+Args:
+    program: A QCO program whose entry ``func.func`` is simulated.
+    dd_package: DD package with enough qubits for the program.
+
+Returns:
+    Matrix DD of the program functionality.
+
+Raises:
+    ValueError: When the program is unsupported for functionality construction.)pb");
+
+  m.def(
+      "simulate",
+      [](const mlir::QCOProgram& program, const dd::VectorDD& initialState,
+         dd::Package& ddPackage, const std::optional<std::uint64_t> seed) {
+        auto func = entryFunc(program);
+        if (!seed.has_value()) {
+          return takeFailureOr(
+              func.getContext(), "cannot simulate this QCO program", [&] {
+                return mlir::qco::simulate(func, initialState, ddPackage);
+              });
+        }
+        auto rng = makeRng(*seed);
+        return takeFailureOr(
+            func.getContext(), "cannot simulate this QCO program", [&] {
+              return mlir::qco::simulate(func, initialState, ddPackage, rng);
+            });
+      },
+      "program"_a, "initial_state"_a, "dd_package"_a, "seed"_a = nb::none(),
+      // Keep the DD package alive while the returned vector DD is alive.
+      nb::keep_alive<0, 3>(),
+      R"pb(Simulate a QCO program on a DD state.
+
+Args:
+    program: A QCO program whose entry ``func.func`` is simulated.
+    initial_state: Input state DD (one reference is consumed).
+    dd_package: DD package with enough qubits for the program.
+    seed: If ``None``, rejects mid-circuit measure/reset. Otherwise seeds the
+        RNG used for collapsing measurements and resets (``0`` = nondeterministic).
+
+Returns:
+    Output state DD.
+
+Raises:
+    ValueError: When the program is unsupported for simulation.)pb");
+
+  // Sampling uses a caller-provided ``dd::Package`` for the call only; the
+  // binding does not share that package across threads. Release the GIL only
+  // around the C++ sample (not entryFunc / exception translation).
+  m.def(
+      "sample",
+      [](const mlir::QCOProgram& program, dd::Package& ddPackage,
+         const size_t shots, const std::optional<uint64_t> seed) {
+        auto func = entryFunc(program);
+        auto rng = makeRng(seed);
+        return takeFailureOr(
+            func.getContext(), "cannot sample this QCO program", [&] {
+              const nb::gil_scoped_release release;
+              return mlir::qco::sample(func, ddPackage, shots, rng);
+            });
+      },
+      "program"_a, "dd_package"_a, "shots"_a = 1024U, "seed"_a = nb::none(),
+      R"pb(Sample final computational-basis outcomes from a QCO program.
+
+Args:
+    program: A QCO program whose entry ``func.func`` is sampled.
+    dd_package: DD package with enough qubits for the program.
+    shots: Number of shots (default 1024).
+    seed: RNG seed. ``None`` (default) or ``0`` selects nondeterministic seeding.
+
+Returns:
+    Histogram of final ``measureAll`` bitstrings.
+
+Raises:
+    ValueError: When the program is unsupported for sampling.)pb");
+
+  m.def(
+      "sample_with_classics",
+      [](const mlir::QCOProgram& program, dd::Package& ddPackage,
+         const size_t shots, const std::optional<uint64_t> seed) {
+        auto func = entryFunc(program);
+        auto rng = makeRng(seed);
+        return takeFailureOr(
+            func.getContext(), "cannot sample this QCO program", [&] {
+              const nb::gil_scoped_release release;
+              return mlir::qco::sampleWithClassics(func, ddPackage, shots, rng);
+            });
+      },
+      "program"_a, "dd_package"_a, "shots"_a = 1024U, "seed"_a = nb::none(),
+      R"pb(Sample final and mid-circuit classical outcomes from a QCO program.
+
+Args:
+    program: A QCO program whose entry ``func.func`` is sampled.
+    dd_package: DD package with enough qubits for the program.
+    shots: Number of shots (default 1024).
+    seed: RNG seed. ``None`` (default) or ``0`` selects nondeterministic seeding.
+
+Returns:
+    A :class:`SampleResult` with ``shots`` and ``classical`` histograms.
+
+Raises:
+    ValueError: When the program is unsupported for sampling.)pb");
+
   m.def("compile_program", &compileProgram, "program"_a, nb::kw_only(),
         "output"_a = mlir::ProgramFormat::QC, "inplace"_a = false,
         "qco_pipeline"_a = "mqt-qco-default", "enable_timing"_a = false,

mlir/include/mlir/Compiler/Programs.h

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/IR/BuiltinOps.h>
 #include <mlir/IR/MLIRContext.h>
 #include <mlir/IR/OwningOpRef.h>
@@ -207,6 +208,9 @@ class QCOProgram final : public Program {
 
   /// Consume this program and convert it to `jeff` MLIR.
   [[nodiscard]] std::optional<JeffProgram> intoJeff() &&;
+
+  /// Return the entry `func.func` (`main` if present, else the first function).
+  [[nodiscard]] std::optional<func::FuncOp> entryFunc() const;
 };
 
 /**