munich-quantum-toolkit · simon1hofmann · Aug 13, 2026 · Aug 13, 2026 · Aug 14, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -59,10 +59,10 @@ releases may include breaking changes.
   including static unitaries, mid-circuit `measure`/`reset`, concrete structured
   and multi-block control flow, extended classical arithmetic and `memref`
   values across calls, symbolic argument bindings, dynamic qubit and qtensor
-  allocation, qtensor ownership and DD-native deallocation, direct local-matrix
-  DD embedding, multi-shot sampling, and Python wrappers ([#1915], [#1973],
-  [#2077], [#2078], [#2079])
-  ([**@simon1hofmann**])
+  allocation, qtensor ownership and DD-native deallocation, mixed-state density
+  simulation and sampling with partial trace, direct local-matrix DD embedding,
+  multi-shot sampling, and Python wrappers ([#1915], [#1973], [#2077], [#2078],
+  [#2079], [#2080]) ([**@simon1hofmann**])
 - ✨ Add target-independent two-qubit gate fusion, target-native post-routing
   synthesis, and operation-capability and static-site conformance ([#1865],
   [#1961], [#1998]) ([**@simon1hofmann**], [**@burgholzer**])
@@ -757,6 +757,7 @@ for previous changelogs._
 
 <!-- PR links -->
 
+[#2080]: https://github.com/munich-quantum-toolkit/core/pull/2080
 [#2079]: https://github.com/munich-quantum-toolkit/core/pull/2079
 [#2078]: https://github.com/munich-quantum-toolkit/core/pull/2078
 [#2077]: https://github.com/munich-quantum-toolkit/core/pull/2077

diff --git a/bindings/mlir/register_mlir.cpp b/bindings/mlir/register_mlir.cpp
@@ -14,13 +14,15 @@
 #include "mlir/Compiler/FoMaCAdapter.h"
 #include "mlir/Compiler/Programs.h"
 #include "mlir/Compiler/Target.h"
-#include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/IR/QCODialect.h"
 #include "mlir/Dialect/QCO/Utils/DDFunctionality.h"
 #include "qiskit/Qiskit.h"
 
+#include <llvm/Support/Casting.h>
 #include <llvm/Support/Error.h>
 #include <llvm/Support/raw_ostream.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/IR/Attributes.h>
 #include <mlir/IR/BuiltinAttributes.h>
 #include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/Diagnostics.h>
@@ -151,26 +153,32 @@ makeQCODDBindings(mlir::func::FuncOp func,
       throw nb::value_error("QCO DD binding argument index is out of range");
     }
 
-    mlir::Value argument = func.getArgument(static_cast<unsigned>(index));
+    const mlir::Value argument = func.getArgument(static_cast<unsigned>(index));
     mlir::Type type = argument.getType();
     mlir::Attribute attribute;
     if (type.isInteger(1)) {
       if (const auto* value = std::get_if<bool>(&binding)) {
         attribute = mlir::BoolAttr::get(func.getContext(), *value);
       }
-    } else if (mlir::isa<mlir::IndexType, mlir::IntegerType>(type)) {
+    } else if (llvm::isa<mlir::IndexType, mlir::IntegerType>(type)) {
       if (const auto* value = std::get_if<int64_t>(&binding)) {
-        attribute = mlir::IntegerAttr::get(type, *value);
+        const auto width = llvm::isa<mlir::IndexType>(type)
+                               ? 64U
+                               : llvm::cast<mlir::IntegerType>(type).getWidth();
+        if (width >= 64U || llvm::APInt(64, static_cast<uint64_t>(*value), true)
+                                .isSignedIntN(width)) {
+          attribute = mlir::IntegerAttr::get(type, *value);
+        }
       }
-    } else if (const auto floatType = mlir::dyn_cast<mlir::FloatType>(type)) {
+    } else if (type.isIntOrFloat()) {
       if (const auto* value = std::get_if<double>(&binding)) {
-        attribute = mlir::FloatAttr::get(floatType, *value);
+        attribute = mlir::FloatAttr::get(type, *value);
       }
     } else if (const auto tensorType =
-                   mlir::dyn_cast<mlir::RankedTensorType>(type);
+                   llvm::dyn_cast<mlir::RankedTensorType>(type);
                tensorType && tensorType.getRank() == 1 &&
                tensorType.isDynamicDim(0) &&
-               mlir::isa<mlir::qco::QubitType>(tensorType.getElementType())) {
+               llvm::isa<mlir::qco::QubitType>(tensorType.getElementType())) {
       if (const auto* value = std::get_if<int64_t>(&binding);
           value != nullptr && *value >= 0) {
         attribute = mlir::IntegerAttr::get(
@@ -187,20 +195,12 @@ makeQCODDBindings(mlir::func::FuncOp func,
   return bindings;
 }
 
-[[nodiscard]] std::mt19937_64 makeRng(const 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<uint64_t>& seed) {
-  if (!seed.has_value()) {
+  if (!seed.has_value() || *seed == 0) {
     std::random_device rd;
     return std::mt19937_64(rd());
   }
-  return makeRng(*seed);
+  return std::mt19937_64(*seed);
 }
 
 /// Run @p fn under a diagnostic handler and raise `ValueError` on failure,
@@ -1028,7 +1028,7 @@ LLVM bitcode.)pb");
                                      func, initialState, ddPackage, bindings);
                                });
         }
-        auto rng = makeRng(*seed);
+        auto rng = makeRng(seed);
         return takeFailureOr(
             func.getContext(), "cannot simulate this QCO program", [&] {
               return mlir::qco::simulate(func, initialState, ddPackage, rng,
@@ -1055,6 +1055,54 @@ LLVM bitcode.)pb");
 Raises:
     ValueError: When the program is unsupported for simulation.)pb");
 
+  m.def(
+      "make_density_matrix",
+      [](const dd::VectorDD& state, const size_t numQubits,
+         dd::Package& ddPackage) {
+        if (numQubits > ddPackage.qubits()) {
+          throw nb::value_error(
+              "num_qubits exceeds the capacity of the DD package");
+        }
+        return mlir::qco::makeDensityMatrix(state, numQubits, ddPackage);
+      },
+      "state"_a, "num_qubits"_a, "dd_package"_a, nb::keep_alive<0, 3>(),
+      R"pb(Construct ``|psi><psi|`` from a pure DD state.
+
+The input vector reference remains owned by the caller. The returned matrix DD
+is referenced and must be released with ``DDPackage.dec_ref_mat``.
+
+Raises:
+    ValueError: When ``num_qubits`` exceeds the DD package capacity.)pb");
+
+  m.def(
+      "simulate_density",
+      [](const mlir::QCOProgram& program, const dd::MatrixDD& initialState,
+         dd::Package& ddPackage, const std::optional<uint64_t> seed,
+         const QCODDBindingMap& pythonBindings) {
+        auto func = entryFunc(program);
+        auto bindings = makeQCODDBindings(func, pythonBindings);
+        if (!seed.has_value()) {
+          return takeFailureOr(func.getContext(),
+                               "cannot density-simulate this QCO program", [&] {
+                                 return mlir::qco::simulateDensity(
+                                     func, initialState, ddPackage, bindings);
+                               });
+        }
+        auto rng = makeRng(seed);
+        return takeFailureOr(
+            func.getContext(), "cannot density-simulate this QCO program", [&] {
+              return mlir::qco::simulateDensity(func, initialState, ddPackage,
+                                                rng, bindings);
+            });
+      },
+      "program"_a, "initial_state"_a, "dd_package"_a, "seed"_a = nb::none(),
+      nb::kw_only(), "bindings"_a = QCODDBindingMap{}, nb::keep_alive<0, 3>(),
+      R"pb(Simulate a QCO program on a density-matrix DD.
+
+Unitary gates evolve ``rho`` as ``U rho U*`` and deallocation performs a
+partial trace, including for entangled qubits. The input matrix reference is
+consumed. Supply ``seed`` for programs containing measurement or reset.)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).
@@ -1082,6 +1130,8 @@ LLVM bitcode.)pb");
       "bindings"_a = QCODDBindingMap{},
       R"pb(Sample final computational-basis outcomes from a QCO program.
 
+The same ``QCOProgram`` must not be sampled concurrently from multiple threads.
+
 Args:
     program: A QCO program whose entry ``func.func`` is sampled.
     dd_package: DD package with enough qubits for the program. Not thread-safe;
@@ -1098,6 +1148,31 @@ LLVM bitcode.)pb");
 Raises:
     ValueError: When the program is unsupported for sampling.)pb");
 
+  m.def(
+      "sample_density",
+      [](const mlir::QCOProgram& program, const dd::MatrixDD& initialState,
+         dd::Package& ddPackage, const size_t shots,
+         const std::optional<uint64_t> seed,
+         const QCODDBindingMap& pythonBindings) {
+        auto func = entryFunc(program);
+        auto bindings = makeQCODDBindings(func, pythonBindings);
+        auto rng = makeRng(seed);
+        return takeFailureOr(
+            func.getContext(), "cannot density-sample this QCO program", [&] {
+              const nb::gil_scoped_release release;
+              return mlir::qco::sampleDensity(func, initialState, ddPackage,
+                                              shots, rng, bindings);
+            });
+      },
+      "program"_a, "initial_state"_a, "dd_package"_a, "shots"_a = 1024U,
+      "seed"_a = nb::none(), nb::kw_only(), "bindings"_a = QCODDBindingMap{},
+      R"pb(Sample a QCO program from a density-matrix DD.
+
+The input matrix reference is consumed. Mixed states and entangled qubit
+deallocation are supported. The DD package is not thread-safe and must not be
+shared across threads while sampling. The same ``QCOProgram`` must not be
+sampled concurrently from multiple threads.)pb");
+
   m.def(
       "sample_with_classics",
       [](const mlir::QCOProgram& program, dd::Package& ddPackage,
@@ -1123,6 +1198,8 @@ LLVM bitcode.)pb");
       "bindings"_a = QCODDBindingMap{},
       R"pb(Sample final and mid-circuit classical outcomes from a QCO program.
 
+The same ``QCOProgram`` must not be sampled concurrently from multiple threads.
+
 Args:
     program: A QCO program whose entry ``func.func`` is sampled.
     dd_package: DD package with enough qubits for the program. Not thread-safe;

diff --git a/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h b/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h
@@ -52,11 +52,12 @@ using DDBindings = DenseMap<Value, Attribute>;
  * - `qtensor.from_elements` / `extract` / `insert` / `dealloc` as linear
  *   bookkeeping over existing input wires
  * - Concrete `qco.if` / `qco.index_switch`, bounded `scf.for` / `scf.while`,
- *   standard `scf.if` / `scf.index_switch` / single-block
- *   `scf.execute_region`, concrete `cf.br` / `cf.cond_br` function CFGs, and
- *   non-recursive `func.call`
+ *   standard `scf.if` / `scf.index_switch` / multi-block
+ *   `scf.execute_region`, concrete `cf.br` / `cf.cond_br` / `cf.switch`
+ *   function CFGs, and non-recursive `func.call`
  * - Concrete integer, index, and floating-point `arith` operations and
- *   one-dimensional `memref` storage over those scalar types
+ *   common `math` operations, plus one-dimensional `memref` storage over those
+ *   scalar types, including aliases passed through `func.call`
  * - `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
@@ -82,7 +83,7 @@ buildFunctionality(func::FuncOp func, dd::Package& dd,
  *
  * @details Same supported unitary op set as @ref buildFunctionality, plus
  * concrete classical control-flow (`qco.if`, `qco.index_switch`, `scf.if`,
- * `scf.index_switch`, and single-block `scf.execute_region`) and static- or
+ * `scf.index_switch`, and multi-block `scf.execute_region`) and static- or
  * concrete dynamic-shape 1-D `memref` registers of integer, index, or
  * floating-point values (`alloc`/`store`/`load`/`dealloc`).
  * `qco.alloc` and `qtensor.alloc` append zero-state wires, while
@@ -128,10 +129,10 @@ FailureOr<dd::VectorDD> simulate(func::FuncOp func, const dd::VectorDD& in,
  * Deterministic control-flow without measure/reset also works on the non-RNG
  * overload. Only one-dimensional qtensors of qubits are supported. Nested
  * regions are walked; `scf.for` with concrete positive step, concrete
- * `scf.while`, concrete `cf.br` / `cf.cond_br` function CFGs, and non-recursive
- * `func.call` are supported. Loops and function CFGs are limited to 10000
- * transitions. Consumes one reference to @p in regardless of whether
- * simulation succeeds or fails.
+ * `scf.while`, concrete `cf.br` / `cf.cond_br` / `cf.switch` function CFGs,
+ * and non-recursive `func.call` are supported. Loops and function CFGs are
+ * limited to 10000 transitions. 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
@@ -145,6 +146,44 @@ FailureOr<dd::VectorDD> simulate(func::FuncOp func, const dd::VectorDD& in,
                                  dd::Package& dd, std::mt19937_64& rng,
                                  const DDBindings& bindings = DDBindings());
 
+/**
+ * @brief Construct the density operator @f$|\psi\rangle\langle\psi|@f$.
+ *
+ * @param state Pure input state; its reference is retained by the caller
+ * @param numQubits Number of active qubits represented by @p state
+ * @param dd The DD package to use
+ * @return A referenced matrix DD representing the pure-state density operator
+ */
+dd::MatrixDD makeDensityMatrix(const dd::VectorDD& state, size_t numQubits,
+                               dd::Package& dd);
+
+/**
+ * @brief Simulate a QCO function using a density-matrix DD.
+ *
+ * @details Unitary operations evolve the state as @f$U\rho U^\dagger@f$.
+ * Qubit and qtensor deallocation performs a physical partial trace, including
+ * for entangled qubits. The RNG overload additionally supports collapsing
+ * measurement and reset. Consumes one reference to @p in regardless of
+ * success or failure.
+ *
+ * @param func The QCO function to simulate
+ * @param in Input density matrix; one reference is consumed
+ * @param dd The DD package to use
+ * @param bindings Concrete values for symbolic scalar function arguments
+ * @return The output density-matrix DD on success
+ */
+FailureOr<dd::MatrixDD>
+simulateDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd,
+                const DDBindings& bindings = DDBindings());
+
+/// @copydoc simulateDensity(func::FuncOp, const dd::MatrixDD&, dd::Package&,
+/// const DDBindings&)
+/// Uses @p rng for collapsing measurement and reset.
+FailureOr<dd::MatrixDD>
+simulateDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd,
+                std::mt19937_64& rng,
+                const DDBindings& bindings = DDBindings());
+
 /**
  * @brief Sample measurement outcomes from a QCO `func.func`.
  *
@@ -187,6 +226,18 @@ sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng,
 FailureOr<std::map<std::string, size_t>>
 sample(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, size_t shots,
        std::mt19937_64& rng, const DDBindings& bindings = DDBindings());
+/**
+ * @brief Sample a QCO function from an input density-matrix DD.
+ *
+ * @details Supports mixed states and entangled qubit deallocation. Each final
+ * sample collapses a referenced copy of the simulated density state. Programs
+ * with mid-circuit measurement or reset are re-simulated per shot. Consumes one
+ * reference to @p in.
+ */
+FailureOr<std::map<std::string, size_t>>
+sampleDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd,
+              size_t shots, std::mt19937_64& rng,
+              const DDBindings& bindings = DDBindings());
 
 /// Histograms produced by @ref sampleWithClassics.
 struct SampleResult {

diff --git a/mlir/lib/Compiler/Programs.cpp b/mlir/lib/Compiler/Programs.cpp
@@ -512,11 +512,11 @@ std::optional<JeffProgram> QCOProgram::intoJeff() && {
 }
 
 std::optional<func::FuncOp> QCOProgram::entryFunc() const {
-  ModuleOp module = mod();
-  if (auto main = module.lookupSymbol<func::FuncOp>("main")) {
+  ModuleOp moduleOp = mod();
+  if (auto main = moduleOp.lookupSymbol<func::FuncOp>("main")) {
     return main;
   }
-  auto funcs = module.getBody()->getOps<func::FuncOp>();
+  auto funcs = moduleOp.getBody()->getOps<func::FuncOp>();
   if (funcs.empty()) {
     return std::nullopt;
   }