diff --git a/CHANGELOG.md b/CHANGELOG.md index 4994358b81..19b6992220 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,8 +39,11 @@ releases may include breaking changes. [#1755], [#1787], [#1815], [#1823], [#1830], [#1886], [#1933], [#1978], [#1979], [#2007]) ([**@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 target-independent two-qubit gate fusion, target-native post-routing synthesis, and operation-capability and static-site conformance ([#1865], [#1961], [#1998]) ([**@simon1hofmann**], [**@burgholzer**]) @@ -744,6 +747,7 @@ for previous 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 diff --git a/bindings/mlir/CMakeLists.txt b/bindings/mlir/CMakeLists.txt index 68e089c228..0c20aa3049 100644 --- a/bindings/mlir/CMakeLists.txt +++ b/bindings/mlir/CMakeLists.txt @@ -24,7 +24,8 @@ if(NOT TARGET ${TARGET_NAME}) LINK_LIBS MQTCompilerFoMaCAdapter MQTCompilerPipeline - MQT::CoreIR) + MQT::CoreIR + MLIRQCODDFunctionality) # install the Python stub file in editable mode for better IDE support if(SKBUILD_STATE STREQUAL "editable") diff --git a/bindings/mlir/register_mlir.cpp b/bindings/mlir/register_mlir.cpp index 65f8dff968..4ec51bd6a4 100644 --- a/bindings/mlir/register_mlir.cpp +++ b/bindings/mlir/register_mlir.cpp @@ -8,14 +8,22 @@ * Licensed under the MIT License */ +#include "dd/Package.hpp" #include "fomac/FoMaC.hpp" // NOLINT(misc-include-cleaner) #include "ir/QuantumComputation.hpp" #include "mlir/Compiler/FoMaCAdapter.h" #include "mlir/Compiler/Programs.h" #include "mlir/Compiler/Target.h" +#include "mlir/Dialect/QCO/Utils/DDFunctionality.h" +#include +#include +#include +#include +#include #include #include // NOLINT(misc-include-cleaner) +#include // NOLINT(misc-include-cleaner) #include // NOLINT(misc-include-cleaner) #include // NOLINT(misc-include-cleaner) #include // NOLINT(misc-include-cleaner) @@ -28,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -100,6 +109,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& 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 +[[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)(); + 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 [[nodiscard]] ProgramType copiedOrConsumed(ProgramType& program, const bool copy) { @@ -772,6 +833,136 @@ LLVM bitcode.)pb"); &BooleanMemberAdapter<&mlir::QIRProgram::writeBitcode>::call, "path"_a, "Write this program as LLVM bitcode."); + nb::module_::import_("mqt.core.dd"); + + nb::class_(m, "SampleResult", + R"pb(Histograms from QCO DD sampling.)pb") + .def_ro("shots", &mlir::qco::SampleResult::shots, + R"pb(Final computational-basis outcome histogram.)pb") + .def_ro("classical", &mlir::qco::SampleResult::classical, + R"pb(Mid-circuit measure-bit histogram (encounter order).)pb"); + + 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 used to build a matrix DD. + 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 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 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. Not thread-safe; + do not share it across threads while sampling (the GIL is released for + the duration of the call). + 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 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. Not thread-safe; + do not share it across threads while sampling (the GIL is released for + the duration of the call). + 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, "target"_a = nb::none(), "qco_pipeline"_a = "mqt-qco-default", diff --git a/mlir/include/mlir/Compiler/Programs.h b/mlir/include/mlir/Compiler/Programs.h index 82af524e5b..3958e81d73 100644 --- a/mlir/include/mlir/Compiler/Programs.h +++ b/mlir/include/mlir/Compiler/Programs.h @@ -10,6 +10,7 @@ #pragma once +#include #include #include #include @@ -241,6 +242,9 @@ class QCOProgram final : public Program { /// Consume this program and convert it to `jeff` MLIR. [[nodiscard]] std::optional intoJeff() &&; + + /// Return the entry `func.func` (`main` if present, else the first function). + [[nodiscard]] std::optional entryFunc() const; }; /** diff --git a/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h b/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h index 43f162ced0..a9e5eb077f 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h @@ -15,6 +15,11 @@ #include #include +#include +#include +#include +#include + 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,19 @@ namespace mlir::qco { FailureOr 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-shape + * 1-D `memref` classical registers (`alloc`/`store`/`load`/`dealloc`). + * Mid-circuit `measure` / `reset` require the RNG overload below. Concrete- + * bound `scf.for` loops and non-recursive single-block `func.call` are + * supported independently of RNG. Only qubit-typed linear values are supported + * (no qtensors). Nested regions are walked; `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, represented as a vector DD; one reference is @@ -66,4 +77,102 @@ FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd); FailureOr 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-shape 1-D `memref` 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 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> 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> +sample(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, + std::size_t shots, std::mt19937_64& rng); + +/// Histograms produced by @ref sampleWithClassics. +struct SampleResult { + /// Final computational-basis outcome histogram. + std::map shots; + /// Mid-circuit measurement-bit histogram (encounter order). + std::map classical; +}; + +/** + * @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. + */ +FailureOr sampleWithClassics(func::FuncOp func, dd::Package& dd, + size_t shots, std::mt19937_64& rng); + +/// @copydoc sampleWithClassics(func::FuncOp, dd::Package&, size_t, +/// std::mt19937_64&) +/// Starts from @p in; one reference is consumed. +FailureOr sampleWithClassics(func::FuncOp func, + const dd::VectorDD& in, + dd::Package& dd, size_t shots, + std::mt19937_64& rng); + } // namespace mlir::qco diff --git a/mlir/lib/Compiler/Programs.cpp b/mlir/lib/Compiler/Programs.cpp index 0d8f1eab5e..6258d372d0 100644 --- a/mlir/lib/Compiler/Programs.cpp +++ b/mlir/lib/Compiler/Programs.cpp @@ -474,6 +474,18 @@ std::optional QCOProgram::intoJeff() && { return JeffProgram(std::move(*this).releaseStorage()); } +std::optional QCOProgram::entryFunc() const { + ModuleOp module = mod(); + if (auto main = module.lookupSymbol("main")) { + return main; + } + auto funcs = module.getBody()->getOps(); + if (funcs.empty()) { + return std::nullopt; + } + return *funcs.begin(); +} + //===----------------------------------------------------------------------===// // JeffProgram //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt b/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt index 0908de519f..1c00210985 100644 --- a/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt +++ b/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt @@ -64,7 +64,9 @@ add_mlir_library( PUBLIC MLIRQCODialect MLIRQCOMatrix + MLIRArithDialect MLIRFuncDialect + MLIRMemRefDialect MQT::CoreDD) mqt_mlir_target_use_project_options(MLIRQCODDFunctionality) diff --git a/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp b/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp index 1630c8e6d7..cd126206f4 100644 --- a/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp +++ b/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp @@ -11,8 +11,10 @@ #include "mlir/Dialect/QCO/Utils/DDFunctionality.h" #include "dd/DDDefinitions.hpp" +#include "dd/GateMatrixDefinitions.hpp" #include "dd/Operations.hpp" #include "dd/Package.hpp" +#include "dd/StateGeneration.hpp" #include "ir/Definitions.hpp" #include "ir/operations/Control.hpp" #include "ir/operations/OpType.hpp" @@ -23,22 +25,34 @@ #include "mlir/Dialect/Utils/Utils.h" #include +#include #include #include #include #include #include +#include +#include +#include +#include +#include #include #include #include +#include #include #include +#include #include #include #include #include +#include #include +#include +#include +#include #include #include @@ -88,11 +102,69 @@ struct QubitMap { } }; +struct ClassicalEnv { + DenseMap bools; + DenseMap indices; + /// Backing storage for static-shape 1-D `memref` classical registers. + DenseMap> memrefs; + + LogicalResult bindFrom(Value source, Value dest, Operation* op) { + if (dest.getType().isInteger(1)) { + const auto it = bools.find(source); + if (it == bools.end()) { + return op->emitError() + << "classical i1 SSA value is not mapped for QCO DD simulation"; + } + bools[dest] = it->second; + return success(); + } + if (isa(dest.getType())) { + const auto it = indices.find(source); + if (it == indices.end()) { + return op->emitError() << "classical index SSA value is not mapped " + "for QCO DD simulation"; + } + indices[dest] = it->second; + return success(); + } + return op->emitError() + << "unsupported classical type for QCO DD simulation: " + << dest.getType(); + } +}; + struct DecodedGate { qc::OpType type = qc::OpType::None; std::vector params; }; +struct WalkState { + // Non-owning handles into the active simulation frame. + QubitMap& qubits; // NOLINT(cppcoreguidelines-avoid-const-or-ref-data-members) + ClassicalEnv& + classical; // NOLINT(cppcoreguidelines-avoid-const-or-ref-data-members) + dd::Package& dd; // NOLINT(cppcoreguidelines-avoid-const-or-ref-data-members) + std::mt19937_64* rng = nullptr; + std::string* classicalBits = nullptr; + DenseSet* activeCalls = nullptr; +}; + +/// Erases @p op from @p set on destruction (used around `func.call`). +struct ActiveCallGuard { + DenseSet* set = nullptr; + Operation* op = nullptr; + + ActiveCallGuard(DenseSet* activeSet, Operation* callee) + : set(activeSet), op(callee) {} + ~ActiveCallGuard() { + if (set != nullptr && op != nullptr) { + set->erase(op); + } + } + ActiveCallGuard(const ActiveCallGuard&) = delete; + ActiveCallGuard& operator=(const ActiveCallGuard&) = delete; +}; + } // namespace /// `std::nullopt` if @p unitary is not a standard gate; failure if its unitary @@ -150,8 +222,12 @@ decodeStandardGate(UnitaryOpInterface unitary) { /// QCO matrices are MSB-first (operand 0 = high bit). [[nodiscard]] static size_t qcoIndexFromDdIndex(const size_t ddIndex, const size_t numQubits) { - const auto shift = static_cast(64 - numQubits); - return llvm::reverseBits(ddIndex) >> shift; + // Dense embed is capped at 12 qubits; guard the shift width for the analyzer. + if (numQubits == 0 || numQubits > 63) { + return ddIndex; + } + const auto shift = static_cast(64U - numQubits); + return llvm::reverseBits(static_cast(ddIndex)) >> shift; } [[nodiscard]] static dd::CMat toCMatInDdBasis(const DynamicMatrix& qcoMatrix, @@ -168,10 +244,50 @@ decodeStandardGate(UnitaryOpInterface unitary) { return out; } +/// Embed a k-qubit QCO/MSB matrix onto @p wires of an n-qubit register. +[[nodiscard]] static DynamicMatrix +embedLocalInNQubitMsb(const DynamicMatrix& local, size_t n, + ArrayRef wires) { + const size_t k = wires.size(); + const auto dimN = static_cast(size_t{1} << n); + DynamicMatrix out(dimN); + const auto dimNSz = static_cast(dimN); + auto bitAt = [](size_t idx, size_t nQ, size_t q) -> size_t { + return (idx >> (nQ - 1 - q)) & 1U; + }; + llvm::SmallDenseSet wireSet; + wireSet.insert(wires.begin(), wires.end()); + for (size_t row = 0; row < dimNSz; ++row) { + for (size_t col = 0; col < dimNSz; ++col) { + bool idleMatch = true; + for (size_t q = 0; q < n; ++q) { + if (wireSet.contains(static_cast(q))) { + continue; + } + if (bitAt(row, n, q) != bitAt(col, n, q)) { + idleMatch = false; + break; + } + } + if (!idleMatch) { + continue; + } + size_t rLoc = 0; + size_t cLoc = 0; + for (size_t i = 0; i < k; ++i) { + rLoc = (rLoc << 1) | bitAt(row, n, wires[i]); + cLoc = (cLoc << 1) | bitAt(col, n, wires[i]); + } + out(static_cast(row), static_cast(col)) = + local(static_cast(rLoc), static_cast(cLoc)); + } + } + return out; +} + template static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, - QubitMap& qubits, dd::Package& dd, - StateDD& state) { + WalkState& walk, StateDD& state) { Operation* op = unitary.getOperation(); if (!unitary.hasCompileTimeKnownUnitaryMatrix()) { return unitary.emitError() @@ -180,12 +296,12 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, if (auto gphase = dyn_cast(op)) { const auto theta = *utils::valueToDouble(gphase.getTheta()); auto id = dd::Package::makeIdent(); - id.w = dd.cn.lookup(std::cos(theta), std::sin(theta)); - state = dd.applyOperation(id, state); + id.w = walk.dd.cn.lookup(std::cos(theta), std::sin(theta)); + state = walk.dd.applyOperation(id, state); return success(); } if (isa(op)) { - return qubits.remapUnitary(unitary); + return walk.qubits.remapUnitary(unitary); } DynamicMatrix local; @@ -194,7 +310,7 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, << "unitary must have a compile-time constant matrix"; } - auto wiresOr = qubits.lookupRange(unitary.getInputQubits(), op); + auto wiresOr = walk.qubits.lookupRange(unitary.getInputQubits(), op); if (failed(wiresOr)) { return failure(); } @@ -203,8 +319,8 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, if (wires.size() == 1) { const dd::GateMatrix mat{local(0, 0), local(0, 1), local(1, 0), local(1, 1)}; - state = dd.applyOperation(dd.makeGateDD(mat, wires[0]), state); - return qubits.remapUnitary(unitary); + state = walk.dd.applyOperation(walk.dd.makeGateDD(mat, wires[0]), state); + return walk.qubits.remapUnitary(unitary); } if (wires.size() == 2) { @@ -215,9 +331,9 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, local(static_cast(row), static_cast(col)); } } - state = dd.applyOperation(dd.makeTwoQubitGateDD(mat, wires[0], wires[1]), - state); - return qubits.remapUnitary(unitary); + state = walk.dd.applyOperation( + walk.dd.makeTwoQubitGateDD(mat, wires[0], wires[1]), state); + return walk.qubits.remapUnitary(unitary); } if (wires.size() == 3) { @@ -228,49 +344,51 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, local(static_cast(row), static_cast(col)); } } - state = dd.applyOperation( - dd.makeThreeQubitGateDD(mat, wires[0], wires[1], wires[2]), state); - return qubits.remapUnitary(unitary); + state = walk.dd.applyOperation( + walk.dd.makeThreeQubitGateDD(mat, wires[0], wires[1], wires[2]), state); + return walk.qubits.remapUnitary(unitary); } - // Map full-width matrices from QCO/MSB order to DD/LSB. Cap at 12 qubits - // (~256 MiB dense `CMat`). - if (qubits.numQubits > 12) { + // Dense embed of a k-qubit QCO/MSB matrix into n wires, then rewrite to + // DD/LSB. Cap at 12 qubits (~256 MiB dense `CMat`). + if (walk.qubits.numQubits > 12) { return unitary.emitError() << "QCO DD matrix fallback supports at most 12 qubits"; } - if (wires.size() != qubits.numQubits || - !llvm::all_of(llvm::enumerate(wires), - [](const auto& it) { return it.value() == it.index(); })) { - return op->emitError() - << "QCO DD matrix fallback supports full-width unitaries on qubits " - "0..n-1"; + DynamicMatrix embedded = local; + const bool fullWidthCanonical = + wires.size() == walk.qubits.numQubits && + llvm::all_of(llvm::enumerate(wires), + [](const auto& it) { return it.value() == it.index(); }); + if (!fullWidthCanonical) { + embedded = embedLocalInNQubitMsb(local, walk.qubits.numQubits, wires); } - state = dd.applyOperation( - dd.makeDDFromMatrix(toCMatInDdBasis(local, qubits.numQubits)), state); - return qubits.remapUnitary(unitary); + state = walk.dd.applyOperation(walk.dd.makeDDFromMatrix(toCMatInDdBasis( + embedded, walk.qubits.numQubits)), + state); + return walk.qubits.remapUnitary(unitary); } template -static LogicalResult -applyDecodedStandard(UnitaryOpInterface unitary, const DecodedGate& gate, - const qc::Controls& controls, QubitMap& qubits, - dd::Package& dd, StateDD& state) { +static LogicalResult applyDecodedStandard(UnitaryOpInterface unitary, + const DecodedGate& gate, + const qc::Controls& controls, + WalkState& walk, StateDD& state) { SmallVector targetVals; for (size_t i = 0; i < unitary.getNumTargets(); ++i) { targetVals.push_back(unitary.getInputTarget(i)); } - auto targets = qubits.lookupRange(targetVals, unitary.getOperation()); + auto targets = walk.qubits.lookupRange(targetVals, unitary.getOperation()); if (failed(targets)) { return failure(); } - state = dd.applyOperation( - getStandardOperationDD(dd, gate.type, gate.params, controls, + state = walk.dd.applyOperation( + getStandardOperationDD(walk.dd, gate.type, gate.params, controls, {targets->begin(), targets->end()}), state); - return qubits.remapUnitary(unitary); + return walk.qubits.remapUnitary(unitary); } static LogicalResult validateReturn(func::ReturnOp returnOp, @@ -298,14 +416,800 @@ static LogicalResult validateReturn(func::ReturnOp returnOp, return success(); } +static LogicalResult recordConstant(arith::ConstantOp constant, + ClassicalEnv& classical) { + // `arith.constant true/false` is a BoolAttr; other integers are IntegerAttr. + if (auto boolAttr = dyn_cast(constant.getValue())) { + classical.bools[constant.getResult()] = boolAttr.getValue(); + return success(); + } + auto attr = dyn_cast(constant.getValue()); + if (!attr) { + return success(); + } + if (constant.getType().isInteger(1)) { + classical.bools[constant.getResult()] = attr.getValue() != 0; + } else if (isa(constant.getType())) { + classical.indices[constant.getResult()] = attr.getInt(); + } + return success(); +} + +static FailureOr lookupBool(Value value, ClassicalEnv& classical, + Operation* op) { + const auto it = classical.bools.find(value); + if (it == classical.bools.end()) { + return op->emitError() + << "classical i1 SSA value is not mapped for QCO DD simulation"; + } + return it->second; +} + +static FailureOr lookupIndex(Value value, ClassicalEnv& classical, + Operation* op) { + const auto it = classical.indices.find(value); + if (it == classical.indices.end()) { + return op->emitError() + << "classical index SSA value is not mapped for QCO DD simulation"; + } + return it->second; +} + +/// Cast a concrete `i1` SSA value to `index` (shared by `extui` / +/// `index_castui`). +static LogicalResult applyI1ToIndex(Value in, Value out, Operation* op, + ClassicalEnv& classical) { + if (!isa(out.getType())) { + return op->emitError() + << "QCO DD simulation only supports casting i1 to index"; + } + if (!in.getType().isInteger(1)) { + return op->emitError() + << "QCO DD simulation only supports casting from i1 to index"; + } + auto bit = lookupBool(in, classical, op); + if (failed(bit)) { + return failure(); + } + classical.indices[out] = *bit ? 1 : 0; + return success(); +} + +[[nodiscard]] static bool isStaticI1MemRef(Type type) { + auto memref = dyn_cast(type); + return memref && memref.getRank() == 1 && memref.hasStaticShape() && + memref.getElementType().isInteger(1); +} + +/// Resolve a static-shape 1-D `i1` memref classical register and a concrete +/// index. +static FailureOr*, int64_t>> +lookupI1MemRefSlot(Value memref, Value index, ClassicalEnv& classical, + Operation* op) { + if (!isStaticI1MemRef(memref.getType())) { + return op->emitError() << "QCO DD simulation only supports static-shape " + "1-D memref classical registers"; + } + auto idx = lookupIndex(index, classical, op); + if (failed(idx)) { + return failure(); + } + auto it = classical.memrefs.find(memref); + if (it == classical.memrefs.end()) { + return op->emitError() + << "classical memref is not mapped for QCO DD simulation"; + } + if (*idx < 0 || static_cast(*idx) >= it->second.size()) { + return op->emitError() + << "classical memref index out of range for QCO DD simulation"; + } + return std::pair{&it->second, *idx}; +} + +static LogicalResult applyMemRefAlloc(memref::AllocOp alloc, + ClassicalEnv& classical) { + if (!isStaticI1MemRef(alloc.getType())) { + return alloc.emitError() << "QCO DD simulation only supports static-shape " + "1-D memref classical registers"; + } + if (!alloc.getDynamicSizes().empty() || !alloc.getSymbolOperands().empty()) { + return alloc.emitError() + << "QCO DD simulation does not support dynamic memref allocation"; + } + auto type = cast(alloc.getType()); + classical.memrefs[alloc.getResult()] = + SmallVector(static_cast(type.getDimSize(0)), false); + return success(); +} + +static LogicalResult applyMemRefStore(memref::StoreOp store, + ClassicalEnv& classical) { + if (store.getIndices().size() != 1) { + return store.emitError() + << "QCO DD simulation only supports 1-D memref.store"; + } + Value value = store.getValue(); + if (!value.getType().isInteger(1)) { + return store.emitError() + << "QCO DD simulation only supports storing i1 into classical " + "registers"; + } + auto bit = lookupBool(value, classical, store); + if (failed(bit)) { + return failure(); + } + auto slot = lookupI1MemRefSlot(store.getMemref(), store.getIndices()[0], + classical, store); + if (failed(slot)) { + return failure(); + } + (*slot->first)[static_cast(slot->second)] = *bit; + return success(); +} + +static LogicalResult applyMemRefLoad(memref::LoadOp load, + ClassicalEnv& classical) { + if (load.getIndices().size() != 1) { + return load.emitError() + << "QCO DD simulation only supports 1-D memref.load"; + } + if (!load.getResult().getType().isInteger(1)) { + return load.emitError() + << "QCO DD simulation only supports loading i1 from classical " + "registers"; + } + auto slot = lookupI1MemRefSlot(load.getMemref(), load.getIndices()[0], + classical, load); + if (failed(slot)) { + return failure(); + } + classical.bools[load.getResult()] = + (*slot->first)[static_cast(slot->second)]; + return success(); +} + +template +static LogicalResult applyBinaryI1(OpTy op, ClassicalEnv& classical, + bool (*combine)(bool, bool)) { + if (!op.getType().isInteger(1)) { + return op.emitError() << "QCO DD simulation only supports i1 " + << op.getOperationName(); + } + auto lhs = lookupBool(op.getLhs(), classical, op); + auto rhs = lookupBool(op.getRhs(), classical, op); + if (failed(lhs) || failed(rhs)) { + return failure(); + } + classical.bools[op.getResult()] = combine(*lhs, *rhs); + return success(); +} + +template +static LogicalResult applyBinaryIndex(OpTy op, ClassicalEnv& classical, + int64_t (*combine)(int64_t, int64_t)) { + if (!isa(op.getType())) { + return op.emitError() << "QCO DD simulation only supports index " + << op.getOperationName(); + } + auto lhs = lookupIndex(op.getLhs(), classical, op); + auto rhs = lookupIndex(op.getRhs(), classical, op); + if (failed(lhs) || failed(rhs)) { + return failure(); + } + classical.indices[op.getResult()] = combine(*lhs, *rhs); + return success(); +} + +static LogicalResult applyClassicalOp(Operation& op, ClassicalEnv& classical) { + return TypeSwitch(&op) + .Case([&](arith::AndIOp andOp) -> LogicalResult { + if (andOp.getType().isInteger(1)) { + return applyBinaryI1(andOp, classical, + [](bool a, bool b) { return a && b; }); + } + return applyBinaryIndex(andOp, classical, + [](int64_t a, int64_t b) { return a & b; }); + }) + .Case([&](arith::OrIOp orOp) -> LogicalResult { + if (orOp.getType().isInteger(1)) { + return applyBinaryI1(orOp, classical, + [](bool a, bool b) { return a || b; }); + } + return applyBinaryIndex(orOp, classical, + [](int64_t a, int64_t b) { return a | b; }); + }) + .Case([&](arith::XOrIOp xorOp) -> LogicalResult { + if (xorOp.getType().isInteger(1)) { + return applyBinaryI1(xorOp, classical, + [](bool a, bool b) { return a != b; }); + } + return applyBinaryIndex(xorOp, classical, + [](int64_t a, int64_t b) { return a ^ b; }); + }) + .Case([&](arith::AddIOp addOp) { + return applyBinaryIndex(addOp, classical, [](int64_t a, int64_t b) { + return static_cast(static_cast(a) + + static_cast(b)); + }); + }) + .Case([&](arith::SubIOp subOp) { + return applyBinaryIndex(subOp, classical, [](int64_t a, int64_t b) { + return static_cast(static_cast(a) - + static_cast(b)); + }); + }) + .Case([&](arith::MulIOp mulOp) { + return applyBinaryIndex(mulOp, classical, [](int64_t a, int64_t b) { + return static_cast(static_cast(a) * + static_cast(b)); + }); + }) + .Case([&](arith::ShLIOp shli) -> LogicalResult { + if (!isa(shli.getType())) { + return shli.emitError() << "QCO DD simulation only supports index " + << arith::ShLIOp::getOperationName(); + } + auto lhs = lookupIndex(shli.getLhs(), classical, shli); + auto rhs = lookupIndex(shli.getRhs(), classical, shli); + if (failed(lhs) || failed(rhs)) { + return failure(); + } + if (*rhs < 0 || *rhs >= 64) { + return shli.emitError() + << "shift amount out of range for QCO DD simulation"; + } + classical.indices[shli.getResult()] = static_cast( + static_cast(*lhs) << static_cast(*rhs)); + return success(); + }) + .Case([&](arith::ShRUIOp shrui) -> LogicalResult { + if (!isa(shrui.getType())) { + return shrui.emitError() << "QCO DD simulation only supports index " + << arith::ShRUIOp::getOperationName(); + } + auto lhs = lookupIndex(shrui.getLhs(), classical, shrui); + auto rhs = lookupIndex(shrui.getRhs(), classical, shrui); + if (failed(lhs) || failed(rhs)) { + return failure(); + } + if (*rhs < 0 || *rhs >= 64) { + return shrui.emitError() + << "shift amount out of range for QCO DD simulation"; + } + classical.indices[shrui.getResult()] = static_cast( + static_cast(*lhs) >> static_cast(*rhs)); + return success(); + }) + .Case([&](arith::CmpIOp cmp) -> LogicalResult { + if (!cmp.getType().isInteger(1)) { + return cmp.emitError() + << "QCO DD simulation only supports cmpi to i1"; + } + FailureOr lhs; + FailureOr rhs; + if (isa(cmp.getLhs().getType())) { + lhs = lookupIndex(cmp.getLhs(), classical, cmp); + rhs = lookupIndex(cmp.getRhs(), classical, cmp); + } else if (cmp.getLhs().getType().isInteger(1)) { + auto lb = lookupBool(cmp.getLhs(), classical, cmp); + auto rb = lookupBool(cmp.getRhs(), classical, cmp); + if (failed(lb) || failed(rb)) { + return failure(); + } + // Unsigned i1 uses 0/1; signed predicates use arith sign-extension + // (true → -1). Equality is identical under either encoding. + const int64_t aU = *lb ? 1 : 0; + const int64_t bU = *rb ? 1 : 0; + const int64_t aS = *lb ? -1 : 0; + const int64_t bS = *rb ? -1 : 0; + bool result = false; + switch (cmp.getPredicate()) { + case arith::CmpIPredicate::eq: + result = aU == bU; + break; + case arith::CmpIPredicate::ne: + result = aU != bU; + break; + case arith::CmpIPredicate::slt: + result = aS < bS; + break; + case arith::CmpIPredicate::sle: + result = aS <= bS; + break; + case arith::CmpIPredicate::sgt: + result = aS > bS; + break; + case arith::CmpIPredicate::sge: + result = aS >= bS; + break; + case arith::CmpIPredicate::ult: + result = static_cast(aU) < static_cast(bU); + break; + case arith::CmpIPredicate::ule: + result = static_cast(aU) <= static_cast(bU); + break; + case arith::CmpIPredicate::ugt: + result = static_cast(aU) > static_cast(bU); + break; + case arith::CmpIPredicate::uge: + result = static_cast(aU) >= static_cast(bU); + break; + } + classical.bools[cmp.getResult()] = result; + return success(); + } else { + return cmp.emitError() + << "QCO DD simulation only supports cmpi on i1 or index"; + } + if (failed(lhs) || failed(rhs)) { + return failure(); + } + const int64_t a = *lhs; + const int64_t b = *rhs; + bool result = false; + switch (cmp.getPredicate()) { + case arith::CmpIPredicate::eq: + result = a == b; + break; + case arith::CmpIPredicate::ne: + result = a != b; + break; + case arith::CmpIPredicate::slt: + result = a < b; + break; + case arith::CmpIPredicate::sle: + result = a <= b; + break; + case arith::CmpIPredicate::sgt: + result = a > b; + break; + case arith::CmpIPredicate::sge: + result = a >= b; + break; + case arith::CmpIPredicate::ult: + result = static_cast(a) < static_cast(b); + break; + case arith::CmpIPredicate::ule: + result = static_cast(a) <= static_cast(b); + break; + case arith::CmpIPredicate::ugt: + result = static_cast(a) > static_cast(b); + break; + case arith::CmpIPredicate::uge: + result = static_cast(a) >= static_cast(b); + break; + } + classical.bools[cmp.getResult()] = result; + return success(); + }) + .Case([&](arith::SelectOp select) -> LogicalResult { + auto cond = lookupBool(select.getCondition(), classical, select); + if (failed(cond)) { + return failure(); + } + if (select.getType().isInteger(1)) { + auto t = lookupBool(select.getTrueValue(), classical, select); + auto f = lookupBool(select.getFalseValue(), classical, select); + if (failed(t) || failed(f)) { + return failure(); + } + classical.bools[select.getResult()] = *cond ? *t : *f; + return success(); + } + if (isa(select.getType())) { + auto t = lookupIndex(select.getTrueValue(), classical, select); + auto f = lookupIndex(select.getFalseValue(), classical, select); + if (failed(t) || failed(f)) { + return failure(); + } + classical.indices[select.getResult()] = *cond ? *t : *f; + return success(); + } + return select.emitError() + << "QCO DD simulation only supports select on i1 or index"; + }) + .Case([&](arith::ExtUIOp ext) { + return applyI1ToIndex(ext.getIn(), ext.getOut(), ext, classical); + }) + .Case([&](arith::IndexCastUIOp cast) { + return applyI1ToIndex(cast.getIn(), cast.getOut(), cast, classical); + }) + .Case([&](arith::TruncIOp trunc) -> LogicalResult { + if (!trunc.getType().isInteger(1)) { + return trunc.emitError() + << "QCO DD simulation only supports trunci to i1"; + } + Value in = trunc.getIn(); + if (!isa(in.getType())) { + return trunc.emitError() + << "QCO DD simulation only supports trunci from index"; + } + auto idx = lookupIndex(in, classical, trunc); + if (failed(idx)) { + return failure(); + } + classical.bools[trunc.getOut()] = (*idx & 1) != 0; + return success(); + }) + .Default([](Operation* unsupported) { + return unsupported->emitError() + << "unsupported classical op for QCO DD simulation: " + << unsupported->getName().getStringRef(); + }); +} + +static LogicalResult bindLinearArgs(ValueRange operands, Block& block, + WalkState& walk, Operation* op) { + if (operands.size() != block.getNumArguments()) { + return op->emitError() + << "region argument count does not match linear operands"; + } + for (auto [operand, arg] : llvm::zip_equal(operands, block.getArguments())) { + if (!isa(arg.getType())) { + return op->emitError() + << "QCO DD simulation does not support qtensor linear region " + "args (qubits only)"; + } + const auto q = walk.qubits.lookup(operand); + if (!q) { + return op->emitError() + << "qubit SSA value is not mapped for QCO DD construction"; + } + walk.qubits.bind(arg, *q); + } + return success(); +} + +/// Bind each source SSA onto the corresponding dest (qubits via `QubitMap`, +/// classical via `ClassicalEnv::bindFrom`). Callers must ensure equal sizes. +static LogicalResult bindValuePairs(ValueRange sources, ValueRange dests, + WalkState& walk, Operation* op) { + for (auto [src, dest] : llvm::zip_equal(sources, dests)) { + if (isa(dest.getType())) { + if (!isa(src.getType())) { + return op->emitError() + << "qubit/classical SSA type mismatch for QCO DD simulation"; + } + const auto q = walk.qubits.lookup(src); + if (!q) { + return op->emitError() + << "qubit SSA value is not mapped for QCO DD construction"; + } + walk.qubits.bind(dest, *q); + } else if (failed(walk.classical.bindFrom(src, dest, op))) { + return failure(); + } + } + return success(); +} + +static LogicalResult bindYieldResults(YieldOp yield, + ValueRange classicalResults, + ValueRange linearResults, + WalkState& walk) { + const size_t expected = classicalResults.size() + linearResults.size(); + if (yield.getNumOperands() != expected) { + return yield.emitError() + << "yield operand count does not match result segments"; + } + size_t idx = 0; + for (Value result : classicalResults) { + if (failed( + walk.classical.bindFrom(yield.getOperand(idx++), result, yield))) { + return failure(); + } + } + for (Value result : linearResults) { + if (!isa(result.getType())) { + return yield.emitError() + << "QCO DD simulation does not support qtensor linear results " + "(qubits only)"; + } + const auto q = walk.qubits.lookup(yield.getOperand(idx++)); + if (!q) { + return yield.emitError() + << "yielded qubit SSA value is not mapped for QCO DD construction"; + } + walk.qubits.bind(result, *q); + } + return success(); +} + +template +static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state); + +template +static LogicalResult walkFunction(func::FuncOp func, WalkState& walkState, + StateDD& state); + template -static LogicalResult applyOp(Operation& op, QubitMap& qubits, dd::Package& dd, - StateDD& state) { +static LogicalResult walkBlock(Block& block, WalkState& walk, StateDD& state) { + for (Operation& op : block.without_terminator()) { + if (failed(applyOp(op, walk, state))) { + return failure(); + } + } + return success(); +} + +template +static LogicalResult +applyRegionBranch(ValueRange linearOperands, Block& block, YieldOp yield, + ValueRange classicalResults, ValueRange linearResults, + WalkState& walk, StateDD& state, Operation* parent) { + if (failed(bindLinearArgs(linearOperands, block, walk, parent))) { + return failure(); + } + if (failed(walkBlock(block, walk, state))) { + return failure(); + } + return bindYieldResults(yield, classicalResults, linearResults, walk); +} + +template +static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { return TypeSwitch(&op) - .template Case( - [](auto) { return success(); }) + .template Case([](auto) { return success(); }) + .template Case([&](arith::ConstantOp constant) { + return recordConstant(constant, walk.classical); + }) + .template Case([&](memref::AllocOp alloc) { + return applyMemRefAlloc(alloc, walk.classical); + }) + .template Case([&](memref::StoreOp store) { + return applyMemRefStore(store, walk.classical); + }) + .template Case([&](memref::LoadOp load) { + return applyMemRefLoad(load, walk.classical); + }) + .template Case([](auto) { return success(); }) + .template Case( + [&](Operation* classicalOp) { + return applyClassicalOp(*classicalOp, walk.classical); + }) .template Case([&](func::ReturnOp returnOp) { - return validateReturn(returnOp, qubits); + return validateReturn(returnOp, walk.qubits); + }) + .template Case([&](MeasureOp measureOp) -> LogicalResult { + if constexpr (!std::is_same_v) { + return measureOp.emitError() + << "measurements are not supported for QCO DD functionality " + "construction"; + } else { + if (walk.rng == nullptr) { + return measureOp.emitError() + << "measurements require simulate(..., rng)"; + } + const auto q = walk.qubits.lookup(measureOp.getQubitIn()); + if (!q) { + return measureOp.emitError() + << "qubit SSA value is not mapped for QCO DD construction"; + } + const char bit = walk.dd.measureOneCollapsing(state, *q, *walk.rng); + walk.classical.bools[measureOp.getResult()] = bit == '1'; + if (walk.classicalBits != nullptr) { + walk.classicalBits->push_back(bit); + } + walk.qubits.bind(measureOp.getQubitOut(), *q); + return success(); + } + }) + .template Case([&](ResetOp resetOp) -> LogicalResult { + if constexpr (!std::is_same_v) { + return resetOp.emitError() + << "resets are not supported for QCO DD functionality " + "construction"; + } else { + if (walk.rng == nullptr) { + return resetOp.emitError() << "resets require simulate(..., rng)"; + } + const auto q = walk.qubits.lookup(resetOp.getQubitIn()); + if (!q) { + return resetOp.emitError() + << "qubit SSA value is not mapped for QCO DD construction"; + } + const char bit = walk.dd.measureOneCollapsing(state, *q, *walk.rng); + if (bit == '1') { + state = walk.dd.applyOperation( + walk.dd.makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), + *q), + state); + } + walk.qubits.bind(resetOp.getQubitOut(), *q); + return success(); + } + }) + .template Case([&](IfOp ifOp) -> LogicalResult { + if constexpr (!std::is_same_v) { + return ifOp.emitError() + << "control-flow is not supported for QCO DD functionality " + "construction"; + } else { + const auto condIt = walk.classical.bools.find(ifOp.getCondition()); + if (condIt == walk.classical.bools.end()) { + return ifOp.emitError() + << "if condition is not a concrete classical value"; + } + Block* block = condIt->second ? ifOp.thenBlock() : ifOp.elseBlock(); + if (block == nullptr) { + return ifOp.emitError() << "if region block is missing"; + } + YieldOp yield = condIt->second ? ifOp.thenYield() : ifOp.elseYield(); + return applyRegionBranch(ifOp.getQubits(), *block, yield, + ifOp.getClassicalResults(), + ifOp.getLinearResults(), walk, state, ifOp); + } + }) + .template Case( + [&](IndexSwitchOp switchOp) -> LogicalResult { + if constexpr (!std::is_same_v) { + return switchOp.emitError() + << "control-flow is not supported for QCO DD " + "functionality construction"; + } else { + const auto idxIt = walk.classical.indices.find(switchOp.getArg()); + if (idxIt == walk.classical.indices.end()) { + return switchOp.emitError() + << "index_switch argument is not a concrete index"; + } + const int64_t selector = idxIt->second; + const auto cases = switchOp.getCases(); + if (switchOp.getDefaultRegion().empty()) { + return switchOp.emitError() + << "index_switch default region is missing or empty"; + } + Block* block = switchOp.getDefaultBlock(); + YieldOp yield = switchOp.getDefaultYield(); + if (block == nullptr) { + return switchOp.emitError() + << "index_switch default region is missing or empty"; + } + for (auto [i, caseValue] : llvm::enumerate(cases)) { + if (caseValue == selector) { + block = switchOp.getCaseBlock(i); + yield = switchOp.getCaseYield(i); + break; + } + } + return applyRegionBranch(switchOp.getTargets(), *block, yield, + switchOp.getClassicalResults(), + switchOp.getLinearResults(), walk, state, + switchOp); + } + }) + .template Case([&](scf::ForOp forOp) -> LogicalResult { + if constexpr (!std::is_same_v) { + return forOp.emitError() + << "scf.for is not supported for QCO DD functionality " + "construction"; + } else { + auto lb = lookupIndex(forOp.getLowerBound(), walk.classical, forOp); + auto ub = lookupIndex(forOp.getUpperBound(), walk.classical, forOp); + auto step = lookupIndex(forOp.getStep(), walk.classical, forOp); + if (failed(lb) || failed(ub) || failed(step)) { + return failure(); + } + if (*step <= 0) { + return forOp.emitError() + << "scf.for step must be positive for QCO DD simulation"; + } + constexpr int64_t maxTrips = 10000; + int64_t trips = 0; + if (*ub > *lb) { + // Use unsigned arithmetic to avoid signed-overflow UB when + // classical bounds are extreme (e.g. INT64_MIN / INT64_MAX). + const auto span = + static_cast(*ub) - static_cast(*lb); + const uint64_t tripsU = + ((span - 1) / static_cast(*step)) + 1; + if (tripsU > static_cast(maxTrips)) { + return forOp.emitError() + << "scf.for trip count exceeds QCO DD simulation limit of " + << maxTrips; + } + trips = static_cast(tripsU); + } + + Block& body = *forOp.getBody(); + SmallVector carried(forOp.getInits().begin(), + forOp.getInits().end()); + + if (trips == 0) { + if (carried.size() != forOp.getNumResults()) { + return forOp.emitError() + << "scf.for result size mismatch during simulation"; + } + return bindValuePairs(carried, forOp.getResults(), walk, forOp); + } + + for (int64_t t = 0; t < trips; ++t) { + const auto offset = + static_cast(t) * static_cast(*step); + walk.classical.indices[body.getArgument(0)] = + static_cast(static_cast(*lb) + offset); + auto iterArgs = body.getArguments().drop_front(); + if (carried.size() != iterArgs.size()) { + return forOp.emitError() + << "scf.for iter_args size mismatch during simulation"; + } + if (failed(bindValuePairs(carried, iterArgs, walk, forOp))) { + return failure(); + } + if (failed(walkBlock(body, walk, state))) { + return failure(); + } + auto yield = dyn_cast(body.getTerminator()); + if (!yield) { + return forOp.emitError() << "scf.for body missing scf.yield"; + } + carried.assign(yield.getOperands().begin(), + yield.getOperands().end()); + } + if (carried.size() != forOp.getNumResults()) { + return forOp.emitError() + << "scf.for result size mismatch during simulation"; + } + return bindValuePairs(carried, forOp.getResults(), walk, forOp); + } + }) + .template Case([&](func::CallOp call) -> LogicalResult { + auto module = call->getParentOfType(); + if (!module) { + return call.emitError() + << "func.call requires a parent ModuleOp for QCO DD " + "simulation"; + } + auto callee = module.lookupSymbol(call.getCallee()); + if (!callee) { + return call.emitError() + << "func.call callee not found: " << call.getCallee(); + } + if (!callee.getBody().hasOneBlock()) { + return call.emitError() + << "func.call callee must have a single-block body"; + } + if (walk.activeCalls == nullptr) { + return call.emitError() + << "internal error: missing active call set for QCO DD"; + } + Operation* calleeOp = callee.getOperation(); + if (!walk.activeCalls->insert(calleeOp).second) { + return call.emitError() + << "recursive func.call is not supported for QCO DD " + "simulation"; + } + ActiveCallGuard guard(walk.activeCalls, calleeOp); + + if (call.getArgOperands().size() != callee.getNumArguments()) { + return call.emitError() + << "func.call operand count does not match callee arguments"; + } + if (failed(bindValuePairs(call.getArgOperands(), callee.getArguments(), + walk, call))) { + return failure(); + } + + // Walk the callee body without its terminator so entry-function-only + // `validateReturn` (canonical wire order) is not applied to callees. + if (failed(walkBlock(callee.getBody().front(), walk, state))) { + return failure(); + } + + // Map callee return operands onto call results via the return op. + auto returnOp = + dyn_cast(callee.getBody().front().getTerminator()); + if (!returnOp) { + return call.emitError() << "callee missing func.return"; + } + if (returnOp.getNumOperands() != call.getNumResults()) { + return call.emitError() + << "func.call result count does not match callee return"; + } + return bindValuePairs(returnOp.getOperands(), call.getResults(), walk, + call); }) .template Case([&](CtrlOp ctrlOp) -> LogicalResult { if (auto inner = utils::getSoleBodyUnitary( @@ -316,7 +1220,7 @@ static LogicalResult applyOp(Operation& op, QubitMap& qubits, dd::Package& dd, } if (*decoded) { auto controlQubits = - qubits.lookupRange(ctrlOp.getControlsIn(), ctrlOp); + walk.qubits.lookupRange(ctrlOp.getControlsIn(), ctrlOp); if (failed(controlQubits)) { return failure(); } @@ -324,11 +1228,11 @@ static LogicalResult applyOp(Operation& op, QubitMap& qubits, dd::Package& dd, for (qc::Qubit q : *controlQubits) { controls.emplace(q); } - return applyDecodedStandard(ctrlOp, **decoded, controls, qubits, dd, + return applyDecodedStandard(ctrlOp, **decoded, controls, walk, state); } } - return applyUnitaryMatrix(ctrlOp, qubits, dd, state); + return applyUnitaryMatrix(ctrlOp, walk, state); }) .template Case( [&](UnitaryOpInterface unitary) -> LogicalResult { @@ -337,10 +1241,9 @@ static LogicalResult applyOp(Operation& op, QubitMap& qubits, dd::Package& dd, return failure(); } if (*decoded) { - return applyDecodedStandard(unitary, **decoded, {}, qubits, dd, - state); + return applyDecodedStandard(unitary, **decoded, {}, walk, state); } - return applyUnitaryMatrix(unitary, qubits, dd, state); + return applyUnitaryMatrix(unitary, walk, state); }) .Default([](Operation* unsupported) { return unsupported->emitError() @@ -350,10 +1253,12 @@ static LogicalResult applyOp(Operation& op, QubitMap& qubits, dd::Package& dd, } template -static LogicalResult walk(func::FuncOp func, QubitMap& qubits, dd::Package& dd, - StateDD& state) { +static LogicalResult walkFunction(func::FuncOp func, WalkState& walkState, + StateDD& state) { + // Function bodies include `func.return` as terminator; region walks skip + // `qco.yield` and bind it separately. for (Operation& op : func.getBody().front()) { - if (failed(applyOp(op, qubits, dd, state))) { + if (failed(applyOp(op, walkState, state))) { return failure(); } } @@ -398,12 +1303,20 @@ FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd) { return failure(); } QubitMap qubits = std::move(*qubitsOr); + ClassicalEnv classical; + DenseSet activeCalls; + WalkState walkState{.qubits = qubits, + .classical = classical, + .dd = dd, + .rng = nullptr, + .classicalBits = nullptr, + .activeCalls = &activeCalls}; dd::MatrixDD state = qubits.numQubits == 0 ? dd::MatrixDD::one() : dd.createInitialMatrix(std::vector(qubits.numQubits, false)); - if (failed(walk(func, qubits, dd, state))) { + if (failed(walkFunction(func, walkState, state))) { if (qubits.numQubits != 0) { dd.decRef(state); } @@ -412,21 +1325,163 @@ FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd) { return state; } -FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, - dd::Package& dd) { +static FailureOr +simulateImpl(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, + std::mt19937_64* rng, std::string* classicalBits) { auto qubitsOr = prepare(func, dd); if (failed(qubitsOr)) { dd.decRef(in); return failure(); } QubitMap qubits = std::move(*qubitsOr); + ClassicalEnv classical; + DenseSet activeCalls; + WalkState walkState{.qubits = qubits, + .classical = classical, + .dd = dd, + .rng = rng, + .classicalBits = classicalBits, + .activeCalls = &activeCalls}; dd::VectorDD state = in; - if (failed(walk(func, qubits, dd, state))) { + if (failed(walkFunction(func, walkState, state))) { dd.decRef(state); return failure(); } return state; } +FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, + dd::Package& dd) { + return simulateImpl(func, in, dd, nullptr, nullptr); +} + +FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, + dd::Package& dd, std::mt19937_64& rng) { + return simulateImpl(func, in, dd, &rng, nullptr); +} + +[[nodiscard]] static bool +requiresDynamicSampling(func::FuncOp func, + DenseSet* visiting = nullptr) { + DenseSet localVisiting; + DenseSet& active = + visiting != nullptr ? *visiting : localVisiting; + Operation* funcOp = func.getOperation(); + if (!active.insert(funcOp).second) { + // Recursive call cycle: treat as dynamic to avoid infinite recursion. + return true; + } + + bool dynamic = false; + func.walk([&](Operation* op) { + // Only stochastic collapse forces per-shot re-simulation. Deterministic + // control-flow can reuse a single simulated state. + if (isa(op)) { + dynamic = true; + return WalkResult::interrupt(); + } + if (auto call = dyn_cast(op)) { + auto module = call->getParentOfType(); + if (!module) { + dynamic = true; + return WalkResult::interrupt(); + } + auto callee = module.lookupSymbol(call.getCallee()); + if (!callee || !callee.getBody().hasOneBlock() || + requiresDynamicSampling(callee, &active)) { + dynamic = true; + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + active.erase(funcOp); + return dynamic; +} + +static FailureOr sampleImpl(func::FuncOp func, + const dd::VectorDD& in, + dd::Package& dd, const size_t shots, + std::mt19937_64& rng, + const bool recordClassics) { + SampleResult result; + if (shots == 0) { + dd.decRef(in); + return result; + } + + if (!requiresDynamicSampling(func)) { + auto stateOr = simulateImpl(func, in, dd, nullptr, nullptr); + if (failed(stateOr)) { + return failure(); + } + dd::VectorDD state = *stateOr; + for (size_t i = 0; i < shots; ++i) { + result.shots[dd.measureAll(state, false, rng)] += 1; + } + dd.decRef(state); + return result; + } + + for (size_t i = 0; i < shots; ++i) { + dd.incRef(in); + std::string classical; + auto stateOr = + simulateImpl(func, in, dd, &rng, recordClassics ? &classical : nullptr); + if (failed(stateOr)) { + dd.decRef(in); + return failure(); + } + dd::VectorDD state = *stateOr; + result.shots[dd.measureAll(state, false, rng)] += 1; + if (recordClassics && !classical.empty()) { + result.classical[classical] += 1; + } + dd.decRef(state); + } + dd.decRef(in); + return result; +} + +FailureOr> +sample(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, + const size_t shots, std::mt19937_64& rng) { + auto result = sampleImpl(func, in, dd, shots, rng, /*recordClassics=*/false); + if (failed(result)) { + return failure(); + } + return std::move(result->shots); +} + +FailureOr sampleWithClassics(func::FuncOp func, + const dd::VectorDD& in, + dd::Package& dd, const size_t shots, + std::mt19937_64& rng) { + return sampleImpl(func, in, dd, shots, rng, /*recordClassics=*/true); +} + +FailureOr> sample(func::FuncOp func, + dd::Package& dd, + const size_t shots, + std::mt19937_64& rng) { + auto qubitsOr = prepare(func, dd); + if (failed(qubitsOr)) { + return failure(); + } + const size_t n = qubitsOr->numQubits; + return sample(func, dd::makeZeroState(n, dd), dd, shots, rng); +} + +FailureOr sampleWithClassics(func::FuncOp func, dd::Package& dd, + const size_t shots, + std::mt19937_64& rng) { + auto qubitsOr = prepare(func, dd); + if (failed(qubitsOr)) { + return failure(); + } + const size_t n = qubitsOr->numQubits; + return sampleWithClassics(func, dd::makeZeroState(n, dd), dd, shots, rng); +} + } // namespace mlir::qco diff --git a/mlir/unittests/Compiler/test_compiler_pipeline.cpp b/mlir/unittests/Compiler/test_compiler_pipeline.cpp index c918894c4b..14f505e6d2 100644 --- a/mlir/unittests/Compiler/test_compiler_pipeline.cpp +++ b/mlir/unittests/Compiler/test_compiler_pipeline.cpp @@ -971,6 +971,33 @@ h q; EXPECT_FALSE(qcoFromString->runPassPipeline("not-a-pass")); EXPECT_FALSE(qcoFromString->str().empty()); + auto entry = qcoFromString->entryFunc(); + ASSERT_TRUE(entry); + EXPECT_EQ(entry->getSymName(), "main"); + + auto firstFuncOnly = QCOProgram::fromMLIRString(R"mlir( +module { + func.func @only() { + %q = qco.static 0 : !qco.qubit + qco.sink %q : !qco.qubit + return + } +} +)mlir"); + ASSERT_TRUE(firstFuncOnly); + auto first = firstFuncOnly->entryFunc(); + ASSERT_TRUE(first); + EXPECT_EQ(first->getSymName(), "only"); + + auto noFunc = QCOProgram::fromMLIRString(R"mlir( +module { + %theta = arith.constant 0.0 : f64 + qco.gphase(%theta) +} +)mlir"); + ASSERT_TRUE(noFunc); + EXPECT_FALSE(noFunc->entryFunc()); + auto baseInput = QCProgram::fromQASMString(qasm); auto adaptiveInput = QCProgram::fromQASMString(qasm); ASSERT_TRUE(baseInput); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp b/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp index 99836cbef5..00108c91b4 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp @@ -9,18 +9,24 @@ */ #include "dd/FunctionalityConstruction.hpp" +#include "dd/GateMatrixDefinitions.hpp" #include "dd/Node.hpp" #include "dd/Package.hpp" #include "dd/Simulation.hpp" #include "dd/StateGeneration.hpp" #include "ir/QuantumComputation.hpp" +#include "ir/operations/OpType.hpp" #include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h" #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/Utils/DDFunctionality.h" #include +#include +#include #include #include +#include +#include #include #include #include @@ -30,14 +36,18 @@ #include #include +#include #include #include #include #include #include #include +#include +#include #include #include +#include using namespace mlir; using namespace qco; @@ -50,14 +60,20 @@ class QCODDFunctionalityTest : public testing::Test { void SetUp() override { DialectRegistry registry; - registry.insert(); + registry.insert(); context = std::make_unique(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); } [[nodiscard]] static func::FuncOp mainFunc(ModuleOp mod) { - return *mod.getBody()->getOps().begin(); + if (auto main = mod.lookupSymbol("main")) { + return main; + } + auto funcs = mod.getBody()->getOps(); + assert(funcs.begin() != funcs.end() && "module must contain a func.func"); + return *funcs.begin(); } template @@ -96,6 +112,79 @@ class QCODDFunctionalityTest : public testing::Test { auto dd = std::make_unique(numQubits); EXPECT_TRUE(failed(buildFunctionality(mainFunc(*mod), *dd))); } + + [[nodiscard]] static dd::VectorDD + basisState(size_t nQubits, llvm::ArrayRef bits, dd::Package& dd) { + return dd::makeBasisState(nQubits, + std::vector(bits.begin(), bits.end()), dd); + } + + [[nodiscard]] static dd::VectorDD oneQubitState(dd::Package& dd) { + return basisState(1, {true}, dd); + } + + static void + expectSimulatesFromZero(func::FuncOp func, size_t nQubits, + llvm::ArrayRef expectedBits, + std::optional seed = std::nullopt) { + auto dd = std::make_unique(nQubits); + auto expected = basisState(nQubits, expectedBits, *dd); + if (seed) { + std::mt19937_64 rng(*seed); + const auto out = + simulate(func, dd::makeZeroState(nQubits, *dd), *dd, rng); + ASSERT_TRUE(succeeded(out)); + EXPECT_EQ(out->getVector(), expected.getVector()); + dd->decRef(*out); + } else { + const auto out = simulate(func, dd::makeZeroState(nQubits, *dd), *dd); + ASSERT_TRUE(succeeded(out)); + EXPECT_EQ(out->getVector(), expected.getVector()); + dd->decRef(*out); + } + dd->decRef(expected); + } + + enum class SampleApi : std::uint8_t { Sample, SampleWithClassics }; + + static void expectSampleHistogram( + func::FuncOp func, size_t nQubits, std::size_t shots, std::uint64_t seed, + StringRef expectedShotKey, SampleApi api = SampleApi::Sample, + std::optional expectedClassicalKey = std::nullopt) { + auto dd = std::make_unique(nQubits); + std::mt19937_64 rng(seed); + if (api == SampleApi::Sample) { + const auto hist = sample(func, *dd, shots, rng); + ASSERT_TRUE(succeeded(hist)); + ASSERT_EQ(hist->size(), 1U); + EXPECT_EQ(hist->begin()->first, expectedShotKey); + EXPECT_EQ(hist->begin()->second, shots); + return; + } + const auto hist = sampleWithClassics(func, *dd, shots, rng); + ASSERT_TRUE(succeeded(hist)); + ASSERT_EQ(hist->shots.size(), 1U); + EXPECT_EQ(hist->shots.begin()->first, expectedShotKey); + EXPECT_EQ(hist->shots.begin()->second, shots); + if (expectedClassicalKey) { + ASSERT_EQ(hist->classical.size(), 1U); + EXPECT_EQ(hist->classical.begin()->first, *expectedClassicalKey); + EXPECT_EQ(hist->classical.begin()->second, shots); + } else { + EXPECT_TRUE(hist->classical.empty()); + } + } + + static void expectBuildAndSimFail(func::FuncOp func, size_t nQubits) { + auto dd = std::make_unique(nQubits); + EXPECT_TRUE(failed(buildFunctionality(func, *dd))); + EXPECT_TRUE(failed(simulate(func, dd::makeZeroState(nQubits, *dd), *dd))); + } + + static void expectSimulateFail(func::FuncOp func, size_t nQubits) { + auto dd = std::make_unique(nQubits); + EXPECT_TRUE(failed(simulate(func, dd::makeZeroState(nQubits, *dd), *dd))); + } }; TEST_F(QCODDFunctionalityTest, MatchesQuantumComputation) { @@ -297,6 +386,34 @@ TEST_F(QCODDFunctionalityTest, DensePaths) { qc.rz(-0.4, 2); expectEqualToQc(mainFunc(*mod), qc); } + { + // Four-qubit dense `inv` on a non-contiguous wire subset (idle q3). + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q0 = b.staticQubit(0); + auto q1 = b.staticQubit(1); + auto q2 = b.staticQubit(2); + auto q3 = b.staticQubit(3); + auto q4 = b.staticQubit(4); + auto outs = + b.inv({q0, q1, q2, q4}, [&](ValueRange t) -> SmallVector { + return {b.rx(0.2, t[0]), b.ry(0.3, t[1]), b.rz(0.4, t[2]), + b.h(t[3])}; + }); + b.sink(outs[0]); + b.sink(outs[1]); + b.sink(outs[2]); + b.sink(q3); + b.sink(outs[3]); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + qc::QuantumComputation qc(5); + qc.rx(-0.2, 0); + qc.ry(-0.3, 1); + qc.rz(-0.4, 2); + qc.h(4); + expectEqualToQc(mainFunc(*mod), qc); + } } TEST_F(QCODDFunctionalityTest, TwoQubitDensePathBeyondFallbackLimit) { @@ -463,6 +580,495 @@ TEST_F(QCODDFunctionalityTest, SimulationConsumesInputReference) { EXPECT_TRUE(zeroQubitDd->getRootSet().empty()); } +TEST_F(QCODDFunctionalityTest, SimulateMeasureCollapsesLikePackage) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.h(b.staticQubit(0)); + std::tie(q, std::ignore) = b.measure(q); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + constexpr std::uint64_t seed = 42; + auto dd = std::make_unique(1); + + std::mt19937_64 refRng(seed); + auto ref = dd::makeZeroState(1, *dd); + ref = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::H), 0), ref); + (void)dd->measureOneCollapsing(ref, 0, refRng); + const auto expected = ref.getVector(); + + std::mt19937_64 rng(seed); + const auto out = + simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd, rng); + ASSERT_TRUE(succeeded(out)); + EXPECT_EQ(out->getVector(), expected); + dd->decRef(*out); + dd->decRef(ref); +} + +TEST_F(QCODDFunctionalityTest, SimulateResetForcesZero) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + q = b.reset(q); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(7); + auto expected = dd::makeZeroState(1, *dd); + const auto out = + simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd, rng); + ASSERT_TRUE(succeeded(out)); + EXPECT_EQ(out->getVector(), expected.getVector()); + dd->decRef(*out); + dd->decRef(expected); +} + +TEST_F(QCODDFunctionalityTest, SimulateIfConstantBranches) { + auto thenMod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + q = b.qcoIf( + true, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + auto elseMod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + q = b.qcoIf( + false, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(thenMod); + ASSERT_TRUE(elseMod); + + // Deterministic constant branches do not require an RNG. + expectSimulatesFromZero(mainFunc(*thenMod), 1, {true}); + expectSimulatesFromZero(mainFunc(*elseMod), 1, {false}); +} + +TEST_F(QCODDFunctionalityTest, SimulateIndexSwitchBranches) { + auto caseMod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + q = b.qcoIndexSwitch(0, q, ArrayRef{0, 1}, + SmallVector>{ + [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }}, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + auto defaultMod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + q = b.qcoIndexSwitch(5, q, ArrayRef{0, 1}, + SmallVector>{ + [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return b.x(arg); }}, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(caseMod); + ASSERT_TRUE(defaultMod); + + expectSimulatesFromZero(mainFunc(*caseMod), 1, {true}); + expectSimulatesFromZero(mainFunc(*defaultMod), 1, {false}); +} + +TEST_F(QCODDFunctionalityTest, SimulateMeasureFeedsIf) { + // |1> measure is deterministic; then-branch identity keeps |1>. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + Value bit; + std::tie(q, bit) = b.measure(q); + q = b.qcoIf( + bit, q, [&](Value arg) { return arg; }, + [&](Value arg) { return b.x(arg); }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + const auto func = mainFunc(*mod); + expectSimulatesFromZero(func, 1, {true}, 99); + expectSampleHistogram(func, 1, 32, 7, "1"); +} + +TEST_F(QCODDFunctionalityTest, SimulateMeasureFeedsIndexSwitch) { + // |1> → measure → index_castui → index_switch case 1 applies X → |0>. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + Value bit; + std::tie(q, bit) = b.measure(q); + auto idx = + arith::IndexCastUIOp::create(b, b.getIndexType(), bit).getResult(); + q = b.qcoIndexSwitch(idx, q, ArrayRef{0, 1}, + SmallVector>{ + [&](Value arg) { return arg; }, + [&](Value arg) { return b.x(arg); }}, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSimulatesFromZero(mainFunc(*mod), 1, {false}, 3); +} + +TEST_F(QCODDFunctionalityTest, SimulateFuncCallAppliesCallee) { + auto mod = parseSourceString(R"mlir( + module { + func.func @apply_x(%q: !qco.qubit) -> !qco.qubit { + %q1 = qco.x %q : !qco.qubit -> !qco.qubit + return %q1 : !qco.qubit + } + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = func.call @apply_x(%q) : (!qco.qubit) -> !qco.qubit + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + auto main = mod->lookupSymbol("main"); + ASSERT_TRUE(main); + + expectSimulatesFromZero(main, 1, {true}); +} + +TEST_F(QCODDFunctionalityTest, RejectsRecursiveFuncCall) { + auto mod = parseSourceString(R"mlir( + module { + func.func @rec(%q: !qco.qubit) -> !qco.qubit { + %q1 = func.call @rec(%q) : (!qco.qubit) -> !qco.qubit + return %q1 : !qco.qubit + } + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = func.call @rec(%q) : (!qco.qubit) -> !qco.qubit + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + auto main = mod->lookupSymbol("main"); + ASSERT_TRUE(main); + + expectSimulateFail(main, 1); +} + +TEST_F(QCODDFunctionalityTest, SimulateScfForAppliesBodyTrips) { + // Three X applications: |0> → |1>. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto results = + b.scfFor(0, 3, 1, ValueRange{q}, + [&](Value /*iv*/, ValueRange iterArgs) -> SmallVector { + return {b.x(iterArgs[0])}; + }); + b.sink(results[0]); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSimulatesFromZero(mainFunc(*mod), 1, {true}); +} + +TEST_F(QCODDFunctionalityTest, AcceptsScfForAtTripCountLimit) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto results = + b.scfFor(0, 10000, 1, ValueRange{q}, + [&](Value /*iv*/, ValueRange iterArgs) -> SmallVector { + return {iterArgs[0]}; + }); + b.sink(results[0]); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSimulatesFromZero(mainFunc(*mod), 1, {false}); +} + +TEST_F(QCODDFunctionalityTest, RejectsScfForTripCountLimit) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto results = + b.scfFor(0, 10001, 1, ValueRange{q}, + [&](Value /*iv*/, ValueRange iterArgs) -> SmallVector { + return {iterArgs[0]}; + }); + b.sink(results[0]); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectBuildAndSimFail(mainFunc(*mod), 1); +} + +TEST_F(QCODDFunctionalityTest, SimulateRicherClassicalArithmetic) { + // idx = (1+2)*3 >> 1 = 4; select(true, idx, 0)=4; cmpi eq 4 → if applies X. + // Also round-trip i1 via extui/trunci. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + auto two = arith::ConstantIndexOp::create(b, 2).getResult(); + auto three = arith::ConstantIndexOp::create(b, 3).getResult(); + auto four = arith::ConstantIndexOp::create(b, 4).getResult(); + auto zero = arith::ConstantIndexOp::create(b, 0).getResult(); + auto sum = arith::AddIOp::create(b, one, two).getResult(); + auto prod = arith::MulIOp::create(b, sum, three).getResult(); + auto shifted = arith::ShRUIOp::create(b, prod, one).getResult(); + auto t = b.boolConstant(true); + auto selected = arith::SelectOp::create(b, t, shifted, zero).getResult(); + auto eq = arith::CmpIOp::create(b, arith::CmpIPredicate::eq, selected, four) + .getResult(); + auto asIndex = arith::ExtUIOp::create(b, b.getIndexType(), eq).getResult(); + auto asBool = + arith::TruncIOp::create(b, b.getI1Type(), asIndex).getResult(); + q = b.qcoIf( + asBool, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + // Exercise subi: 4-4=0 unused for branching but must succeed. + (void)arith::SubIOp::create(b, four, four); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSimulatesFromZero(mainFunc(*mod), 1, {true}); +} + +TEST_F(QCODDFunctionalityTest, SimulateAndiOriXoriShliClassical) { + // Pack two measure bits (from |1>,|0>) as index = bit0 | (bit1 << 1) = 1, + // then switch case 1 applies X on an idle |0> target → |1>. + // Also exercise andi / xori on the measured bits. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q0 = b.x(b.staticQubit(0)); + auto q1 = b.staticQubit(1); + auto q2 = b.staticQubit(2); + Value bit0; + Value bit1; + std::tie(q0, bit0) = b.measure(q0); + std::tie(q1, bit1) = b.measure(q1); + auto i0 = + arith::IndexCastUIOp::create(b, b.getIndexType(), bit0).getResult(); + auto i1 = + arith::IndexCastUIOp::create(b, b.getIndexType(), bit1).getResult(); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + auto shifted = arith::ShLIOp::create(b, i1, one).getResult(); + auto packed = arith::OrIOp::create(b, i0, shifted).getResult(); + auto t = b.boolConstant(true); + auto anded = arith::AndIOp::create(b, bit0, t).getResult(); + // bit0 ^ true flips the measured-1 bit to false; keep the value live. + auto xored = arith::XOrIOp::create(b, anded, t).getResult(); + (void)xored; + q2 = b.qcoIndexSwitch(packed, q2, ArrayRef{0, 1, 2}, + SmallVector>{ + [&](Value arg) { return arg; }, + [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }}, + [&](Value arg) { return arg; }); + b.sink(q0); + b.sink(q1); + b.sink(q2); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSimulatesFromZero(mainFunc(*mod), 3, {true, false, true}, 11); +} + +TEST_F(QCODDFunctionalityTest, SampleUnitaryXIsDeterministic) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + const auto func = mainFunc(*mod); + expectSampleHistogram(func, 1, 64, 1, "1"); + expectSampleHistogram(func, 1, 16, 1, "1", SampleApi::SampleWithClassics); +} + +TEST_F(QCODDFunctionalityTest, SampleWithClassicsRecordsMeasureBits) { + // |1> → measure (bit 1) → if then X → |0>. Classical key "1" every shot. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + Value bit; + std::tie(q, bit) = b.measure(q); + q = b.qcoIf( + bit, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSampleHistogram(mainFunc(*mod), 1, 32, 9, "0", + SampleApi::SampleWithClassics, "1"); +} + +TEST_F(QCODDFunctionalityTest, SimulateClassicalMemRefRegister) { + // measure into memref c[0], then qcoIf loads c[0] and applies X → |0>. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto c = b.allocClassicalBitRegister(1); + auto q = b.x(b.staticQubit(0)); + Value bit; + std::tie(q, bit) = b.measure(q, c, 0); + auto results = b.qcoIf( + c, 0, ValueRange{q}, + [&](ValueRange args) { return SmallVector{b.x(args[0])}; }, + [&](ValueRange args) { return SmallVector{args[0]}; }); + b.sink(results[0]); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSampleHistogram(mainFunc(*mod), 1, 16, 11, "0", + SampleApi::SampleWithClassics, "1"); +} + +TEST_F(QCODDFunctionalityTest, SampleCombinedForMeasureIfIndexSwitch) { + // Drivers-style CF stack on static wires: three X in `scf.for` → |1>, + // measure, keep via `qco.if`, then identity `index_switch`. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto forResults = + b.scfFor(0, 3, 1, ValueRange{q}, + [&](Value /*iv*/, ValueRange iterArgs) -> SmallVector { + return {b.x(iterArgs[0])}; + }); + q = forResults[0]; + Value bit; + std::tie(q, bit) = b.measure(q); + q = b.qcoIf( + bit, q, [&](Value arg) { return arg; }, + [&](Value arg) { return b.x(arg); }); + const auto identity = [](Value arg) { return arg; }; + q = b.qcoIndexSwitch(0, q, ArrayRef{0}, + SmallVector>{identity}, + identity); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSampleHistogram(mainFunc(*mod), 1, 24, 13, "1", + SampleApi::SampleWithClassics, "1"); +} + +TEST_F(QCODDFunctionalityTest, SampleFromInputStateConsumesReference) { + auto unitary = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + b.sink(q); + return b.intConstant(0); + }); + auto withReset = buildModule([](QCOProgramBuilder& b) { + auto q = b.reset(b.staticQubit(0)); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(unitary); + ASSERT_TRUE(withReset); + + auto dd = std::make_unique(1); + auto& roots = dd->getRootSet(); + std::mt19937_64 rng(5); + + // Static path: input |1> sampled without mid-circuit collapse. + for (size_t i = 0; i < 3; ++i) { + auto in = oneQubitState(*dd); + const auto hist = sample(mainFunc(*unitary), in, *dd, /*shots=*/8, rng); + ASSERT_TRUE(succeeded(hist)); + ASSERT_EQ(hist->size(), 1U); + EXPECT_EQ(hist->begin()->first, "1"); + EXPECT_EQ(hist->begin()->second, 8U); + EXPECT_TRUE(roots.empty()); + } + + // Dynamic path: reset forces per-shot re-simulation from input |1|. + for (size_t i = 0; i < 3; ++i) { + auto in = oneQubitState(*dd); + const auto hist = sample(mainFunc(*withReset), in, *dd, /*shots=*/4, rng); + ASSERT_TRUE(succeeded(hist)); + ASSERT_EQ(hist->size(), 1U); + EXPECT_EQ(hist->begin()->first, "0"); + EXPECT_EQ(hist->begin()->second, 4U); + EXPECT_TRUE(roots.empty()); + } + + // shots == 0 still consumes the input reference. + { + auto in = dd::makeZeroState(1, *dd); + const auto hist = sample(mainFunc(*unitary), in, *dd, /*shots=*/0, rng); + ASSERT_TRUE(succeeded(hist)); + EXPECT_TRUE(hist->empty()); + EXPECT_TRUE(roots.empty()); + } +} + +TEST_F(QCODDFunctionalityTest, SampleConstantIfUsesStaticPath) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + q = b.qcoIf( + true, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectSampleHistogram(mainFunc(*mod), 1, 16, 2, "1"); +} + +TEST_F(QCODDFunctionalityTest, SampleHadamardApproximatelyBalanced) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.h(b.staticQubit(0)); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(42); + constexpr std::size_t shots = 2000; + const auto hist = sample(mainFunc(*mod), *dd, shots, rng); + ASSERT_TRUE(succeeded(hist)); + ASSERT_EQ(hist->size(), 2U); + EXPECT_EQ(hist->at("0") + hist->at("1"), shots); + EXPECT_NEAR(static_cast(hist->at("0")), shots / 2.0, 150.0); +} + +TEST_F(QCODDFunctionalityTest, RejectsOutOfRangeShift) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + auto bad = arith::ConstantIndexOp::create(b, 64).getResult(); + auto shifted = arith::ShLIOp::create(b, one, bad).getResult(); + q = b.qcoIndexSwitch( + shifted, q, ArrayRef{0}, + SmallVector>{[&](Value arg) { return arg; }}, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + expectBuildAndSimFail(mainFunc(*mod), 1); +} + TEST_F(QCODDFunctionalityTest, Rejects) { { auto mod = buildModule([](QCOProgramBuilder& b) { @@ -472,10 +1078,17 @@ TEST_F(QCODDFunctionalityTest, Rejects) { return b.intConstant(0); }); ASSERT_TRUE(mod); - auto dd = std::make_unique(1); - EXPECT_TRUE(failed(buildFunctionality(mainFunc(*mod), *dd))); - EXPECT_TRUE( - failed(simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd))); + expectBuildAndSimFail(mainFunc(*mod), 1); + } + + { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q0 = b.reset(b.staticQubit(0)); + b.sink(q0); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + expectBuildAndSimFail(mainFunc(*mod), 1); } { @@ -488,10 +1101,7 @@ TEST_F(QCODDFunctionalityTest, Rejects) { return b.intConstant(0); }); ASSERT_TRUE(mod); - auto dd = std::make_unique(1); - EXPECT_TRUE(failed(buildFunctionality(mainFunc(*mod), *dd))); - EXPECT_TRUE( - failed(simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd))); + expectBuildAndSimFail(mainFunc(*mod), 1); } expectMlirFails(1, R"mlir( @@ -593,4 +1203,505 @@ TEST_F(QCODDFunctionalityTest, Rejects) { EXPECT_TRUE(failed(buildFunctionality(func, *dd))); } +TEST_F(QCODDFunctionalityTest, ClassicalCmpSelectAndIndexBitwise) { + // Hit cmpi predicates (incl. i1), i1 select, and index andi/ori/xori. + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto two = arith::ConstantIndexOp::create(b, 2).getResult(); + auto three = arith::ConstantIndexOp::create(b, 3).getResult(); + auto zero = arith::ConstantIndexOp::create(b, 0).getResult(); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + + auto ne = arith::CmpIOp::create(b, arith::CmpIPredicate::ne, two, three) + .getResult(); + auto slt = arith::CmpIOp::create(b, arith::CmpIPredicate::slt, two, three) + .getResult(); + auto sle = arith::CmpIOp::create(b, arith::CmpIPredicate::sle, two, three) + .getResult(); + auto sgt = arith::CmpIOp::create(b, arith::CmpIPredicate::sgt, three, two) + .getResult(); + auto sge = arith::CmpIOp::create(b, arith::CmpIPredicate::sge, three, two) + .getResult(); + auto ult = arith::CmpIOp::create(b, arith::CmpIPredicate::ult, two, three) + .getResult(); + auto ule = arith::CmpIOp::create(b, arith::CmpIPredicate::ule, two, three) + .getResult(); + auto ugt = arith::CmpIOp::create(b, arith::CmpIPredicate::ugt, three, two) + .getResult(); + auto uge = arith::CmpIOp::create(b, arith::CmpIPredicate::uge, three, two) + .getResult(); + auto all = arith::AndIOp::create( + b, ne, + arith::AndIOp::create( + b, slt, + arith::AndIOp::create( + b, sle, + arith::AndIOp::create( + b, sgt, + arith::AndIOp::create( + b, sge, + arith::AndIOp::create( + b, ult, + arith::AndIOp::create( + b, ule, + arith::AndIOp::create(b, ugt, uge) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult(); + auto t = b.boolConstant(true); + auto f = b.boolConstant(false); + auto selected = arith::SelectOp::create(b, all, t, f).getResult(); + auto i1Ne = arith::CmpIOp::create(b, arith::CmpIPredicate::ne, selected, f) + .getResult(); + // i1 signed preds use sign-extension (true≡-1); unsigned use 0/1. + auto i1Eq = + arith::CmpIOp::create(b, arith::CmpIPredicate::eq, t, t).getResult(); + auto i1Slt = + arith::CmpIOp::create(b, arith::CmpIPredicate::slt, t, f).getResult(); + auto i1Sle = + arith::CmpIOp::create(b, arith::CmpIPredicate::sle, t, t).getResult(); + auto i1Sgt = + arith::CmpIOp::create(b, arith::CmpIPredicate::sgt, f, t).getResult(); + auto i1Sge = + arith::CmpIOp::create(b, arith::CmpIPredicate::sge, f, t).getResult(); + auto i1Ult = + arith::CmpIOp::create(b, arith::CmpIPredicate::ult, f, t).getResult(); + auto i1Ule = + arith::CmpIOp::create(b, arith::CmpIPredicate::ule, t, t).getResult(); + auto i1Ugt = + arith::CmpIOp::create(b, arith::CmpIPredicate::ugt, t, f).getResult(); + auto i1Uge = + arith::CmpIOp::create(b, arith::CmpIPredicate::uge, t, f).getResult(); + auto i1All = + arith::AndIOp::create( + b, i1Ne, + arith::AndIOp::create( + b, i1Eq, + arith::AndIOp::create( + b, i1Slt, + arith::AndIOp::create( + b, i1Sle, + arith::AndIOp::create( + b, i1Sgt, + arith::AndIOp::create( + b, i1Sge, + arith::AndIOp::create( + b, i1Ult, + arith::AndIOp::create( + b, i1Ule, + arith::AndIOp::create(b, i1Ugt, i1Uge) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult()) + .getResult(); + auto orI1 = arith::OrIOp::create(b, i1All, f).getResult(); + + auto masked = arith::AndIOp::create(b, three, one).getResult(); // 1 + auto ored = arith::OrIOp::create(b, masked, zero).getResult(); // 1 + auto xored = arith::XOrIOp::create(b, ored, zero).getResult(); // 1 + auto eqOne = arith::CmpIOp::create(b, arith::CmpIPredicate::eq, xored, one) + .getResult(); + auto cond = arith::AndIOp::create(b, orI1, eqOne).getResult(); + + q = b.qcoIf( + cond, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + expectSimulatesFromZero(mainFunc(*mod), 1, {true}); +} + +TEST_F(QCODDFunctionalityTest, ClassicalBindThroughScfForAndCall) { + // Exercise ClassicalEnv::bindFrom via scf.for iter_args and func.call. + auto mod = parseSourceString(R"mlir( + module { + func.func @flip_if(%q: !qco.qubit, %bit: i1) -> (!qco.qubit, i1) { + %q1 = qco.if %bit args(%qin = %q) -> (!qco.qubit) { + %qx = qco.x %qin : !qco.qubit -> !qco.qubit + qco.yield %qx : !qco.qubit + } else args(%qin = %q) { + qco.yield %qin : !qco.qubit + } + return %q1, %bit : !qco.qubit, i1 + } + func.func @main() { + %q = qco.static 0 : !qco.qubit + %true = arith.constant true + %c0 = arith.constant 0 : index + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %q1, %bit1 = scf.for %iv = %c0 to %c2 step %c1 + iter_args(%qarg = %q, %barg = %true) -> (!qco.qubit, i1) { + %q2, %bout = func.call @flip_if(%qarg, %barg) + : (!qco.qubit, i1) -> (!qco.qubit, i1) + scf.yield %q2, %bout : !qco.qubit, i1 + } + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + // two flips with always-true bit: |0> -> |1> -> |0> + expectSimulatesFromZero(mainFunc(*mod), 1, {false}); +} + +TEST_F(QCODDFunctionalityTest, ScfForZeroTripsAndRejectsBadStep) { + auto zeroTrips = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + auto results = + b.scfFor(3, 3, 1, ValueRange{q}, + [&](Value /*iv*/, ValueRange iterArgs) -> SmallVector { + return {b.h(iterArgs[0])}; + }); + b.sink(results[0]); + return b.intConstant(0); + }); + ASSERT_TRUE(zeroTrips); + expectSimulatesFromZero(mainFunc(*zeroTrips), 1, {true}); + + auto badStep = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto results = + b.scfFor(0, 3, 0, ValueRange{q}, + [&](Value /*iv*/, ValueRange iterArgs) -> SmallVector { + return {iterArgs[0]}; + }); + b.sink(results[0]); + return b.intConstant(0); + }); + ASSERT_TRUE(badStep); + expectBuildAndSimFail(mainFunc(*badStep), 1); +} + +TEST_F(QCODDFunctionalityTest, ClassicalMemRefErrorsAndDealloc) { + auto ok = buildModule([](QCOProgramBuilder& b) { + auto c = b.allocClassicalBitRegister(2); + auto q = b.x(b.staticQubit(0)); + Value bit; + std::tie(q, bit) = b.measure(q, c, 1); + auto loaded = memref::LoadOp::create( + b, c, ValueRange{arith::ConstantIndexOp::create(b, 1)}) + .getResult(); + q = b.qcoIf( + loaded, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + memref::DeallocOp::create(b, c); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(ok); + expectSimulatesFromZero(mainFunc(*ok), 1, {false}, /*seed=*/5); + + // Wrong element type / rank rejected. + expectMlirFails(0, R"mlir( + module { + func.func @main() { + %c = memref.alloc() : memref<2xi32> + memref.dealloc %c : memref<2xi32> + return + } + } + )mlir"); + + auto oob = buildModule([](QCOProgramBuilder& b) { + auto c = b.allocClassicalBitRegister(1); + auto q = b.staticQubit(0); + auto bit = b.boolConstant(true); + memref::StoreOp::create(b, bit, c, + ValueRange{arith::ConstantIndexOp::create(b, 3)}); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(oob); + expectSimulateFail(mainFunc(*oob), 1); + + auto badStoreRank = parseSourceString(R"mlir( + module { + func.func @main() { + %c = memref.alloc() : memref<2x2xi1> + %t = arith.constant true + %i0 = arith.constant 0 : index + memref.store %t, %c[%i0, %i0] : memref<2x2xi1> + memref.dealloc %c : memref<2x2xi1> + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(badStoreRank); + expectSimulateFail(mainFunc(*badStoreRank), 0); +} + +TEST_F(QCODDFunctionalityTest, RejectsUnmappedClassicalAndBadControlFlow) { + // if condition not concrete + auto badIf = parseSourceString(R"mlir( + module { + func.func @main(%cond: i1) { + %q = qco.static 0 : !qco.qubit + %q1 = qco.if %cond args(%qin = %q) -> (!qco.qubit) { + qco.yield %qin : !qco.qubit + } else args(%qin = %q) { + qco.yield %qin : !qco.qubit + } + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(badIf); + expectSimulateFail(mainFunc(*badIf), 1); + + auto badSwitch = parseSourceString(R"mlir( + module { + func.func @main(%idx: index) { + %q = qco.static 0 : !qco.qubit + %q1 = qco.index_switch %idx -> !qco.qubit + default args(%arg0 = %q) { + qco.yield %arg0 : !qco.qubit + } + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(badSwitch); + expectSimulateFail(mainFunc(*badSwitch), 1); + + // Declaration without a body (covers callee lookup / single-block checks). + auto missingBody = parseSourceString(R"mlir( + module { + func.func private @missing(%q: !qco.qubit) -> !qco.qubit + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = func.call @missing(%q) : (!qco.qubit) -> !qco.qubit + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(missingBody); + expectSimulateFail(mainFunc(*missingBody), 1); + + // Unsupported classical op + auto div = parseSourceString(R"mlir( + module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %d = arith.divui %c2, %c1 : index + %q1 = qco.index_switch %d -> !qco.qubit + default args(%arg0 = %q) { + qco.yield %arg0 : !qco.qubit + } + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(div); + expectSimulateFail(mainFunc(*div), 1); + + auto shruiBad = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + auto bad = arith::ConstantIndexOp::create(b, 64).getResult(); + auto shifted = arith::ShRUIOp::create(b, one, bad).getResult(); + q = b.qcoIndexSwitch( + shifted, q, ArrayRef{0}, + SmallVector>{[&](Value arg) { return arg; }}, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(shruiBad); + expectBuildAndSimFail(mainFunc(*shruiBad), 1); +} + +TEST_F(QCODDFunctionalityTest, RejectsDenseFallbackAboveTwelveQubits) { + // 4-qubit inv on a 13-qubit register forces the dense embed path and fails. + auto mod = buildModule([](QCOProgramBuilder& b) { + SmallVector qs; + qs.reserve(13); + for (unsigned i = 0; i < 13; ++i) { + qs.push_back(b.staticQubit(i)); + } + auto outs = b.inv({qs[0], qs[1], qs[2], qs[3]}, + [&](ValueRange t) -> SmallVector { + return {b.x(t[0]), t[1], t[2], t[3]}; + }); + for (unsigned i = 0; i < 4; ++i) { + b.sink(outs[i]); + } + for (unsigned i = 4; i < 13; ++i) { + b.sink(qs[i]); + } + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + expectBuildAndSimFail(mainFunc(*mod), 13); +} + +TEST_F(QCODDFunctionalityTest, ClassicalErrorPathsAndCalleeMeasureSample) { + // Unmapped classical args fail in bindFrom through func.call. + auto unmapped = parseSourceString(R"mlir( + module { + func.func @use(%q: !qco.qubit, %b: i1, %i: index) -> !qco.qubit { + return %q : !qco.qubit + } + func.func @main(%b: i1, %i: index) { + %q = qco.static 0 : !qco.qubit + %q1 = func.call @use(%q, %b, %i) : (!qco.qubit, i1, index) -> !qco.qubit + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(unmapped); + expectSimulateFail(mainFunc(*unmapped), 1); + + // Unsupported classical type for bindFrom (i64). + auto badType = parseSourceString(R"mlir( + module { + func.func @use(%q: !qco.qubit, %x: i64) -> !qco.qubit { + return %q : !qco.qubit + } + func.func @main() { + %q = qco.static 0 : !qco.qubit + %x = arith.constant 1 : i64 + %q1 = func.call @use(%q, %x) : (!qco.qubit, i64) -> !qco.qubit + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(badType); + expectSimulateFail(mainFunc(*badType), 1); + + // Non-index shifts / bad select / bad trunci / bad cmpi result type. + expectMlirFails(1, R"mlir( + module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %t = arith.constant true + %f = arith.constant false + %s = arith.shli %t, %f : i1 + qco.sink %q : !qco.qubit + return + } + } + )mlir"); + expectMlirFails(1, R"mlir( + module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %t = arith.constant true + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %s = arith.select %t, %c0, %c1 : i64 + qco.sink %q : !qco.qubit + return + } + } + )mlir"); + expectMlirFails(1, R"mlir( + module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %c = arith.constant 1 : i64 + %w = arith.trunci %c : i64 to i1 + qco.sink %q : !qco.qubit + return + } + } + )mlir"); + + // Unmapped memref / dynamic alloc. + auto unmappedMem = parseSourceString(R"mlir( + module { + func.func @main(%c: memref<1xi1>) { + %i0 = arith.constant 0 : index + %v = memref.load %c[%i0] : memref<1xi1> + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(unmappedMem); + expectSimulateFail(mainFunc(*unmappedMem), 0); + + auto dynAlloc = parseSourceString(R"mlir( + module { + func.func @main(%n: index) { + %c = memref.alloc(%n) : memref + memref.dealloc %c : memref + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(dynAlloc); + expectSimulateFail(mainFunc(*dynAlloc), 0); + + // IntegerAttr i1 (non-BoolAttr) constant recording + index select. + auto intAttrI1 = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + auto bit = arith::ConstantOp::create(b, IntegerAttr::get(b.getI1Type(), 1)) + .getResult(); + auto zero = arith::ConstantIndexOp::create(b, 0).getResult(); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + auto idx = arith::SelectOp::create(b, bit, one, zero).getResult(); + q = b.qcoIndexSwitch(idx, q, ArrayRef{1}, + SmallVector>{ + [&](Value arg) { return b.x(arg); }}, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(intAttrI1); + expectSimulatesFromZero(mainFunc(*intAttrI1), 1, {true}); + + // Measure inside callee forces dynamic per-shot sampling. + auto calleeMeasure = parseSourceString(R"mlir( + module { + func.func @meas(%q: !qco.qubit) -> !qco.qubit { + %q1, %b = qco.measure %q : !qco.qubit + return %q1 : !qco.qubit + } + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = qco.x %q : !qco.qubit -> !qco.qubit + %q2 = func.call @meas(%q1) : (!qco.qubit) -> !qco.qubit + qco.sink %q2 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(calleeMeasure); + expectSampleHistogram(mainFunc(*calleeMeasure), 1, 8, /*seed=*/9, "1", + SampleApi::SampleWithClassics, + /*expectedClassicalKey=*/StringRef("1")); +} + } // namespace diff --git a/python/mqt/core/mlir.pyi b/python/mqt/core/mlir.pyi index c8162511af..ac3ba13eb3 100644 --- a/python/mqt/core/mlir.pyi +++ b/python/mqt/core/mlir.pyi @@ -15,6 +15,7 @@ from typing import Literal, overload import qiskit +import mqt.core.dd import mqt.core.fomac import mqt.core.ir @@ -505,6 +506,90 @@ class QIRProgram(Program): def write_bitcode(self, path: str | os.PathLike) -> None: """Write this program as LLVM bitcode.""" +class SampleResult: + """Histograms from QCO DD sampling.""" + + @property + def shots(self) -> dict[str, int]: + """Final computational-basis outcome histogram.""" + + @property + def classical(self) -> dict[str, int]: + """Mid-circuit measure-bit histogram (encounter order).""" + +def build_functionality(program: QCOProgram, dd_package: mqt.core.dd.DDPackage) -> mqt.core.dd.MatrixDD: + """Build a matrix DD for a static unitary QCO program. + + Args: + program: A QCO program whose entry ``func.func`` is used to build a matrix DD. + 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. + """ + +def simulate( + program: QCOProgram, initial_state: mqt.core.dd.VectorDD, dd_package: mqt.core.dd.DDPackage, seed: int | None = None +) -> mqt.core.dd.VectorDD: + """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. + """ + +def sample( + program: QCOProgram, dd_package: mqt.core.dd.DDPackage, shots: int = 1024, seed: int | None = None +) -> dict[str, int]: + """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. Not thread-safe; + do not share it across threads while sampling (the GIL is released for + the duration of the call). + 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. + """ + +def sample_with_classics( + program: QCOProgram, dd_package: mqt.core.dd.DDPackage, shots: int = 1024, seed: int | None = None +) -> SampleResult: + """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. Not thread-safe; + do not share it across threads while sampling (the GIL is released for + the duration of the call). + 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. + """ + @overload def compile_program( program: str diff --git a/test/python/test_qco_dd.py b/test/python/test_qco_dd.py new file mode 100644 index 0000000000..d7b1400764 --- /dev/null +++ b/test/python/test_qco_dd.py @@ -0,0 +1,110 @@ +# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +# All rights reserved. +# +# SPDX-License-Identifier: MIT +# +# Licensed under the MIT License + +"""Tests for QCO DD Python bindings.""" + +from __future__ import annotations + +import numpy as np +import pytest + +dd = pytest.importorskip("mqt.core.dd") +mlir = pytest.importorskip("mqt.core.mlir") + + +def _x_program() -> mlir.QCOProgram: + return mlir.QCOProgram.from_mlir_str(""" +module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = qco.x %q : !qco.qubit -> !qco.qubit + qco.sink %q1 : !qco.qubit + return + } +} +""") + + +def _measure_program() -> mlir.QCOProgram: + return mlir.QCOProgram.from_mlir_str(""" +module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = qco.x %q : !qco.qubit -> !qco.qubit + %q2, %bit = qco.measure %q1 : !qco.qubit + %q3 = qco.if %bit args(%q_in = %q2) -> (!qco.qubit) { + %qx = qco.x %q_in : !qco.qubit -> !qco.qubit + qco.yield %qx : !qco.qubit + } else args(%q_in = %q2) { + qco.yield %q_in : !qco.qubit + } + qco.sink %q3 : !qco.qubit + return + } +} +""") + + +def test_unitary_x_build_simulate_and_sample() -> None: + """X on |0>: unitary matrix, simulation to |1>, deterministic sampling.""" + program = _x_program() + package = dd.DDPackage(1) + matrix = mlir.build_functionality(program, package) + package.dec_ref_mat(matrix) + + zero = package.zero_state(1) + out = mlir.simulate(program, zero, package) + expected = package.computational_basis_state(1, [True]) + assert np.allclose(out.get_vector(), expected.get_vector()) + package.dec_ref_vec(out) + package.dec_ref_vec(expected) + + assert mlir.sample(program, package, shots=32, seed=1) == {"1": 32} + result = mlir.sample_with_classics(program, package, shots=16, seed=2) + assert result.shots == {"1": 16} + assert result.classical == {} + + +def test_simulate_measure_requires_seed() -> None: + """Simulate without seed rejects measure/reset; with seed it succeeds.""" + program = _measure_program() + package = dd.DDPackage(1) + + zero = package.zero_state(1) + with pytest.raises(ValueError, match=r"cannot simulate|measure"): + mlir.simulate(program, zero, package) + + zero = package.zero_state(1) + out = mlir.simulate(program, zero, package, seed=3) + expected = package.computational_basis_state(1, [False]) + assert np.allclose(out.get_vector(), expected.get_vector()) + package.dec_ref_vec(out) + package.dec_ref_vec(expected) + + +def test_entry_func_required() -> None: + """Programs without a func.func raise ValueError via entryFunc.""" + # Top-level qco op satisfies dialect checks but provides no entry function. + program = mlir.QCOProgram.from_mlir_str(""" +module { + %theta = arith.constant 0.0 : f64 + qco.gphase(%theta) +} +""") + package = dd.DDPackage(1) + with pytest.raises(ValueError, match=r"no func\.func"): + mlir.build_functionality(program, package) + + +def test_sample_with_classics_records_midcircuit_measure() -> None: + """Measure then classically controlled X records classical bit '1'.""" + program = _measure_program() + package = dd.DDPackage(1) + result = mlir.sample_with_classics(program, package, shots=20, seed=3) + assert result.shots == {"0": 20} + assert result.classical == {"1": 20}