✨ Extend QCO DD construction/simulation beyond static unitaries

CHANGELOG.md

@@ -20,8 +20,11 @@ releases may include breaking changes.
   [#1755], [#1787], [#1815], [#1823], [#1830], [#1886], [#1933])
   ([**@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**])
@@ -692,6 +695,7 @@ changelogs._
 
 <!-- PR links -->
 
+[#1973]: https://github.com/munich-quantum-toolkit/core/pull/1973
 [#1967]: https://github.com/munich-quantum-toolkit/core/pull/1967
 [#1965]: https://github.com/munich-quantum-toolkit/core/pull/1965
 [#1961]: https://github.com/munich-quantum-toolkit/core/pull/1961

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,17 @@
  * 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 <mlir/Dialect/Func/IR/FuncOps.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 +27,10 @@
 
 #include <cctype>
 #include <cstddef>
+#include <cstdint>
 #include <filesystem>
 #include <optional>
+#include <random>
 #include <span>
 #include <stdexcept>
 #include <string>
@@ -95,6 +103,32 @@ 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 to simulate");
+  }
+  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);
+}
+
+template <typename T>
+[[nodiscard]] T takeFailureOr(mlir::FailureOr<T>&& result,
+                              const char* message) {
+  if (mlir::failed(result)) {
+    throw nb::value_error(message);
+  }
+  return *std::move(result);
+}
+
 template <class ProgramType>
 [[nodiscard]] ProgramType copiedOrConsumed(ProgramType& program,
                                            const bool copy) {
@@ -513,6 +547,114 @@ 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) {
+        return takeFailureOr(
+            mlir::qco::buildFunctionality(entryFunc(program), ddPackage),
+            "cannot build DD functionality for this QCO program");
+      },
+      "program"_a, "dd_package"_a,
+      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(
+              mlir::qco::simulate(func, initialState, ddPackage),
+              "cannot simulate this QCO program");
+        }
+        auto rng = makeRng(*seed);
+        return takeFailureOr(
+            mlir::qco::simulate(func, initialState, ddPackage, rng),
+            "cannot simulate this QCO program");
+      },
+      "program"_a, "initial_state"_a, "dd_package"_a, "seed"_a = nb::none(),
+      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");
+
+  m.def(
+      "sample",
+      [](const mlir::QCOProgram& program, dd::Package& ddPackage,
+         const std::size_t shots, const std::uint64_t seed) {
+        auto rng = makeRng(seed);
+        return takeFailureOr(
+            mlir::qco::sample(entryFunc(program), ddPackage, shots, rng),
+            "cannot sample this QCO program");
+      },
+      "program"_a, "dd_package"_a, "shots"_a = 1024U, "seed"_a = 0U,
+      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 (``0`` = nondeterministic).
+
+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 std::size_t shots, const std::uint64_t seed) {
+        auto rng = makeRng(seed);
+        return takeFailureOr(mlir::qco::sampleWithClassics(
+                                 entryFunc(program), ddPackage, shots, rng),
+                             "cannot sample this QCO program");
+      },
+      "program"_a, "dd_package"_a, "shots"_a = 1024U, "seed"_a = 0U,
+      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 (``0`` = nondeterministic).
+
+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;
 };
 
 /**

mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h

@@ -15,6 +15,11 @@
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Support/LogicalResult.h>
 
+#include <cstddef>
+#include <map>
+#include <random>
+#include <string>
+
 namespace mlir::qco {
 
 /**
@@ -31,15 +36,15 @@ namespace mlir::qco {
  * - `ctrl` with a sole standard-gate body (same sparse path)
  * - Other `UnitaryOpInterface` ops with a compile-time known matrix (`inv`,
  *   compound `ctrl`, ...), including `gphase` and `barrier`
- * - Skips: `static`, `sink`, `arith.constant`; `func.return` accepts qubit
- *   results only in canonical wire order
+ * - `qco.static` establishes the wire map (or qubit-typed `func` args if none);
+ *   `sink` is ignored; `arith.constant` is ignored for matrix construction;
+ *   `func.return` accepts qubit results only in canonical wire order
  *
  * Known one-, two-, and three-qubit matrices are constructed directly as DD
- * gates. The dense matrix fallback accepts full-width unitaries on wires
- * `0..n-1` and rewrites the QCO/MSB-first basis into the DD package's
- * LSB-first indexing. Only this full-width fallback is limited to 12 qubits
- * (dense `2^n × 2^n` storage). Measurements, resets, symbolic parameters, and
- * control-flow ops are not supported.
+ * gates. Larger compile-time unitaries (including partial wire subsets) use a
+ * dense embed into the full register, rewritten from QCO/MSB to DD/LSB, limited
+ * to 12 qubits (`2^n × 2^n` storage). Measurements, resets, symbolic
+ * parameters, and control-flow ops are not supported.
  *
  * @param func The QCO function to construct the functionality for
  * @param dd The DD package to use (must hold at least the function's qubits)
@@ -48,13 +53,17 @@ namespace mlir::qco {
 FailureOr<dd::MatrixDD> buildFunctionality(func::FuncOp func, dd::Package& dd);
 
 /**
- * @brief Simulate a static unitary QCO `func.func` on a given input state.
+ * @brief Simulate a QCO `func.func` on a given input state without stochastic
+ * collapse.
  *
- * @details Same supported op set and limitations as @ref buildFunctionality.
- * Mirrors @ref dd::simulate: sequentially applies unitaries to @p in via
- * decision-diagram multiplication. Only purely quantum, measurement-free
- * programs are supported. Consumes one reference to @p in regardless of
- * whether simulation succeeds or fails.
+ * @details Same supported unitary op set as @ref buildFunctionality, plus
+ * concrete classical control-flow (`qco.if` / `qco.index_switch` with
+ * compile-time or previously recorded classical selectors) and static 1-D
+ * `memref<?xi1>` classical registers (`alloc`/`store`/`load`/`dealloc`).
+ * Mid-circuit `measure` / `reset` require the RNG overload below. Only
+ * qubit-typed linear values are supported (no qtensors). Nested regions are
+ * walked; loops and multi-block function bodies remain unsupported. Consumes
+ * one reference to @p in regardless of whether simulation succeeds or fails.
  *
  * @param func The QCO function to simulate
  * @param in The input state, represented as a vector DD; one reference is
@@ -66,4 +75,95 @@ FailureOr<dd::MatrixDD> buildFunctionality(func::FuncOp func, dd::Package& dd);
 FailureOr<dd::VectorDD> simulate(func::FuncOp func, const dd::VectorDD& in,
                                  dd::Package& dd);
 
+/**
+ * @brief Simulate a QCO `func.func` that may contain measurements, resets, and
+ * concrete control-flow.
+ *
+ * @details Supports the unitary op set of @ref buildFunctionality, plus
+ * `qco.measure` / `qco.reset` (collapsing via @p rng) and `qco.if` /
+ * `qco.index_switch` when the branch selector is a concrete classical SSA value
+ * (`arith.constant` `i1`/`index`, a prior measurement, `arith.index_castui`,
+ * `arith.extui`/`trunci` between `i1` and `index`, `arith.cmpi`,
+ * `arith.select`, `arith.addi`/`subi`/`muli`,
+ * `andi`/`ori`/`xori`/`shli`/`shrui` on those values). Classical registers as
+ * static 1-D `memref<?xi1>` with `memref.alloc` / `store` / `load` /
+ * `dealloc` are supported. Deterministic
+ * control-flow without measure/reset also works on the non-RNG overload. Only
+ * qubit-typed linear values are supported (no qtensors). Nested regions are
+ * walked; `scf.for` with concrete positive step and at most 10000 trips and
+ * non-recursive single-block `func.call` are supported; `scf.while` and
+ * multi-block function bodies remain unsupported. Consumes one reference to
+ * @p in regardless of whether simulation succeeds or fails.
+ *
+ * @param func The QCO function to simulate
+ * @param in The input state; one reference is consumed
+ * @param dd The DD package to use
+ * @param rng RNG used for collapsing measurements and resets
+ * @return The output statevector DD on success, or failure for unsupported
+ *         programs
+ */
+FailureOr<dd::VectorDD> simulate(func::FuncOp func, const dd::VectorDD& in,
+                                 dd::Package& dd, std::mt19937_64& rng);
+
+/**
+ * @brief Sample measurement outcomes from a QCO `func.func`.
+ *
+ * @details Starts from the all-zero state and draws @p shots bitstrings via
+ * `Package::measureAll` (qubit `n-1` … `0`, same as @ref dd::sample). Programs
+ * without `measure` / `reset` are simulated once and sampled without
+ * collapsing (including deterministic control-flow). Programs with mid-circuit
+ * `measure` / `reset` are re-simulated per shot with @p rng. Histograms are
+ * final computational-basis bitstrings, not classical mid-circuit records.
+ *
+ * @param func The QCO function to sample
+ * @param dd The DD package to use
+ * @param shots Number of shots
+ * @param rng RNG for collapsing measurements and non-collapsing sampling
+ * @return Histogram of outcome strings on success, or failure for unsupported
+ *         programs
+ */
+FailureOr<std::map<std::string, std::size_t>> sample(func::FuncOp func,
+                                                     dd::Package& dd,
+                                                     std::size_t shots,
+                                                     std::mt19937_64& rng);
+
+/**
+ * @brief Sample measurement outcomes from a QCO `func.func` on a given input.
+ *
+ * @details Same as the zero-state overload, but starts from @p in. Consumes one
+ * reference to @p in (the static path keeps that state for all shots; the
+ * dynamic path clones per shot).
+ *
+ * @param func The QCO function to sample
+ * @param in Input state; one reference is consumed
+ * @param dd The DD package to use
+ * @param shots Number of shots
+ * @param rng RNG for collapsing measurements and non-collapsing sampling
+ * @return Histogram of outcome strings on success, or failure for unsupported
+ *         programs
+ */
+FailureOr<std::map<std::string, std::size_t>>
+sample(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd,
+       std::size_t shots, std::mt19937_64& rng);
+
+/**
+ * @brief Sample final and mid-circuit classical outcomes from a QCO
+ * `func.func`.
+ *
+ * @details Like @ref sample, but also histograms collapsing mid-circuit
+ * measurement bits in encounter order into @c SampleResult::classical.
+ * Programs without mid-circuit measures leave @c classical empty.
+ */
+struct SampleResult {
+  std::map<std::string, size_t> shots;
+  std::map<std::string, size_t> classical;
+};
+
+FailureOr<SampleResult> sampleWithClassics(func::FuncOp func, dd::Package& dd,
+                                           size_t shots, std::mt19937_64& rng);
+FailureOr<SampleResult> sampleWithClassics(func::FuncOp func,
+                                           const dd::VectorDD& in,
+                                           dd::Package& dd, size_t shots,
+                                           std::mt19937_64& rng);
+
 } // namespace mlir::qco