diff --git a/CHANGELOG.md b/CHANGELOG.md index 6e46ff7146..9cbd2d5db3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,10 +29,15 @@ releases may include breaking changes. [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**], [**@J4MMlE**]) - ✨ Add decision diagram-based construction, simulation, and sampling for QCO - programs, including mid-circuit `measure` / `reset`, concrete `if` / - `index_switch` / `scf.for` / `func.call`, classical SSA and CBit registers, - dense multi-wire embedding, output-aware multi-shot sampling, and Python - bindings ([#1915], [#1973], [#2077]) ([**@simon1hofmann**]) + programs, including mid-circuit `measure` / `reset`, concrete QCO, SCF, and + multi-block CFG control flow, non-recursive calls, bound parameters, classical + integer, floating-point, and common math operations, CBit registers, + one-dimensional memrefs, dynamic quantum allocation and separable + deallocation, qtensors, dense multi-wire embedding, Python input bindings, + optional sampling input states, output-aware multi-shot sampling, + density-matrix simulation and sampling with physical partial trace, and Python + bindings ([#1915], [#1973], [#2077], [#2078], [#2079], [#2080]) + ([**@simon1hofmann**], [**@burgholzer**]) - ✨ Add immutable MLIR compiler targets, QDMI device integration, and target compilation through C++, Python, and `mqt-cc` ([#1687], [#1993], [#1999], [#2049]) ([**@MatthiasReumann**], [**@simon1hofmann**], [**@burgholzer**]) @@ -131,8 +136,10 @@ releases may include breaking changes. - 💥 Remove `MQT::CoreAlgorithms`, its fixed-circuit factories, and the legacy DD package evaluation. MQT Core provides no direct replacement ([#2214]) ([**@burgholzer**]) -- 💥 Remove the unowned decision-diagram approximation algorithm and - density-matrix support from MQT Core ([#1466], [#2154]) ([**@burgholzer**]) +- 💥 Remove the unowned decision-diagram approximation algorithm and the legacy + raw density-matrix and noise APIs. Compiler-backed QCO density simulation + remains available through generic matrix DDs ([#1466], [#2154]) + ([**@burgholzer**]) - 💥 Make `nlohmann_json` an implementation detail and replace JSON-typed decision-diagram statistics APIs with strings and streams ([#2138]) ([**@denialhaag**]) @@ -914,6 +921,9 @@ for previous changelogs._ [#2105]: https://github.com/munich-quantum-toolkit/core/pull/2105 [#2084]: https://github.com/munich-quantum-toolkit/core/pull/2084 [#2082]: https://github.com/munich-quantum-toolkit/core/pull/2082 +[#2080]: https://github.com/munich-quantum-toolkit/core/pull/2080 +[#2079]: https://github.com/munich-quantum-toolkit/core/pull/2079 +[#2078]: https://github.com/munich-quantum-toolkit/core/pull/2078 [#2077]: https://github.com/munich-quantum-toolkit/core/pull/2077 [#2074]: https://github.com/munich-quantum-toolkit/core/pull/2074 [#2066]: https://github.com/munich-quantum-toolkit/core/pull/2066 diff --git a/UPGRADING.md b/UPGRADING.md index f8bb98513a..895e995cb0 100644 --- a/UPGRADING.md +++ b/UPGRADING.md @@ -105,16 +105,19 @@ The Python bindings depend on `nanobind-backend`, which supplies the interpreter-specific nanobind runtime. This dependency does not change the C++ API or the Python import paths. -### Removal of DD approximation and density-matrix support +### Removal of DD approximation and legacy density/noise APIs MQT Core no longer provides the decision-diagram approximation algorithm. The algorithm had no production owner in the MQT ecosystem. Remove uses of the `dd/Approximation.hpp` header, the `dd::ApproximationMetadata` type, and the `dd::approximate` function. MQT Core does not provide a replacement. -MQT Core also no longer provides density-matrix decision diagrams or the noise -operations that depended on them. Consumers must provide this functionality or -use another implementation. +MQT Core also no longer provides the legacy raw density-matrix DD types or the +noise operations that depended on them. Compiler-backed QCO density simulation +remains available through generic matrix DDs: use +`mqt.core.mlir.make_density_matrix`, then call `QCOProgram.simulate_density` or +`QCOProgram.sample_density`. Consumers of the removed raw APIs must migrate to +QCO or another implementation. ### Private `nlohmann_json` dependency diff --git a/bindings/mlir/register_mlir.cpp b/bindings/mlir/register_mlir.cpp index e14f43c696..c3cf138842 100644 --- a/bindings/mlir/register_mlir.cpp +++ b/bindings/mlir/register_mlir.cpp @@ -8,12 +8,14 @@ * Licensed under the MIT License */ +#include "dd/Edge.hpp" #include "dd/Node.hpp" #include "dd/Package.hpp" #include "mlir/Compiler/Programs.h" #include "mlir/Compiler/QDMIAdapter.h" #include "mlir/Compiler/Target.h" #include "mlir/Dialect/MQT/IR/MQTDialect.h" +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/Utils/DDFunctionality.h" #include "qdmi/Client.hpp" #include "qdmi/driver/SessionConfig.hpp" @@ -22,6 +24,8 @@ #include #include #include +#include +#include #include #include #include @@ -39,6 +43,7 @@ #include #include #include +#include #include #include #include @@ -48,6 +53,7 @@ #include #include #include +#include #include namespace mqt { @@ -139,6 +145,66 @@ entryFunc(const mlir::QCOProgram& program) { return func; } +using QCODDBindingValue = std::variant; +using QCODDBindingMap = std::map; + +[[nodiscard]] static mlir::qco::DDBindings +makeQCODDBindings(mlir::func::FuncOp func, + const QCODDBindingMap& pythonBindings) { + mlir::qco::DDBindings bindings; + for (const auto& [index, binding] : pythonBindings) { + if (index >= func.getNumArguments()) { + throw nb::value_error("QCO DD binding argument index is out of range"); + } + + mlir::Value argument = func.getArgument(static_cast(index)); + const mlir::Type type = argument.getType(); + mlir::Attribute attribute; + if (type.isInteger(1)) { + if (const auto* value = std::get_if(&binding)) { + attribute = mlir::BoolAttr::get(func.getContext(), *value); + } + } else if (mlir::isa(type)) { + if (const auto* value = std::get_if(&binding)) { + attribute = mlir::IntegerAttr::get(type, *value); + } + } else if (const auto floatType = mlir::dyn_cast(type)) { + if (const auto* value = std::get_if(&binding)) { + attribute = mlir::FloatAttr::get(floatType, *value); + } + } else if (const auto tensorType = + mlir::dyn_cast(type); + tensorType && tensorType.getRank() == 1 && + tensorType.isDynamicDim(0) && + mlir::isa(tensorType.getElementType())) { + if (const auto* value = std::get_if(&binding); + value != nullptr && *value >= 0) { + attribute = mlir::IntegerAttr::get( + mlir::IndexType::get(func.getContext()), *value); + } + } + + if (!attribute) { + throw nb::value_error( + "QCO DD binding value does not match the entry argument type"); + } + bindings[argument] = attribute; + } + return bindings; +} + +template +static void requireLiveReference(const dd::Edge& state, + dd::Package& ddPackage, + const char* argumentName = "initial_state") { + if (dd::Edge::trackingRequired(state) && + !ddPackage.getRootSet().contains(state)) { + std::string message(argumentName); + message.append(" must have a live reference in dd_package"); + throw nb::value_error(message.c_str()); + } +} + [[nodiscard]] static std::mt19937_64 makeRng(const uint64_t seed) { if (seed == 0) { return std::mt19937_64(std::random_device{}()); @@ -993,20 +1059,24 @@ LLVM bitcode.)pb"); qcoProgram.def( "build_functionality", - [](const mlir::QCOProgram& program, dd::Package& ddPackage) { + [](const mlir::QCOProgram& program, dd::Package& ddPackage, + const QCODDBindingMap& pythonBindings) { auto func = entryFunc(program); + const auto bindings = makeQCODDBindings(func, pythonBindings); return takeFailureOr( func.getContext(), - "cannot build DD functionality for this QCO program", - [&] { return mlir::qco::buildFunctionality(func, ddPackage); }); + "cannot build DD functionality for this QCO program", [&] { + return mlir::qco::buildFunctionality(func, ddPackage, bindings); + }); }, - "dd_package"_a, + "dd_package"_a, nb::kw_only(), "bindings"_a = QCODDBindingMap{}, // Keep the DD package alive while the returned matrix DD is alive. nb::keep_alive<0, 2>(), R"pb(Build a matrix DD for a static unitary QCO program. Args: dd_package: DD package with enough qubits for the program. + bindings: Concrete entry-argument values keyed by zero-based argument index. Returns: Matrix DD of the program functionality. @@ -1017,20 +1087,20 @@ LLVM bitcode.)pb"); qcoProgram.def( "simulate", [](const mlir::QCOProgram& program, const dd::VectorDD& initialState, - dd::Package& ddPackage, const uint64_t seed) { - if (dd::VectorDD::trackingRequired(initialState) && - !ddPackage.getRootSet().contains(initialState)) { - throw nb::value_error( - "initial_state must have a live reference in dd_package"); - } + dd::Package& ddPackage, const uint64_t seed, + const QCODDBindingMap& pythonBindings) { + requireLiveReference(initialState, ddPackage); auto func = entryFunc(program); + const auto bindings = makeQCODDBindings(func, pythonBindings); auto rng = makeRng(seed); return takeFailureOr( func.getContext(), "cannot simulate this QCO program", [&] { - return mlir::qco::simulate(func, initialState, ddPackage, rng); + return mlir::qco::simulate(func, initialState, ddPackage, rng, + bindings); }); }, - "initial_state"_a, "dd_package"_a, "seed"_a = 0U, + "initial_state"_a, "dd_package"_a, "seed"_a = 0U, nb::kw_only(), + "bindings"_a = QCODDBindingMap{}, // 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. @@ -1042,6 +1112,7 @@ LLVM bitcode.)pb"); dd_package: DD package with enough qubits for the program. seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other value produces reproducible measurement and reset results. + bindings: Concrete entry-argument values keyed by zero-based argument index. Returns: Output state DD. @@ -1050,17 +1121,93 @@ LLVM bitcode.)pb"); ValueError: When ``initial_state`` has no live reference in ``dd_package``, has too few qubits, or the program is unsupported for simulation.)pb"); + m.def( + "make_density_matrix", + [](const dd::VectorDD& state, const size_t numQubits, + dd::Package& ddPackage) { + requireLiveReference(state, ddPackage, "state"); + try { + return mlir::qco::makeDensityMatrix(state, numQubits, ddPackage); + } catch (const std::invalid_argument& error) { + throw nb::value_error(error.what()); + } + }, + "state"_a, "num_qubits"_a, "dd_package"_a, nb::keep_alive<0, 3>(), + R"pb(Construct ``|psi>(), + R"pb(Simulate a QCO program on a density-matrix DD. + +Args: + initial_state: Input density-matrix DD with a live reference in + ``dd_package``. It represents exactly the program's inferred initial + quantum register; skipped DD levels denote identity factors within + that register. A valid input reference is consumed. + dd_package: DD package with enough qubits for the program. + seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other + value produces reproducible measurement and reset results. + bindings: Concrete entry-argument values keyed by zero-based argument index. + +Returns: + Output density-matrix DD. + +Raises: + ValueError: When ``initial_state`` has no live reference in ``dd_package`` + or the program is unsupported for simulation.)pb"); + qcoProgram.def( "sample", [](const mlir::QCOProgram& program, dd::Package& ddPackage, - const size_t shots, const uint64_t seed) { + const size_t shots, const uint64_t seed, + const std::optional& initialState, + const QCODDBindingMap& pythonBindings) { auto func = entryFunc(program); + const auto bindings = makeQCODDBindings(func, pythonBindings); auto rng = makeRng(seed); return takeFailureOr( - func.getContext(), "cannot sample this QCO program", - [&] { return mlir::qco::sample(func, ddPackage, shots, rng); }); + func.getContext(), "cannot sample this QCO program", [&] { + if (initialState) { + requireLiveReference(*initialState, ddPackage); + return mlir::qco::sample(func, *initialState, ddPackage, shots, + rng, bindings); + } + return mlir::qco::sample(func, ddPackage, shots, rng, bindings); + }); }, - "dd_package"_a, "shots"_a = 1024U, "seed"_a = 0U, + "dd_package"_a, "shots"_a = 1024U, "seed"_a = 0U, nb::kw_only(), + "initial_state"_a = nb::none(), "bindings"_a = QCODDBindingMap{}, R"pb(Sample the declared outputs of a QCO program. Args: @@ -1068,13 +1215,55 @@ LLVM bitcode.)pb"); shots: Number of shots (default 1024). seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other value produces reproducible results. + initial_state: Optional input state with a live reference in ``dd_package``. + A valid input reference is consumed. + bindings: Concrete entry-argument values keyed by zero-based argument index. + +Returns: + Histogram of returned CBit registers in return order, each MSB first. If + no CBit result exists, final ``measureAll`` bitstrings instead. + +Raises: + ValueError: When ``initial_state`` has no live reference in ``dd_package`` + or the program is unsupported for sampling.)pb"); + + qcoProgram.def( + "sample_density", + [](const mlir::QCOProgram& program, const dd::MatrixDD& initialState, + dd::Package& ddPackage, const size_t shots, const uint64_t seed, + const QCODDBindingMap& pythonBindings) { + requireLiveReference(initialState, ddPackage); + auto func = entryFunc(program); + const auto bindings = makeQCODDBindings(func, pythonBindings); + auto rng = makeRng(seed); + return takeFailureOr( + func.getContext(), "cannot density-sample this QCO program", [&] { + return mlir::qco::sampleDensity(func, initialState, ddPackage, + shots, rng, bindings); + }); + }, + "initial_state"_a, "dd_package"_a, "shots"_a = 1024U, "seed"_a = 0U, + nb::kw_only(), "bindings"_a = QCODDBindingMap{}, + R"pb(Sample the declared outputs of a QCO program from a density-matrix DD. + +Args: + initial_state: Input density-matrix DD with a live reference in + ``dd_package``. It represents exactly the program's inferred initial + quantum register; skipped DD levels denote identity factors within + that register. A valid input reference is consumed. + dd_package: DD package with enough qubits for the program. + shots: Number of shots (default 1024). + seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other + value produces reproducible results. + bindings: Concrete entry-argument values keyed by zero-based argument index. Returns: Histogram of returned CBit registers in return order, each MSB first. If no CBit result exists, final ``measureAll`` bitstrings instead. Raises: - ValueError: When the program is unsupported for sampling.)pb"); + ValueError: When ``initial_state`` has no live reference in ``dd_package`` + or 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, diff --git a/docs/dd_package.md b/docs/dd_package.md index 83186f5c8d..e69bd1fdd7 100644 --- a/docs/dd_package.md +++ b/docs/dd_package.md @@ -184,6 +184,11 @@ zero_state_dd = dd.zero_state(qc.num_qubits) out_state_dd = simulate(qc, zero_state_dd, dd) ``` +Density simulation uses generic matrix DDs. Construct a pure density operator +with {py:func}`~mqt.core.mlir.make_density_matrix`, then use +{py:meth}`~mqt.core.mlir.QCOProgram.simulate_density` or +{py:meth}`~mqt.core.mlir.QCOProgram.sample_density`. + If the [Graphviz](https://www.graphviz.org/) library is installed, the `graphviz` Python package can be used to visualize resulting decision diagram via the {py:meth}`~mqt.core.dd.VectorDD.to_dot` method. To directly, generate diff --git a/include/mqt-core/dd/Package.hpp b/include/mqt-core/dd/Package.hpp index 1faa104bab..6517199b80 100644 --- a/include/mqt-core/dd/Package.hpp +++ b/include/mqt-core/dd/Package.hpp @@ -1463,7 +1463,7 @@ class Package { * mapped to the interval [0,1] (as opposed to the interval [0,2^N]). */ mCachedEdge trace(const mEdge& a, const std::vector& eliminate, - std::size_t level, std::size_t alreadyEliminated = 0); + const std::vector& keptBefore); /** * @brief Recursively checks if a given matrix is close to the identity diff --git a/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h b/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h index f3506f845a..67265cb28d 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/DDFunctionality.h @@ -12,7 +12,10 @@ #include "dd/Package_fwd.hpp" +#include #include +#include +#include #include #include @@ -22,39 +25,107 @@ namespace mlir::qco { +/** + * @brief Concrete values for symbolic QCO DD inputs. + * + * Integer and `f64` attributes bind scalar function arguments. An integer + * attribute bound to a dynamic one-dimensional qtensor argument gives its + * runtime extent. Bindings for other values are rejected. + */ +using DDBindings = DenseMap; + +/** + * @brief Determine the DD package capacity required by a concrete QCO function. + * + * @details Counts static/input wires and statically sized entry-block quantum + * allocations. Dynamic QTensor extents, allocations outside the entry block, + * and calls that may allocate qubits are rejected because their peak capacity + * cannot be determined without executing the program. + * + * @param func The QCO function to inspect + * @return The required number of DD wires, or failure when the capacity cannot + * be determined statically + */ +FailureOr getNumQubits(func::FuncOp func); + /** * @brief Sequentially build a matrix DD for a static unitary QCO `func.func`. * - * @details Walks the entry block of @p func, maps `qco.static` SSA values to - * wire indices (or, if none are present, qubit-typed function arguments as - * wires `0..n-1`), assigns entry-block `qco.alloc` operations subsequent - * wires, and applies unitary operations via decision-diagram multiplication. + * @details Walks the concrete control-flow path through @p func, maps + * `qco.static` SSA values to wire indices (or, if none are present, + * qubit-typed function arguments as wires `0..n-1`), assigns entry-block + * `qco.alloc` operations subsequent wires, and applies unitary operations via + * decision-diagram multiplication. * * Supported programs: - * - Standard single-, two-, and three-qubit gates with compile-time constant + * - Standard single-, two-, and three-qubit gates with constant or bound * parameters (sparse DD path) * - `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` + * - QTensor bookkeeping over existing input wires + * - Concrete QCO and SCF control flow, multi-block ControlFlow CFGs, and + * non-recursive calls + * - Concrete integer, index, `f64`, and common Math operations and + * one-dimensional memrefs of those scalar types * - `qco.static` establishes the wire map (or qubit-typed `func` args if none), - * followed by entry-block `qco.alloc`; `sink` is ignored; `arith.constant` - * is ignored for matrix construction; `func.return` accepts qubit results - * only in canonical wire order + * followed by entry-block `qco.alloc`; `sink` is ignored; returned qubits + * and qtensors must preserve canonical wire order * * Known one-, two-, and three-qubit matrices are constructed directly as DD * gates. Larger compile-time unitaries are embedded directly into a DD over * their target wires, so idle register qubits do not enlarge the local matrix. - * Measurements, resets, symbolic parameters, and control-flow ops are not - * supported. + * Measurements, resets, unbound parameters, and non-concrete control flow are + * not supported. * * @pre The containing module has passed MLIR verification and * `qco::verifyLinearity`. * * @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) + * @param bindings Concrete scalar values and dynamic QTensor extents for entry + * arguments * @return The matrix DD on success, or failure for unsupported programs */ -FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd); +FailureOr +buildFunctionality(func::FuncOp func, dd::Package& dd, + const DDBindings& bindings = DDBindings()); + +/** + * @brief Simulate a QCO `func.func` on a given input state without stochastic + * collapse. + * + * @details Same supported unitary op set as @ref buildFunctionality, plus + * concrete QCO and standard SCF control flow and static- or concrete + * dynamic-shape one-dimensional memrefs of integer, index, or `f64` values and + * CBit registers. `qco.alloc` and `qtensor.alloc` append zero-state + * wires. QTensor + * extraction, insertion, deallocation, and transport through regions are + * tracked with linear value semantics. Deallocating a separable QTensor + * removes its wires from vector DDs; deallocating an entangled wire is + * rejected. QTensor sizes and indices must be concrete; dynamic qtensor + * arguments require an extent in @p bindings. + * Mid-circuit `measure` / `reset` require the RNG overload below. Concrete- + * bound `scf.for` and `scf.while` loops, multi-block `scf.execute_region`, and + * non-recursive multi-block `func.call` are supported independently of RNG. + * A shared 10000-step budget bounds loop iterations and CFG transitions across + * nested regions and calls. + * Consumes one reference to @p in regardless of success or failure. + * + * @pre The containing module has passed MLIR verification and + * `qco::verifyLinearity`. + * + * @param func The QCO function to simulate + * @param in The input state, which must span at least the function's qubits; + * higher wires are preserved; one reference is consumed + * @param dd The DD package to use (must hold at least the function's qubits) + * @param bindings Concrete values for symbolic function arguments + * @return The output statevector DD on success, or failure for unsupported + * programs + */ +FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, + dd::Package& dd, + const DDBindings& bindings = DDBindings()); /** * @brief Simulate a QCO `func.func` that may contain measurements, resets, and @@ -63,17 +134,11 @@ FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd); * @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` integer/index, a prior measurement, a `cbit.load`, - * `arith.extui`, `arith.index_castui`, `arith.cmpi`, `arith.select`, - * `arith.addi` / `subi` / `muli`, or `andi` / `ori` / `xori` / `shli` / - * `shrui` on those values). The simulation tracks CBit initialization, loads, - * and stores. Only qubit-typed linear values are supported (no qtensors). - * Nested regions are walked; direct `scf.for` execution with concrete positive - * steps and non-recursive single-block `func.call` are supported. A shared - * 10000-step budget bounds loop iterations across nested loops and calls; - * `scf.while` and multi-block function bodies remain unsupported. - * Consumes one reference to @p in regardless of whether simulation succeeds or - * fails. + * (`arith.constant`, a prior measurement, integer and `f64` + * arithmetic, comparisons, casts, shifts, and `arith.select`). Dynamic quantum + * allocation, qtensors, memrefs, CBit registers, loops, regions, and calls are + * supported as in the non-RNG overload. + * Consumes one reference to @p in regardless of success or failure. * * @pre The containing module has passed MLIR verification and * `qco::verifyLinearity`. @@ -83,11 +148,76 @@ FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd); * higher wires are preserved; one reference is consumed * @param dd The DD package to use * @param rng RNG used for collapsing measurements and resets + * @param bindings Concrete scalar values and dynamic QTensor extents for entry + * arguments * @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); + dd::Package& dd, std::mt19937_64& rng, + const DDBindings& bindings = DDBindings()); + +/** + * @brief Simulate a QCO function for state extraction. + * + * @details Terminal measurements that only populate returned CBit registers + * are deferred and do not collapse the returned state. Quantum deallocations + * are treated as lifetime markers so the complete circuit state is retained. + * Mid-circuit measurements and resets are executed with @p rng. Use + * @ref getNumQubits to size @p dd; programs whose peak capacity cannot be + * determined statically are rejected. + * + * @param func The QCO function to simulate + * @param dd The DD package to use + * @param rng RNG used for non-terminal measurements and resets + * @return The output statevector DD on success + */ +FailureOr simulateStatevector(func::FuncOp func, dd::Package& dd, + std::mt19937_64& rng); + +/** + * @brief Construct the density operator @f$|\psi\rangle\langle\psi|@f$. + * + * @param state Pure input state; its reference is retained by the caller + * @param numQubits Number of active qubits represented by @p state + * @param dd The DD package to use + * @return A referenced matrix DD representing the pure-state density operator + * @throws std::invalid_argument If @p numQubits does not cover the highest DD + * level in @p state or exceeds the capacity of @p dd + */ +dd::MatrixDD makeDensityMatrix(const dd::VectorDD& state, size_t numQubits, + dd::Package& dd); + +/** + * @brief Simulate a QCO function using a density-matrix DD. + * + * @details Unitary operations evolve the state as @f$U\rho U^\dagger@f$. + * Qubit and qtensor deallocation performs a physical partial trace, including + * for entangled qubits. The RNG overload additionally supports collapsing + * measurement and reset. Consumes one reference to @p in regardless of + * success or failure. The input represents exactly the function's inferred + * initial quantum register. Matrix DDs do not encode a logical extent, so + * skipped levels are interpreted as identity factors within that register. + * + * @param func The QCO function to simulate + * @param in Input density matrix over the inferred initial quantum register; + * one reference is consumed + * @param dd The DD package to use + * @param bindings Concrete values for symbolic function arguments + * @return The output density-matrix DD on success, or failure for unsupported + * programs + */ +FailureOr +simulateDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd, + const DDBindings& bindings = DDBindings()); + +/// @copydoc simulateDensity(func::FuncOp, const dd::MatrixDD&, dd::Package&, +/// const DDBindings&) +/// Uses @p rng for collapsing measurement and reset. +FailureOr +simulateDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd, + std::mt19937_64& rng, + const DDBindings& bindings = DDBindings()); /** * @brief Sample measurement outcomes from a QCO `func.func`. @@ -99,7 +229,11 @@ FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, * basis sampling via `Package::measureAll` (qubit `n-1` … `0`). Terminal entry- * block measurements that only produce returned CBit cells are sampled from * one DD evolution; resets and execution-dependent measurements are executed - * once per shot. + * once per shot. With returned CBit registers, deallocations needed to encode + * the result are treated as lifetime markers. Deallocated separable QTensor + * wires are omitted from fallback-basis outcomes. Multi-block functions are + * executed once per shot and support fallback-basis sampling only, not CBit + * return values. * * @pre The containing module has passed MLIR verification and * `qco::verifyLinearity`. @@ -108,10 +242,72 @@ FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, * @param dd The DD package to use * @param shots Number of shots * @param rng RNG for collapsing measurements and non-collapsing sampling + * @param bindings Concrete scalar values and dynamic QTensor extents for entry + * arguments * @return Histogram of outcome strings on success, or failure for unsupported * programs */ FailureOr> -sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng); +sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng, + const DDBindings& bindings = DDBindings()); +/** + * @brief Sample every allocated qubit of a QCO function. + * + * @details Same as @ref sample, but when the function has no returned CBit + * registers, quantum deallocations are treated as end-of-program lifetime + * markers. This preserves the full-width output distribution expected by + * external circuit formats that lower terminal cleanup to `qtensor.dealloc`. + * + * @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 + * @param bindings Concrete values for symbolic function arguments + * @return Histogram of full-width outcome strings on success, or failure for + * unsupported programs + */ +FailureOr> +sampleAllQubits(func::FuncOp func, dd::Package& dd, size_t shots, + std::mt19937_64& rng, + const DDBindings& bindings = DDBindings()); + +/** + * @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 regardless of success, failure, or the number of shots. + * The non-dynamic path evolves the input once; the dynamic path clones it for + * each 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 + * @param bindings Concrete values for symbolic function arguments + * @return Histogram of outcome strings on success, or failure for unsupported + * programs + */ +FailureOr> +sample(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, size_t shots, + std::mt19937_64& rng, const DDBindings& bindings = DDBindings()); + +/** + * @brief Sample a QCO function from an input density-matrix DD. + * + * @details Supports mixed states and entangled qubit deallocation. Outcome + * encoding follows @ref sample: returned CBit registers take precedence over + * final computational-basis sampling. Each final sample collapses a referenced + * copy of the simulated density state. Programs with mid-circuit measurement + * or reset are re-simulated per shot. Consumes one reference to @p in + * regardless of success, failure, or @p shots. The input represents exactly + * the function's inferred initial quantum register. Matrix DDs do not encode a + * logical extent, so skipped levels are interpreted as identity factors within + * that register. + */ +FailureOr> +sampleDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd, + size_t shots, std::mt19937_64& rng, + const DDBindings& bindings = DDBindings()); } // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt b/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt index ddccd4eeb8..a3e45c5f0a 100644 --- a/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt +++ b/mlir/lib/Dialect/QCO/Utils/CMakeLists.txt @@ -69,7 +69,10 @@ add_mlir_library( MLIRQCODialect MLIRQCOMatrix MLIRArithDialect + MLIRControlFlowDialect MLIRFuncDialect + MLIRMathDialect + MLIRMemRefDialect MLIRSCFDialect MQT::CoreDD PRIVATE diff --git a/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp b/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp index a9e52d540e..1696127276 100644 --- a/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp +++ b/mlir/lib/Dialect/QCO/Utils/DDFunctionality.cpp @@ -11,7 +11,10 @@ #include "mlir/Dialect/QCO/Utils/DDFunctionality.h" #include "dd/CachedEdge.hpp" +#include "dd/Complex.hpp" +#include "dd/ComplexValue.hpp" #include "dd/DDDefinitions.hpp" +#include "dd/Edge.hpp" #include "dd/GateMatrixDefinitions.hpp" #include "dd/Node.hpp" #include "dd/Operations.hpp" @@ -29,14 +32,22 @@ #include "mlir/Dialect/QCO/IR/QCOInterfaces.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Utils/Matrix.h" +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include +#include #include +#include #include #include #include +#include #include +#include #include +#include +#include #include #include #include @@ -49,21 +60,28 @@ #include #include +#include #include #include #include +#include #include #include #include #include +#include #include +#include #include #include +#include #include namespace mlir::qco { namespace { +constexpr size_t MAX_CONTROL_FLOW_STEPS = 10'000; + struct QubitMap { DenseMap qubits; size_t numQubits = 0; @@ -105,6 +123,55 @@ struct QubitMap { } return out; } + + void releaseWire(const qc::Qubit released) { + SmallVector aliases; + for (auto& [value, wire] : qubits) { + if (wire == released) { + aliases.push_back(value); + } else if (wire > released) { + --wire; + } + } + for (Value alias : aliases) { + qubits.erase(alias); + } + --numQubits; + } +}; + +/// Physical wires stored at each tensor index; extracted positions are empty. +using TensorSlots = SmallVector>; + +struct TensorMap { + DenseMap tensors; + + void bind(Value value, TensorSlots slots) { + tensors[value] = std::move(slots); + } + + [[nodiscard]] const TensorSlots* lookup(Value value) const { + const auto it = tensors.find(value); + return it == tensors.end() ? nullptr : &it->second; + } + + void erase(Value value) { tensors.erase(value); } + + void releaseWire(const qc::Qubit released) { + for (auto& [value, slots] : tensors) { + (void)value; + for (auto& wire : slots) { + if (!wire) { + continue; + } + if (*wire == released) { + wire.reset(); + } else if (*wire > released) { + --*wire; + } + } + } + } }; struct ClassicalEnv { @@ -113,11 +180,46 @@ struct ClassicalEnv { std::optional deferredWire; }; using RegisterState = std::vector; + using Scalar = std::variant; - DenseMap scalars; + DenseMap values; DenseMap deferredMeasurements; /// Shared storage preserves CBit register identity across `func.call`. DenseMap> registers; + /// Shared storage preserves caller-visible writes through `func.call`. + DenseMap>> memrefs; + + void releaseWire(const qc::Qubit released) { + for (auto& [value, wire] : deferredMeasurements) { + (void)value; + if (wire > released) { + --wire; + } + } + + DenseSet updated; + for (auto& [value, reg] : registers) { + (void)value; + if (!updated.insert(reg.get()).second) { + continue; + } + for (auto& cell : *reg) { + if (cell.deferredWire && *cell.deferredWire > released) { + --*cell.deferredWire; + } + } + } + } + + LogicalResult bindFrom(Value source, Value dest, Operation* op) { + const auto it = values.find(source); + if (it == values.end()) { + return op->emitError() + << "classical SSA value is not mapped for QCO DD simulation"; + } + values[dest] = it->second; + return success(); + } }; struct DecodedGate { @@ -125,13 +227,21 @@ struct DecodedGate { std::vector params; }; +enum class DeallocationMode : std::uint8_t { + Apply, + PreserveAll, + PreserveDeferred, +}; + struct WalkState { QubitMap* qubits; + TensorMap* tensors; ClassicalEnv* classical; dd::Package* dd; std::mt19937_64* rng = nullptr; const DenseSet* deferredMeasurements = nullptr; - size_t remainingExecutionSteps = 10'000; + DeallocationMode deallocationMode = DeallocationMode::Apply; + size_t remainingExecutionSteps = MAX_CONTROL_FLOW_STEPS; DenseSet activeCalls; }; struct LoopRange { @@ -143,12 +253,48 @@ struct SamplingPlan { SmallVector outputs; DenseSet deferredMeasurements; }; + +/// Distinguishes a density operator from a functionality matrix. +struct DensityState { + dd::MatrixDD matrix; +}; } // namespace +[[nodiscard]] static bool isQTensorType(Type type) { + const auto tensorType = dyn_cast(type); + return tensorType && tensorType.getRank() == 1 && + isa(tensorType.getElementType()); +} + +template +static FailureOr lookupScalar(Value value, const ClassicalEnv& classical, + Operation* op) { + const auto it = classical.values.find(value); + if (it == classical.values.end() || !std::holds_alternative(it->second)) { + return op->emitError() + << "classical SSA value is not mapped for QCO DD simulation"; + } + return std::get(it->second); +} + +static FailureOr +resolveDouble(Value value, const ClassicalEnv& classical, Operation* op) { + if (const auto it = classical.values.find(value); + it != classical.values.end() && + std::holds_alternative(it->second)) { + return std::get(it->second); + } + if (const auto constant = mqt::valueToDouble(value)) { + return *constant; + } + return op->emitError() + << "floating-point SSA value has no concrete QCO DD binding"; +} + /// `std::nullopt` if @p unitary is not a standard gate; failure if its unitary -/// matrix is not known at compile time. +/// parameters are not concrete. static FailureOr> -decodeStandardGate(UnitaryOpInterface unitary) { +decodeStandardGate(UnitaryOpInterface unitary, const ClassicalEnv& classical) { Operation* op = unitary.getOperation(); const auto type = TypeSwitch(op) @@ -185,15 +331,13 @@ decodeStandardGate(UnitaryOpInterface unitary) { if (type == qc::OpType::None) { return std::optional{std::nullopt}; } - if (!unitary.hasCompileTimeKnownUnitaryMatrix()) { - return unitary.emitError() - << "unitary must have a compile-time constant matrix"; - } - DecodedGate decoded{.type = type, .params = {}}; for (Value param : unitary.getParameters()) { - decoded.params.push_back( - static_cast(*mlir::mqt::valueToDouble(param))); + auto concrete = resolveDouble(param, classical, op); + if (failed(concrete)) { + return failure(); + } + decoded.params.push_back(static_cast(*concrete)); } return std::optional{std::move(decoded)}; } @@ -246,24 +390,92 @@ static dd::MatrixDD makeEmbeddedLocalDD(dd::Package& dd, return {.p = root.p, .w = dd.cn.lookup(root.w)}; } +static dd::mCachedEdge +buildDensityMatrix(const dd::VectorDD& ket, const dd::VectorDD& bra, + const int64_t level, dd::Package& dd, + std::map, + dd::mCachedEdge>& cache) { + if (ket.isZeroTerminal() || bra.isZeroTerminal()) { + return dd::mCachedEdge::zero(); + } + const auto weight = + static_cast(ket.w) * dd::ComplexNumbers::conj(bra.w); + if (level < 0) { + return dd::mCachedEdge::terminal(weight); + } + + const auto key = std::tuple{ket.p, bra.p, level}; + if (const auto cached = cache.find(key); cached != cache.end()) { + return {cached->second.p, cached->second.w * weight}; + } + const auto child = [level](const dd::VectorDD& edge, + const size_t index) -> dd::VectorDD { + if (!edge.isTerminal() && std::cmp_equal(edge.p->v, level)) { + return edge.p->e[index]; + } + return index == 0 ? dd::VectorDD{.p = edge.p, .w = dd::Complex::one()} + : dd::VectorDD::zero(); + }; + const auto ketZero = child(ket, 0); + const auto ketOne = child(ket, 1); + const auto braZero = child(bra, 0); + const auto braOne = child(bra, 1); + auto result = dd.makeDDNode( + static_cast(level), + {buildDensityMatrix(ketZero, braZero, level - 1, dd, cache), + buildDensityMatrix(ketZero, braOne, level - 1, dd, cache), + buildDensityMatrix(ketOne, braZero, level - 1, dd, cache), + buildDensityMatrix(ketOne, braOne, level - 1, dd, cache)}); + cache.try_emplace(key, result); + result.w = result.w * weight; + return result; +} + +static void applyStateOperation(const dd::MatrixDD& operation, dd::Package& dd, + dd::VectorDD& state) { + state = dd.applyOperation(operation, state); +} + +static void applyStateOperation(const dd::MatrixDD& operation, dd::Package& dd, + dd::MatrixDD& state) { + state = dd.applyOperation(operation, state); +} + +static void applyStateOperation(const dd::MatrixDD& operation, dd::Package& dd, + DensityState& state) { + const auto left = dd.multiply(operation, state.matrix); + const auto adjoint = dd.conjugateTranspose(operation); + auto result = dd.multiply(left, adjoint); + dd.incRef(result); + dd.decRef(state.matrix); + state.matrix = result; + dd.garbageCollect(); +} + template static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, WalkState& walk, StateDD& state) { Operation* op = unitary.getOperation(); - if (!unitary.hasCompileTimeKnownUnitaryMatrix()) { - return unitary.emitError() - << "unitary must have a compile-time constant matrix"; - } if (auto gphase = dyn_cast(op)) { - const auto theta = *mlir::mqt::valueToDouble(gphase.getTheta()); + if constexpr (std::is_same_v) { + return success(); + } + auto theta = resolveDouble(gphase.getTheta(), *walk.classical, op); + if (failed(theta)) { + return failure(); + } auto id = dd::Package::makeIdent(); - id.w = walk.dd->cn.lookup(std::cos(theta), std::sin(theta)); - state = walk.dd->applyOperation(id, state); + id.w = walk.dd->cn.lookup(std::cos(*theta), std::sin(*theta)); + applyStateOperation(id, *walk.dd, state); return success(); } if (isa(op)) { return walk.qubits->remapUnitary(unitary); } + if (!unitary.hasCompileTimeKnownUnitaryMatrix()) { + return unitary.emitError() + << "unitary must have a compile-time constant matrix"; + } DynamicMatrix local; if (!unitary.getUnitaryMatrixDynamic(local)) { @@ -285,7 +497,7 @@ 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 = walk.dd->applyOperation(walk.dd->makeGateDD(mat, wires[0]), state); + applyStateOperation(walk.dd->makeGateDD(mat, wires[0]), *walk.dd, state); return walk.qubits->remapUnitary(unitary); } @@ -297,8 +509,8 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, local(static_cast(row), static_cast(col)); } } - state = walk.dd->applyOperation( - walk.dd->makeTwoQubitGateDD(mat, wires[0], wires[1]), state); + applyStateOperation(walk.dd->makeTwoQubitGateDD(mat, wires[0], wires[1]), + *walk.dd, state); return walk.qubits->remapUnitary(unitary); } @@ -310,15 +522,15 @@ static LogicalResult applyUnitaryMatrix(UnitaryOpInterface unitary, local(static_cast(row), static_cast(col)); } } - state = walk.dd->applyOperation( + applyStateOperation( walk.dd->makeThreeQubitGateDD(mat, wires[0], wires[1], wires[2]), - state); + *walk.dd, state); return walk.qubits->remapUnitary(unitary); } - state = walk.dd->applyOperation( + applyStateOperation( makeEmbeddedLocalDD(*walk.dd, local, walk.qubits->numQubits, wires), - state); + *walk.dd, state); return walk.qubits->remapUnitary(unitary); } @@ -335,17 +547,33 @@ static LogicalResult applyDecodedStandard(UnitaryOpInterface unitary, if (failed(targets)) { return failure(); } - state = walk.dd->applyOperation( + applyStateOperation( getStandardOperationDD(*walk.dd, gate.type, gate.params, controls, {targets->begin(), targets->end()}), - state); + *walk.dd, state); return walk.qubits->remapUnitary(unitary); } static LogicalResult validateReturn(func::ReturnOp returnOp, - const QubitMap& qubits) { + const QubitMap& qubits, + const TensorMap& tensors) { qc::Qubit expected = 0; for (Value value : returnOp.getOperands()) { + if (isQTensorType(value.getType())) { + const auto* slots = tensors.lookup(value); + if (slots == nullptr) { + return returnOp.emitError() + << "returned qtensor is not mapped for QCO DD simulation"; + } + for (const auto wire : *slots) { + if (!wire || *wire != expected) { + return returnOp.emitError() + << "returned qubits must preserve canonical wire order"; + } + ++expected; + } + continue; + } if (!isa(value.getType())) { continue; } @@ -367,49 +595,152 @@ static LogicalResult validateReturn(func::ReturnOp returnOp, return success(); } -static void bindInteger(Value result, const llvm::APInt& value, - ClassicalEnv& classical) { - classical.scalars[result] = IntegerAttr::get(result.getType(), value); +static LogicalResult recordConstant(arith::ConstantOp constant, + ClassicalEnv& classical) { + if (auto attr = dyn_cast(constant.getValue())) { + classical.values[constant.getResult()] = attr.getValue(); + } else if (auto attr = dyn_cast(constant.getValue())) { + if (!constant.getType().isF64()) { + return constant.emitError() + << "QCO DD simulation only supports f64 classical values"; + } + classical.values[constant.getResult()] = attr.getValue().convertToDouble(); + } else if (auto attr = dyn_cast(constant.getValue())) { + if (constant.getType().isInteger(1)) { + classical.values[constant.getResult()] = attr.getValue() != 0; + } else if (isa(constant.getType())) { + classical.values[constant.getResult()] = attr.getInt(); + } else if (isa(constant.getType())) { + classical.values[constant.getResult()] = attr.getValue(); + } + } + return success(); } -static FailureOr -lookupInteger(Value value, ClassicalEnv& classical, Operation* op) { - const auto it = classical.scalars.find(value); - if (it == classical.scalars.end()) { - return op->emitError() << "classical SSA value is not mapped for QCO DD " - "simulation: " - << value.getType(); +static LogicalResult applyBindings(func::FuncOp func, + const DDBindings& bindings, + ClassicalEnv& classical) { + for (const auto& [value, attr] : bindings) { + auto argument = dyn_cast(value); + if (!argument || argument.getOwner() != &func.getBody().front()) { + return func.emitError() + << "QCO DD bindings must target entry-block arguments"; + } + const Type type = value.getType(); + if (isQTensorType(type)) { + if (cast(type).isDynamicDim(0) && + isa(attr)) { + continue; + } + } else if (type.isInteger(1)) { + if (auto boolean = dyn_cast(attr)) { + classical.values[value] = boolean.getValue(); + continue; + } + if (auto integer = dyn_cast(attr)) { + classical.values[value] = integer.getValue() != 0; + continue; + } + } else if (isa(type)) { + if (auto integer = dyn_cast(attr)) { + classical.values[value] = integer.getInt(); + continue; + } + } else if (auto integerType = dyn_cast(type)) { + if (auto integer = dyn_cast(attr)) { + classical.values[value] = + integer.getValue().sextOrTrunc(integerType.getWidth()); + continue; + } + } else if (type.isF64()) { + if (auto floating = dyn_cast(attr); + floating && floating.getType().isF64()) { + classical.values[value] = floating.getValue().convertToDouble(); + continue; + } + } + return func.emitError() << "QCO DD binding attribute " << attr + << " does not match argument type " << type; } - return it->second.getValue(); + return success(); } -static FailureOr lookupBool(Value value, ClassicalEnv& classical, +static FailureOr lookupBool(Value value, const ClassicalEnv& classical, Operation* op) { - auto result = lookupInteger(value, classical, op); - if (failed(result)) { - return failure(); - } - return !result->isZero(); + return lookupScalar(value, classical, op); } -static FailureOr lookupIndex(Value value, ClassicalEnv& classical, - Operation* op) { - auto result = lookupInteger(value, classical, op); - if (failed(result)) { - return failure(); +static FailureOr +lookupIndex(Value value, const ClassicalEnv& classical, Operation* op) { + return lookupScalar(value, classical, op); +} + +static FailureOr lookupFloat(Value value, const ClassicalEnv& classical, + Operation* op) { + return lookupScalar(value, classical, op); +} + +static FailureOr +lookupInteger(Value value, const ClassicalEnv& classical, Operation* op) { + if (value.getType().isInteger(1)) { + auto bit = lookupBool(value, classical, op); + if (failed(bit)) { + return failure(); + } + return llvm::APInt(1, static_cast(*bit)); + } + if (isa(value.getType())) { + auto index = lookupIndex(value, classical, op); + if (failed(index)) { + return failure(); + } + return llvm::APInt(64, static_cast(*index)); } - return result->getSExtValue(); + if (isa(value.getType())) { + return lookupScalar(value, classical, op); + } + return op->emitError() << "expected an integer or index SSA value"; } -static LogicalResult applyUnsignedIndexCast(Value in, Value out, Operation* op, - ClassicalEnv& classical) { - auto value = lookupInteger(in, classical, op); - if (failed(value)) { +[[nodiscard]] static bool evaluateCmp(arith::CmpIPredicate predicate, + const llvm::APInt& lhs, + const llvm::APInt& rhs) { + switch (predicate) { + case arith::CmpIPredicate::eq: + return lhs == rhs; + case arith::CmpIPredicate::ne: + return lhs != rhs; + case arith::CmpIPredicate::slt: + return lhs.slt(rhs); + case arith::CmpIPredicate::sle: + return lhs.sle(rhs); + case arith::CmpIPredicate::sgt: + return lhs.sgt(rhs); + case arith::CmpIPredicate::sge: + return lhs.sge(rhs); + case arith::CmpIPredicate::ult: + return lhs.ult(rhs); + case arith::CmpIPredicate::ule: + return lhs.ule(rhs); + case arith::CmpIPredicate::ugt: + return lhs.ugt(rhs); + case arith::CmpIPredicate::uge: + return lhs.uge(rhs); + } + llvm_unreachable("unknown arith.cmpi predicate"); +} + +static LogicalResult bindInteger(Value dest, const llvm::APInt& value, + ClassicalEnv& classical) { + if (dest.getType().isInteger(1)) { + classical.values[dest] = value[0]; + } else if (isa(dest.getType())) { + classical.values[dest] = static_cast(value.getZExtValue()); + } else if (auto type = dyn_cast(dest.getType())) { + classical.values[dest] = value.zextOrTrunc(type.getWidth()); + } else { return failure(); } - const auto integerType = dyn_cast(out.getType()); - const unsigned width = integerType ? integerType.getWidth() : 64U; - bindInteger(out, value->zextOrTrunc(width), classical); return success(); } @@ -428,7 +759,7 @@ static LogicalResult allocateRegister(cbit::AllocOp alloc, static FailureOr resolveRegisterIndex(Value index, cbit::RegisterType type, - ClassicalEnv& classical, + const ClassicalEnv& classical, Operation* op) { auto resolved = lookupIndex(index, classical, op); if (failed(resolved)) { @@ -485,85 +816,501 @@ static LogicalResult loadRegister(cbit::LoadOp load, ClassicalEnv& classical) { if (!cell.value) { return load.emitError() << "read from an undefined CBit register element"; } - bindInteger(load.getResult(), llvm::APInt(1, *cell.value ? 1 : 0), classical); + return bindInteger(load.getResult(), + llvm::APInt(1, static_cast(*cell.value)), + classical); +} + +[[nodiscard]] static bool isSupportedClassicalType(Type type) { + return isa(type) || type.isF64(); +} + +static FailureOr +lookupMemRefSlot(Value memref, ValueRange indices, ClassicalEnv& classical, + Operation* op) { + const auto type = dyn_cast(memref.getType()); + if (!type || type.getRank() != 1 || indices.size() != 1 || + !isSupportedClassicalType(type.getElementType())) { + return op->emitError() + << "QCO DD simulation only supports one-dimensional memrefs of " + "integer, index, or f64 values"; + } + auto index = lookupIndex(indices[0], classical, op); + if (failed(index)) { + return failure(); + } + const auto it = classical.memrefs.find(memref); + if (it == classical.memrefs.end()) { + return op->emitError() + << "classical memref is not mapped for QCO DD simulation"; + } + if (*index < 0 || static_cast(*index) >= it->second->size()) { + return op->emitError() + << "classical memref index out of range for QCO DD simulation"; + } + return &(*it->second)[static_cast(*index)]; +} + +static ClassicalEnv::Scalar zeroScalar(Type type) { + if (type.isInteger(1)) { + return false; + } + if (isa(type)) { + return int64_t{0}; + } + if (auto integer = dyn_cast(type)) { + return llvm::APInt(integer.getWidth(), 0); + } + return 0.0; +} + +static LogicalResult applyMemRefAlloc(memref::AllocOp alloc, + ClassicalEnv& classical) { + const auto type = dyn_cast(alloc.getType()); + if (!type || type.getRank() != 1 || + !isSupportedClassicalType(type.getElementType())) { + return alloc.emitError() + << "QCO DD simulation only supports one-dimensional memrefs of " + "integer, index, or f64 values"; + } + if (!alloc.getSymbolOperands().empty()) { + return alloc.emitError() + << "QCO DD simulation does not support symbolic memref operands"; + } + int64_t size = type.getDimSize(0); + if (type.isDynamicDim(0)) { + if (alloc.getDynamicSizes().size() != 1) { + return alloc.emitError() << "dynamic 1-D memref requires one size"; + } + auto dynamicSize = + lookupIndex(alloc.getDynamicSizes()[0], classical, alloc); + if (failed(dynamicSize)) { + return failure(); + } + size = *dynamicSize; + } + if (size < 0) { + return alloc.emitError() << "classical memref size must be non-negative"; + } + classical.memrefs[alloc.getResult()] = + std::make_shared>( + static_cast(size), zeroScalar(type.getElementType())); + return success(); +} + +static LogicalResult applyMemRefStore(memref::StoreOp store, + ClassicalEnv& classical) { + auto slot = + lookupMemRefSlot(store.getMemref(), store.getIndices(), classical, store); + const auto value = classical.values.find(store.getValue()); + if (failed(slot) || value == classical.values.end()) { + if (value == classical.values.end()) { + store.emitError() + << "stored classical value is not mapped for QCO DD simulation"; + } + return failure(); + } + **slot = value->second; + return success(); +} + +static LogicalResult applyMemRefLoad(memref::LoadOp load, + ClassicalEnv& classical) { + auto slot = + lookupMemRefSlot(load.getMemref(), load.getIndices(), classical, load); + if (failed(slot)) { + return failure(); + } + classical.values[load.getResult()] = **slot; return success(); } -static LogicalResult applyBinaryInteger(Operation& op, - ClassicalEnv& classical) { - auto lhs = lookupInteger(op.getOperand(0), classical, &op); - auto rhs = lookupInteger(op.getOperand(1), classical, &op); +template +static LogicalResult applyBinaryInteger(OpTy op, ClassicalEnv& classical, + Combine combine) { + auto lhs = lookupInteger(op.getLhs(), classical, op); + auto rhs = lookupInteger(op.getRhs(), classical, op); if (failed(lhs) || failed(rhs)) { return failure(); } - llvm::APInt result = *lhs; - if (isa(&op)) { - result &= *rhs; - } else if (isa(&op)) { - result |= *rhs; - } else if (isa(&op)) { - result ^= *rhs; - } else if (isa(&op)) { - result += *rhs; - } else if (isa(&op)) { - result -= *rhs; - } else if (isa(&op)) { - result *= *rhs; - } else { - if (rhs->isNegative() || rhs->uge(lhs->getBitWidth())) { - return op.emitError() - << "shift amount out of range for QCO DD simulation"; - } - const auto amount = static_cast(rhs->getZExtValue()); - result = isa(&op) ? lhs->shl(amount) : lhs->lshr(amount); + return bindInteger(op.getResult(), combine(*lhs, *rhs), classical); +} + +template +static LogicalResult applyBinaryFloat(OpTy op, ClassicalEnv& classical, + Combine combine) { + auto lhs = lookupFloat(op.getLhs(), classical, op); + auto rhs = lookupFloat(op.getRhs(), classical, op); + if (failed(lhs) || failed(rhs)) { + return failure(); + } + classical.values[op.getResult()] = combine(*lhs, *rhs); + return success(); +} + +template +static LogicalResult applyUnaryFloat(OpTy op, ClassicalEnv& classical, + Apply apply) { + auto value = lookupFloat(op.getOperand(), classical, op); + if (failed(value)) { + return failure(); } - bindInteger(op.getResult(0), result, classical); + classical.values[op.getResult()] = apply(*value); return success(); } +template +static LogicalResult applyDivision(OpTy op, ClassicalEnv& classical, + Combine combine) { + auto rhs = lookupInteger(op.getRhs(), classical, op); + if (failed(rhs)) { + return failure(); + } + if (rhs->isZero()) { + return op.emitError() << "division by zero during QCO DD simulation"; + } + auto lhs = lookupInteger(op.getLhs(), classical, op); + if (failed(lhs)) { + return failure(); + } + return bindInteger(op.getResult(), combine(*lhs, *rhs), classical); +} + +template +static LogicalResult applyShift(OpTy op, ClassicalEnv& classical, Shift shift) { + auto lhs = lookupInteger(op.getLhs(), classical, op); + auto rhs = lookupInteger(op.getRhs(), classical, op); + if (failed(lhs) || failed(rhs)) { + return failure(); + } + if (rhs->uge(lhs->getBitWidth())) { + return op.emitError() << "shift amount out of range for QCO DD simulation"; + } + return bindInteger(op.getResult(), shift(*lhs, rhs->getZExtValue()), + classical); +} + +static LogicalResult applyIntegerCast(Value in, Value out, Operation* op, + ClassicalEnv& classical, bool isSigned) { + auto value = lookupInteger(in, classical, op); + if (failed(value)) { + return failure(); + } + const unsigned width = isa(out.getType()) + ? 64U + : cast(out.getType()).getWidth(); + if (width > value->getBitWidth()) { + *value = isSigned ? value->sext(width) : value->zext(width); + } else if (width < value->getBitWidth()) { + *value = value->trunc(width); + } + return bindInteger(out, *value, classical); +} + static LogicalResult applyClassicalOp(Operation& op, ClassicalEnv& classical) { + const auto isUnsupportedFloat = [](Type type) { + return isa(type) && !type.isF64(); + }; + if (llvm::any_of(op.getOperandTypes(), isUnsupportedFloat) || + llvm::any_of(op.getResultTypes(), isUnsupportedFloat)) { + return op.emitError() + << "QCO DD simulation only supports f64 classical values"; + } return TypeSwitch(&op) - .Case( - [&](Operation* binary) { - return applyBinaryInteger(*binary, classical); - }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs & rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs | rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs ^ rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs + rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs - rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs * rhs; + }); + }) + .Case([&](auto value) { + return applyDivision( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.udiv(rhs); + }); + }) + .Case([&](auto value) { + return applyDivision( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.sdiv(rhs); + }); + }) + .Case([&](auto value) { + return applyDivision( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.urem(rhs); + }); + }) + .Case([&](auto value) { + return applyDivision( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.srem(rhs); + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.sgt(rhs) ? lhs : rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.slt(rhs) ? lhs : rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.ugt(rhs) ? lhs : rhs; + }); + }) + .Case([&](auto value) { + return applyBinaryInteger( + value, classical, + [](const llvm::APInt& lhs, const llvm::APInt& rhs) { + return lhs.ult(rhs) ? lhs : rhs; + }); + }) + .Case([&](auto value) { + return applyShift(value, classical, + [](const llvm::APInt& lhs, const uint64_t rhs) { + return lhs.shl(rhs); + }); + }) + .Case([&](auto value) { + return applyShift(value, classical, + [](const llvm::APInt& lhs, const uint64_t rhs) { + return lhs.lshr(rhs); + }); + }) + .Case([&](auto value) { + return applyShift(value, classical, + [](const llvm::APInt& lhs, const uint64_t rhs) { + return lhs.ashr(rhs); + }); + }) .Case([&](arith::CmpIOp cmp) -> LogicalResult { auto lhs = lookupInteger(cmp.getLhs(), classical, cmp); auto rhs = lookupInteger(cmp.getRhs(), classical, cmp); if (failed(lhs) || failed(rhs)) { return failure(); } - bindInteger(cmp.getResult(), - llvm::APInt(1, arith::applyCmpPredicate(cmp.getPredicate(), - *lhs, *rhs)), - classical); + classical.values[cmp.getResult()] = + evaluateCmp(cmp.getPredicate(), *lhs, *rhs); return success(); }) .Case([&](arith::SelectOp select) -> LogicalResult { - auto cond = lookupBool(select.getCondition(), classical, select); - if (failed(cond)) { + auto condition = lookupBool(select.getCondition(), classical, select); + if (failed(condition)) { return failure(); } - if (!isa(select.getType())) { - return select.emitError() - << "QCO DD simulation only supports integer or index select"; + Value selected = + *condition ? select.getTrueValue() : select.getFalseValue(); + return classical.bindFrom(selected, select.getResult(), select); + }) + .Case([&](arith::ExtUIOp ext) { + return applyIntegerCast(ext.getIn(), ext.getOut(), ext, classical, + false); + }) + .Case([&](arith::ExtSIOp cast) { + return applyIntegerCast(cast.getIn(), cast.getOut(), cast, classical, + true); + }) + .Case([&](arith::IndexCastUIOp cast) { + return applyIntegerCast(cast.getIn(), cast.getOut(), cast, classical, + false); + }) + .Case([&](arith::IndexCastOp cast) { + return applyIntegerCast(cast.getIn(), cast.getOut(), cast, classical, + true); + }) + .Case([&](arith::TruncIOp cast) { + return applyIntegerCast(cast.getIn(), cast.getOut(), cast, classical, + false); + }) + .Case([&](auto value) { + return applyBinaryFloat( + value, classical, [](double lhs, double rhs) { return lhs + rhs; }); + }) + .Case([&](auto value) { + return applyBinaryFloat( + value, classical, [](double lhs, double rhs) { return lhs - rhs; }); + }) + .Case([&](auto value) { + return applyBinaryFloat( + value, classical, [](double lhs, double rhs) { return lhs * rhs; }); + }) + .Case([&](auto value) { + return applyBinaryFloat( + value, classical, [](double lhs, double rhs) { return lhs / rhs; }); + }) + .Case([&](auto value) { + return applyBinaryFloat(value, classical, [](double lhs, double rhs) { + return std::fmod(lhs, rhs); + }); + }) + .Case([&](auto value) { + return applyBinaryFloat(value, classical, [](double lhs, double rhs) { + if (std::isnan(lhs) || std::isnan(rhs)) { + return std::numeric_limits::quiet_NaN(); + } + return std::fmax(lhs, rhs); + }); + }) + .Case([&](auto value) { + return applyBinaryFloat(value, classical, [](double lhs, double rhs) { + if (std::isnan(lhs) || std::isnan(rhs)) { + return std::numeric_limits::quiet_NaN(); + } + return std::fmin(lhs, rhs); + }); + }) + .Case([&](auto value) { + return applyBinaryFloat(value, classical, [](double lhs, double rhs) { + return std::fmax(lhs, rhs); + }); + }) + .Case([&](auto value) { + return applyBinaryFloat(value, classical, [](double lhs, double rhs) { + return std::fmin(lhs, rhs); + }); + }) + .Case([&](arith::NegFOp neg) -> LogicalResult { + auto value = lookupFloat(neg.getOperand(), classical, neg); + if (failed(value)) { + return failure(); } - auto t = lookupInteger(select.getTrueValue(), classical, select); - auto f = lookupInteger(select.getFalseValue(), classical, select); - if (failed(t) || failed(f)) { + classical.values[neg.getResult()] = -*value; + return success(); + }) + .Case([&](arith::CmpFOp cmp) -> LogicalResult { + auto lhs = lookupFloat(cmp.getLhs(), classical, cmp); + auto rhs = lookupFloat(cmp.getRhs(), classical, cmp); + if (failed(lhs) || failed(rhs)) { return failure(); } - bindInteger(select.getResult(), *cond ? *t : *f, classical); + classical.values[cmp.getResult()] = arith::applyCmpPredicate( + cmp.getPredicate(), llvm::APFloat(*lhs), llvm::APFloat(*rhs)); return success(); }) - .Case([&](arith::IndexCastUIOp cast) { - return applyUnsignedIndexCast(cast.getIn(), cast.getOut(), cast, - classical); + .Case( + [&](Operation* castOp) -> LogicalResult { + auto value = + lookupInteger(castOp->getOperand(0), classical, castOp); + if (failed(value)) { + return failure(); + } + classical.values[castOp->getResult(0)] = + value->roundToDouble(isa(castOp)); + return success(); + }) + .Case( + [&](Operation* castOp) -> LogicalResult { + auto value = lookupFloat(castOp->getOperand(0), classical, castOp); + if (failed(value)) { + return failure(); + } + Value out = castOp->getResult(0); + const unsigned width = cast(out.getType()).getWidth(); + const bool isSigned = isa(castOp); + llvm::APSInt result(width, /*isUnsigned=*/!isSigned); + bool exact = false; + const auto status = llvm::APFloat(*value).convertToInteger( + result, llvm::APFloat::rmTowardZero, &exact); + if ((status & llvm::APFloat::opInvalidOp) != 0) { + return castOp->emitError() + << "floating-point value is outside the destination " + "integer range during QCO DD simulation"; + } + return bindInteger(out, result, classical); + }) + .Case([&](auto value) { + return applyUnaryFloat( + value, classical, [](double operand) { return std::abs(operand); }); + }) + .Case([&](auto value) { + return applyUnaryFloat(value, classical, [](double operand) { + return std::ceil(operand); + }); + }) + .Case([&](auto value) { + return applyUnaryFloat( + value, classical, [](double operand) { return std::cos(operand); }); + }) + .Case([&](auto value) { + return applyUnaryFloat( + value, classical, [](double operand) { return std::exp(operand); }); + }) + .Case([&](auto value) { + return applyUnaryFloat(value, classical, [](double operand) { + return std::floor(operand); + }); + }) + .Case([&](auto value) { + return applyUnaryFloat( + value, classical, [](double operand) { return std::log(operand); }); }) - .Case([&](arith::ExtUIOp cast) { - return applyUnsignedIndexCast(cast.getIn(), cast.getOut(), cast, - classical); + .Case([&](auto value) { + return applyUnaryFloat( + value, classical, [](double operand) { return std::sin(operand); }); + }) + .Case([&](auto value) { + return applyUnaryFloat(value, classical, [](double operand) { + return std::sqrt(operand); + }); + }) + .Case([&](auto value) { + return applyUnaryFloat( + value, classical, [](double operand) { return std::tan(operand); }); + }) + .Case([&](auto value) { + return applyBinaryFloat(value, classical, [](double lhs, double rhs) { + return std::pow(lhs, rhs); + }); }) .Default([](Operation* unsupported) { return unsupported->emitError() @@ -609,17 +1356,39 @@ resolveLoop(scf::ForOp forOp, ClassicalEnv& classical, size_t remainingSteps) { } static LogicalResult bindValuePairs(ValueRange sources, ValueRange dests, + WalkState& walk, Operation* op); + +static LogicalResult bindLinearArgs(ValueRange operands, Block& block, WalkState& walk, Operation* op) { - const QubitMap sourceQubits = *walk.qubits; - const ClassicalEnv sourceClassical = *walk.classical; - for (auto [src, dest] : llvm::zip_equal(sources, dests)) { - if (isa(dest.getType())) { + for (Value arg : block.getArguments()) { + if (!isa(arg.getType()) && !isQTensorType(arg.getType())) { + return op->emitError() + << "unsupported linear region argument for QCO DD simulation"; + } + } + return bindValuePairs(operands, block.getArguments(), walk, op); +} + +static LogicalResult bindValuePairs(ValueRange sources, ValueRange dests, + WalkState& walk, Operation* op) { + const QubitMap sourceQubits = *walk.qubits; + const TensorMap sourceTensors = *walk.tensors; + const ClassicalEnv sourceClassical = *walk.classical; + for (auto [src, dest] : llvm::zip_equal(sources, dests)) { + if (isa(dest.getType())) { const auto q = sourceQubits.lookup(src); if (!q) { return op->emitError() << "qubit SSA value is not mapped for QCO DD construction"; } walk.qubits->bind(dest, *q); + } else if (isQTensorType(dest.getType())) { + const auto* slots = sourceTensors.lookup(src); + if (slots == nullptr) { + return op->emitError() + << "qtensor SSA value is not mapped for QCO DD simulation"; + } + walk.tensors->bind(dest, *slots); } else if (isa(dest.getType())) { const auto it = sourceClassical.registers.find(src); if (it == sourceClassical.registers.end()) { @@ -627,21 +1396,49 @@ static LogicalResult bindValuePairs(ValueRange sources, ValueRange dests, << "CBit register is not mapped for QCO DD simulation"; } walk.classical->registers[dest] = it->second; + } else if (isa(dest.getType())) { + const auto it = sourceClassical.memrefs.find(src); + if (it == sourceClassical.memrefs.end()) { + return op->emitError() + << "classical memref is not mapped for QCO DD simulation"; + } + walk.classical->memrefs[dest] = it->second; } else { - const auto value = sourceClassical.scalars.find(src); - if (value == sourceClassical.scalars.end()) { + const auto value = sourceClassical.values.find(src); + if (value == sourceClassical.values.end()) { return op->emitError() << "classical SSA value is not mapped for QCO DD simulation"; } - walk.classical->scalars[dest] = value->second; + walk.classical->values[dest] = value->second; } } return success(); } +static LogicalResult bindYieldResults(YieldOp yield, + ValueRange classicalResults, + ValueRange linearResults, + WalkState& walk) { + const size_t numClassical = classicalResults.size(); + if (yield.getNumOperands() != numClassical + linearResults.size()) { + return yield.emitError() + << "yield operand count does not match operation results"; + } + if (failed(bindValuePairs(yield.getOperands().take_front(numClassical), + classicalResults, walk, yield))) { + return failure(); + } + return bindValuePairs(yield.getOperands().drop_front(numClassical), + linearResults, walk, yield); +} + template static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state); +template +static FailureOr +walkFunctionBody(func::FuncOp func, WalkState& walk, StateDD& state); + template static LogicalResult walkBlock(Block& block, WalkState& walk, StateDD& state) { for (Operation& op : block.without_terminator()) { @@ -653,29 +1450,483 @@ static LogicalResult walkBlock(Block& block, WalkState& walk, StateDD& state) { } template -static LogicalResult applyRegionBranch(ValueRange linearOperands, Block& block, - WalkState& walk, StateDD& state, - Operation* parent) { - if (failed( - bindValuePairs(linearOperands, block.getArguments(), walk, parent))) { +static FailureOr +walkConcreteCFG(Block& entry, WalkState& walk, StateDD& state, + function_ref isExit, StringRef scope) { + Block* block = &entry; + while (true) { + if (failed(walkBlock(*block, walk, state))) { + return failure(); + } + Operation* terminator = block->getTerminator(); + if (isExit(terminator)) { + return terminator; + } + + Block* successor = nullptr; + ValueRange successorOperands; + if (auto branch = dyn_cast(terminator)) { + successor = branch.getDest(); + successorOperands = branch.getDestOperands(); + } else if (auto branch = dyn_cast(terminator)) { + auto condition = + lookupBool(branch.getCondition(), *walk.classical, branch); + if (failed(condition)) { + return failure(); + } + successor = *condition ? branch.getTrueDest() : branch.getFalseDest(); + successorOperands = *condition ? branch.getTrueDestOperands() + : branch.getFalseDestOperands(); + } else if (auto switchOp = dyn_cast(terminator)) { + auto flag = lookupInteger(switchOp.getFlag(), *walk.classical, switchOp); + if (failed(flag)) { + return failure(); + } + successor = switchOp.getDefaultDestination(); + successorOperands = switchOp.getDefaultOperands(); + if (const auto caseValues = switchOp.getCaseValues()) { + for (auto [caseValue, destination, operands] : llvm::zip_equal( + caseValues->getValues(), + switchOp.getCaseDestinations(), switchOp.getCaseOperands())) { + if (*flag == caseValue) { + successor = destination; + successorOperands = operands; + break; + } + } + } + } else { + return terminator->emitError() << "unsupported " << scope + << " CFG terminator for QCO DD simulation"; + } + if (successorOperands.size() != successor->getNumArguments()) { + return terminator->emitError() + << scope + << " CFG successor operand count does not match block arguments"; + } + if (failed(bindValuePairs(successorOperands, successor->getArguments(), + walk, terminator))) { + return failure(); + } + if (walk.remainingExecutionSteps == 0) { + return terminator->emitError() + << "QCO DD execution exceeds the limit of 10000 control-flow " + "steps"; + } + --walk.remainingExecutionSteps; + block = successor; + } +} + +template +static LogicalResult +applyRegionBranch(ValueRange linearOperands, Block& block, + 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(); } - auto yield = cast(block.getTerminator()); - return bindValuePairs(yield.getOperands(), parent->getResults(), walk, yield); + return bindYieldResults(cast(block.getTerminator()), + classicalResults, linearResults, walk); +} + +template +static LogicalResult applyScfRegion(Region& region, ValueRange results, + WalkState& walk, StateDD& state, + Operation* parent) { + if (region.empty()) { + return parent->emitError() << "SCF region is empty"; + } + auto terminator = walkConcreteCFG( + region.front(), walk, state, + [](Operation* op) { return isa(op); }, "SCF region"); + if (failed(terminator)) { + return failure(); + } + auto yield = dyn_cast(*terminator); + if (!yield || yield.getNumOperands() != results.size()) { + return parent->emitError() + << "SCF region must yield one value for each result"; + } + return bindValuePairs(yield.getOperands(), results, walk, parent); +} + +template +static FailureOr allocateZeroQubits(size_t count, WalkState& walk, + StateDD& state, + Operation* op) { + static_assert(std::is_same_v || + std::is_same_v); + if (count == 0) { + return op->emitError() << "quantum allocation size must be positive"; + } + if (walk.qubits->numQubits > walk.dd->qubits() || + count > walk.dd->qubits() - walk.qubits->numQubits) { + return op->emitError() << "DD package has " << walk.dd->qubits() + << " qubits but allocation requires " + << walk.qubits->numQubits + count; + } + + const size_t first = walk.qubits->numQubits; + if constexpr (std::is_same_v) { + auto zeros = dd::makeZeroState(count, *walk.dd, first); + auto extended = walk.dd->kronecker(zeros, state, first, /*incIdx=*/false); + walk.dd->incRef(extended); + walk.dd->decRef(zeros); + walk.dd->decRef(state); + state = extended; + } else { + auto extended = state.matrix; + for (size_t i = 0; i < count; ++i) { + extended = walk.dd->makeDDNode( + static_cast(first + i), + {extended, dd::MatrixDD::zero(), dd::MatrixDD::zero(), + dd::MatrixDD::zero()}); + } + walk.dd->incRef(extended); + walk.dd->decRef(state.matrix); + state.matrix = extended; + walk.dd->garbageCollect(); + } + + TensorSlots slots; + slots.reserve(count); + for (size_t i = 0; i < count; ++i) { + slots.emplace_back(static_cast(first + i)); + } + walk.qubits->numQubits += count; + return slots; +} + +/// Project @p wire onto one basis state and remove its DD level. +static dd::VectorDD projectAndRemoveWire(const dd::VectorDD& root, + const qc::Qubit wire, + const bool projectOne, + dd::Package& dd) { + DenseMap projectedNodes; + const auto project = [&](const auto& self, + const dd::VectorDD& edge) -> dd::VectorDD { + if (edge.isZeroTerminal()) { + return edge; + } + // A skipped vector-DD level represents a qubit fixed to zero. + if (edge.isTerminal() || edge.p->v < wire) { + return projectOne ? dd::VectorDD::zero() : edge; + } + + dd::VectorDD projected; + if (const auto cached = projectedNodes.find(edge.p); + cached != projectedNodes.end()) { + projected = cached->second; + } else if (edge.p->v == wire) { + projected = edge.p->e[projectOne ? 1U : 0U]; + } else { + std::array edges{self(self, edge.p->e[0]), + self(self, edge.p->e[1])}; + projected = dd.makeDDNode( + static_cast(edge.p->v - 1U), edges); + } + projectedNodes.try_emplace(edge.p, projected); + projected.w = dd.cn.lookup(projected.w * edge.w); + return projected; + }; + return project(project, root); +} + +static LogicalResult deallocateWire(const qc::Qubit wire, WalkState& walk, + dd::VectorDD& state, Operation* op) { + if (wire >= walk.qubits->numQubits) { + return op->emitError() + << "deallocated wire is outside the simulated register"; + } + const auto zero = projectAndRemoveWire(state, wire, false, *walk.dd); + const auto one = projectAndRemoveWire(state, wire, true, *walk.dd); + if (zero.isZeroTerminal() && one.isZeroTerminal()) { + return op->emitError() << "cannot deallocate a zero-norm quantum state"; + } + if (!zero.isZeroTerminal() && !one.isZeroTerminal() && zero.p != one.p) { + return op->emitError() + << "deallocating an entangled qubit is not supported by " + "statevector QCO DD simulation"; + } + + const auto zeroWeight = static_cast(zero.w); + const auto oneWeight = static_cast(one.w); + const auto norm = std::sqrt(zeroWeight.mag2() + oneWeight.mag2()); + auto reduced = dd::VectorDD{.p = zero.isZeroTerminal() ? one.p : zero.p, + .w = walk.dd->cn.lookup(norm)}; + walk.dd->incRef(reduced); + walk.dd->decRef(state); + state = reduced; + walk.qubits->releaseWire(wire); + walk.tensors->releaseWire(wire); + walk.classical->releaseWire(wire); + return success(); +} + +static LogicalResult deallocateWire(const qc::Qubit wire, WalkState& walk, + DensityState& state, Operation* op) { + if (wire >= walk.qubits->numQubits) { + return op->emitError() + << "deallocated wire is outside the simulated register"; + } + std::vector eliminate(walk.qubits->numQubits, false); + eliminate[wire] = true; + auto reduced = walk.dd->partialTrace(state.matrix, eliminate); + // `partialTrace` normalizes by two for each removed matrix level. A physical + // partial trace does not, so restore the removed factor. + reduced.w = + walk.dd->cn.lookup(static_cast(reduced.w) * 2.0); + walk.dd->incRef(reduced); + walk.dd->decRef(state.matrix); + state.matrix = reduced; + walk.dd->garbageCollect(); + walk.qubits->releaseWire(wire); + walk.tensors->releaseWire(wire); + walk.classical->releaseWire(wire); + return success(); +} + +static double densityTrace(const dd::MatrixDD& density, const size_t numQubits, + dd::Package& dd) { + const auto normalized = dd.trace(density, numQubits); + return std::ldexp(normalized.r, static_cast(numQubits)); +} + +static FailureOr measureDensity(DensityState& state, const qc::Qubit wire, + const size_t numQubits, dd::Package& dd, + std::mt19937_64& rng, + Operation* diagnosticOp) { + const auto project = [&](const dd::GateMatrix& projector) { + const auto gate = dd.makeGateDD(projector, wire); + return dd.multiply(dd.multiply(gate, state.matrix), gate); + }; + auto zero = project(dd::MEAS_ZERO_MAT); + auto one = project(dd::MEAS_ONE_MAT); + const double rawZero = densityTrace(zero, numQubits, dd); + const double rawOne = densityTrace(one, numQubits, dd); + constexpr double tolerance = 1e-10; + if (!std::isfinite(rawZero) || !std::isfinite(rawOne) || + rawZero < -tolerance || rawOne < -tolerance) { + return diagnosticOp->emitError() + << "density matrix has invalid measurement probabilities"; + } + const double pzero = std::max(0.0, rawZero); + const double pone = std::max(0.0, rawOne); + const double sum = pzero + pone; + if (!std::isfinite(sum) || std::abs(sum - 1.0) > tolerance) { + return diagnosticOp->emitError() + << "density matrix must have unit trace for measurement"; + } + + std::uniform_real_distribution distribution(0.0, sum); + const bool measuredOne = distribution(rng) >= pzero; + auto collapsed = measuredOne ? one : zero; + const double probability = measuredOne ? pone : pzero; + if (!(probability > 0.0)) { + return diagnosticOp->emitError() + << "density measurement selected a zero-probability outcome"; + } + collapsed.w = + dd.cn.lookup(static_cast(collapsed.w) / probability); + dd.incRef(collapsed); + dd.decRef(state.matrix); + state.matrix = collapsed; + dd.garbageCollect(); + return measuredOne ? '1' : '0'; +} + +static FailureOr measureState(dd::VectorDD& state, const qc::Qubit wire, + const size_t /*numQubits*/, dd::Package& dd, + std::mt19937_64& rng, + Operation* /*diagnosticOp*/) { + return dd.measureOneCollapsing(state, wire, rng); +} + +static FailureOr measureState(DensityState& state, const qc::Qubit wire, + const size_t numQubits, dd::Package& dd, + std::mt19937_64& rng, + Operation* diagnosticOp) { + return measureDensity(state, wire, numQubits, dd, rng, diagnosticOp); +} + +static FailureOr +measureAllDensity(DensityState& state, const size_t numQubits, dd::Package& dd, + std::mt19937_64& rng, Operation* diagnosticOp) { + std::string result(numQubits, '0'); + for (size_t i = numQubits; i > 0; --i) { + const auto wire = static_cast(i - 1); + auto measured = + measureDensity(state, wire, numQubits, dd, rng, diagnosticOp); + if (failed(measured)) { + return failure(); + } + result[numQubits - i] = *measured; + } + return result; } 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) { - if (auto attr = dyn_cast(constant.getValue())) { - walk.classical->scalars[constant.getResult()] = attr; + return recordConstant(constant, *walk.classical); + }) + .template Case([&](AllocOp alloc) -> LogicalResult { + if constexpr (std::is_same_v) { + if (!walk.qubits->lookup(alloc.getResult())) { + return alloc.emitError() + << "dynamic qubit allocation is not supported for QCO DD " + "functionality construction"; + } + return success(); + } else { + auto slots = allocateZeroQubits(1, walk, state, alloc); + if (failed(slots)) { + return failure(); + } + walk.qubits->bind(alloc.getResult(), *slots->front()); + return success(); } - return success(); + }) + .template Case( + [&](qtensor::AllocOp alloc) -> LogicalResult { + if constexpr (std::is_same_v) { + return alloc.emitError() + << "qtensor allocation is not supported for QCO DD " + "functionality construction"; + } else { + auto size = lookupIndex(alloc.getSize(), *walk.classical, alloc); + if (failed(size)) { + return failure(); + } + if (*size <= 0) { + return alloc.emitError() + << "qtensor allocation size must be positive"; + } + auto slots = allocateZeroQubits(static_cast(*size), walk, + state, alloc); + if (failed(slots)) { + return failure(); + } + walk.tensors->bind(alloc.getResult(), std::move(*slots)); + return success(); + } + }) + .template Case( + [&](qtensor::FromElementsOp fromElements) -> LogicalResult { + auto wires = walk.qubits->lookupRange(fromElements.getElements(), + fromElements); + if (failed(wires)) { + return failure(); + } + TensorSlots slots; + slots.reserve(wires->size()); + for (const qc::Qubit wire : *wires) { + slots.emplace_back(wire); + } + walk.tensors->bind(fromElements.getResult(), std::move(slots)); + return success(); + }) + .template Case( + [&](qtensor::ExtractOp extract) -> LogicalResult { + const auto* input = walk.tensors->lookup(extract.getTensor()); + auto index = + lookupIndex(extract.getIndex(), *walk.classical, extract); + if (input == nullptr || failed(index)) { + if (input == nullptr) { + extract.emitError() + << "qtensor is not mapped for QCO DD simulation"; + } + return failure(); + } + if (*index < 0 || static_cast(*index) >= input->size()) { + return extract.emitError() << "qtensor index out of range"; + } + TensorSlots output = *input; + auto& wire = output[static_cast(*index)]; + if (!wire) { + return extract.emitError() + << "qtensor element has already been extracted"; + } + walk.qubits->bind(extract.getResult(), *wire); + wire.reset(); + walk.tensors->bind(extract.getOutTensor(), std::move(output)); + return success(); + }) + .template Case( + [&](qtensor::InsertOp insert) -> LogicalResult { + const auto* input = walk.tensors->lookup(insert.getDest()); + const auto wire = walk.qubits->lookup(insert.getScalar()); + auto index = + lookupIndex(insert.getIndex(), *walk.classical, insert); + if (input == nullptr || !wire || failed(index)) { + if (input == nullptr || !wire) { + insert.emitError() + << "qtensor or qubit is not mapped for QCO DD simulation"; + } + return failure(); + } + if (*index < 0 || static_cast(*index) >= input->size()) { + return insert.emitError() << "qtensor index out of range"; + } + TensorSlots output = *input; + output[static_cast(*index)] = wire; + walk.tensors->bind(insert.getResult(), std::move(output)); + return success(); + }) + .template Case( + [&](qtensor::DeallocOp dealloc) -> LogicalResult { + const auto* tracked = walk.tensors->lookup(dealloc.getTensor()); + if (tracked == nullptr) { + return dealloc.emitError() + << "qtensor is not mapped for QCO DD simulation"; + } + TensorSlots slots = *tracked; + walk.tensors->erase(dealloc.getTensor()); + const bool containsDeferredMeasurement = + llvm::any_of(slots, [&](const std::optional wire) { + return wire && + llvm::any_of(walk.classical->deferredMeasurements, + [&](const auto& measurement) { + return measurement.second == *wire; + }); + }); + if (walk.deallocationMode == DeallocationMode::PreserveAll || + (walk.deallocationMode == DeallocationMode::PreserveDeferred && + containsDeferredMeasurement)) { + return success(); + } + if constexpr (!std::is_same_v) { + SmallVector wires; + for (const auto wire : slots) { + if (wire) { + wires.push_back(*wire); + } + } + llvm::sort(wires, [](qc::Qubit lhs, qc::Qubit rhs) { + return lhs > rhs; + }); + for (const qc::Qubit wire : wires) { + if (failed(deallocateWire(wire, walk, state, dealloc))) { + return failure(); + } + } + } + return success(); + }) + .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([&](cbit::AllocOp alloc) { return allocateRegister(alloc, *walk.classical); @@ -686,18 +1937,28 @@ static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { .template Case([&](cbit::StoreOp store) { return storeRegister(store, *walk.classical); }) - .template Case( + .template Case([](auto) { return success(); }) + .template Case< + arith::AndIOp, arith::OrIOp, arith::XOrIOp, arith::AddIOp, + arith::SubIOp, arith::MulIOp, arith::DivUIOp, arith::DivSIOp, + arith::RemUIOp, arith::RemSIOp, arith::MaxSIOp, arith::MinSIOp, + arith::MaxUIOp, arith::MinUIOp, arith::ShLIOp, arith::ShRUIOp, + arith::ShRSIOp, arith::CmpIOp, arith::SelectOp, arith::ExtUIOp, + arith::ExtSIOp, arith::IndexCastUIOp, arith::IndexCastOp, + arith::TruncIOp, arith::AddFOp, arith::SubFOp, arith::MulFOp, + arith::DivFOp, arith::RemFOp, arith::MaximumFOp, arith::MinimumFOp, + arith::MaxNumFOp, arith::MinNumFOp, arith::NegFOp, arith::CmpFOp, + arith::SIToFPOp, arith::UIToFPOp, arith::FPToSIOp, arith::FPToUIOp, + math::AbsFOp, math::CeilOp, math::CosOp, math::ExpOp, math::FloorOp, + math::LogOp, math::SinOp, math::SqrtOp, math::TanOp, math::PowFOp>( [&](Operation* classicalOp) { return applyClassicalOp(*classicalOp, *walk.classical); }) .template Case([&](func::ReturnOp returnOp) { - return validateReturn(returnOp, *walk.qubits); + return validateReturn(returnOp, *walk.qubits, *walk.tensors); }) .template Case([&](MeasureOp measureOp) -> LogicalResult { - if constexpr (!std::is_same_v) { + if constexpr (std::is_same_v) { return measureOp.emitError() << "measurements are not supported for QCO DD functionality " "construction"; @@ -719,15 +1980,19 @@ static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { walk.qubits->bind(measureOp.getQubitOut(), *q); return success(); } - const char bit = walk.dd->measureOneCollapsing(state, *q, *walk.rng); - bindInteger(measureOp.getResult(), llvm::APInt(1, bit == '1'), - *walk.classical); + auto measured = measureState(state, *q, walk.qubits->numQubits, + *walk.dd, *walk.rng, measureOp); + if (failed(measured)) { + return failure(); + } + const char bit = *measured; + walk.classical->values[measureOp.getResult()] = bit == '1'; walk.qubits->bind(measureOp.getQubitOut(), *q); return success(); } }) .template Case([&](ResetOp resetOp) -> LogicalResult { - if constexpr (!std::is_same_v) { + if constexpr (std::is_same_v) { return resetOp.emitError() << "resets are not supported for QCO DD functionality " "construction"; @@ -740,107 +2005,192 @@ static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { return resetOp.emitError() << "qubit SSA value is not mapped for QCO DD construction"; } - const char bit = walk.dd->measureOneCollapsing(state, *q, *walk.rng); + auto measured = measureState(state, *q, walk.qubits->numQubits, + *walk.dd, *walk.rng, resetOp); + if (failed(measured)) { + return failure(); + } + const char bit = *measured; if (bit == '1') { - state = walk.dd->applyOperation( + applyStateOperation( walk.dd->makeGateDD( dd::opToSingleQubitGateMatrix(qc::OpType::X), *q), - state); + *walk.dd, 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 { - auto condition = - lookupBool(ifOp.getCondition(), *walk.classical, ifOp); - if (failed(condition)) { - return failure(); - } - Block* block = *condition ? ifOp.thenBlock() : ifOp.elseBlock(); - return applyRegionBranch(ifOp.getQubits(), *block, walk, state, ifOp); + auto condition = lookupBool(ifOp.getCondition(), *walk.classical, ifOp); + if (failed(condition)) { + return failure(); } + Block* block = *condition ? ifOp.thenBlock() : ifOp.elseBlock(); + if (block == nullptr) { + return ifOp.emitError() << "selected qco.if region is empty"; + } + return applyRegionBranch(ifOp.getQubits(), *block, + 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 { - auto index = - lookupIndex(switchOp.getArg(), *walk.classical, switchOp); - if (failed(index)) { - return failure(); + auto selector = + lookupIndex(switchOp.getArg(), *walk.classical, switchOp); + if (failed(selector)) { + return failure(); + } + Block* block = switchOp.getDefaultBlock(); + for (auto [i, caseValue] : llvm::enumerate(switchOp.getCases())) { + if (caseValue == *selector) { + block = switchOp.getCaseBlock(i); + break; } - const auto cases = switchOp.getCases(); - Block* block = switchOp.getDefaultBlock(); - for (auto [i, caseValue] : llvm::enumerate(cases)) { - if (caseValue == *index) { - block = switchOp.getCaseBlock(i); - break; - } + } + if (block == nullptr) { + return switchOp.emitError() + << "selected qco.index_switch region is empty"; + } + return applyRegionBranch( + switchOp.getTargets(), *block, switchOp.getClassicalResults(), + switchOp.getLinearResults(), walk, state, switchOp); + }) + .template Case([&](scf::IfOp ifOp) -> LogicalResult { + auto condition = lookupBool(ifOp.getCondition(), *walk.classical, ifOp); + if (failed(condition)) { + return failure(); + } + Region& selected = + *condition ? ifOp.getThenRegion() : ifOp.getElseRegion(); + if (selected.empty()) { + return ifOp.getNumResults() == 0 + ? success() + : ifOp.emitError() + << "selected empty scf.if region has results"; + } + return applyScfRegion(selected, ifOp.getResults(), walk, state, ifOp); + }) + .template Case( + [&](scf::IndexSwitchOp switchOp) -> LogicalResult { + auto selector = + lookupIndex(switchOp.getArg(), *walk.classical, switchOp); + if (failed(selector)) { + return failure(); + } + Region* selected = &switchOp.getDefaultRegion(); + for (auto [i, value] : llvm::enumerate(switchOp.getCases())) { + if (value == *selector) { + selected = &switchOp.getCaseRegions()[i]; + break; } - return applyRegionBranch(switchOp.getTargets(), *block, walk, - state, switchOp); } + return applyScfRegion(*selected, switchOp.getResults(), walk, state, + switchOp); + }) + .template Case( + [&](scf::ExecuteRegionOp execute) -> LogicalResult { + return applyScfRegion(execute.getRegion(), execute.getResults(), + walk, state, execute); }) .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 range = - resolveLoop(forOp, *walk.classical, walk.remainingExecutionSteps); - if (failed(range)) { - return failure(); - } + auto range = + resolveLoop(forOp, *walk.classical, walk.remainingExecutionSteps); + if (failed(range)) { + return failure(); + } - Block& body = *forOp.getBody(); - SmallVector carried(forOp.getInits().begin(), - forOp.getInits().end()); + Block& body = *forOp.getBody(); + SmallVector carried(forOp.getInits().begin(), + forOp.getInits().end()); - for (size_t t = 0; t < range->trips; - ++t, range->induction += range->step) { - if (walk.remainingExecutionSteps == 0) { - return forOp.emitError( - "QCO DD execution exceeds the limit of 10000 control-flow " - "steps"); - } - --walk.remainingExecutionSteps; - auto iterArgs = body.getArguments().drop_front(); - if (failed(bindValuePairs(carried, iterArgs, walk, forOp))) { - return failure(); - } - bindInteger( - body.getArgument(0), - range->induction.trunc(range->induction.getBitWidth() - 1), - *walk.classical); - if (failed(walkBlock(body, walk, state))) { - return failure(); - } - auto yield = cast(body.getTerminator()); - carried.assign(yield.getOperands().begin(), - yield.getOperands().end()); + for (size_t t = 0; t < range->trips; + ++t, range->induction += range->step) { + if (walk.remainingExecutionSteps == 0) { + return forOp.emitError( + "QCO DD execution exceeds the limit of 10000 control-flow " + "steps"); + } + --walk.remainingExecutionSteps; + auto iterArgs = body.getArguments().drop_front(); + if (failed(bindValuePairs(carried, iterArgs, walk, forOp))) { + return failure(); } - return bindValuePairs(carried, forOp.getResults(), walk, forOp); + if (failed(bindInteger( + body.getArgument(0), + range->induction.trunc(range->induction.getBitWidth() - 1), + *walk.classical))) { + return failure(); + } + if (failed(walkBlock(body, walk, state))) { + return failure(); + } + auto yield = cast(body.getTerminator()); + carried.assign(yield.getOperands().begin(), + yield.getOperands().end()); + } + return bindValuePairs(carried, forOp.getResults(), walk, forOp); + }) + .template Case([&](scf::WhileOp whileOp) -> LogicalResult { + if (!whileOp.getBefore().hasOneBlock() || + !whileOp.getAfter().hasOneBlock()) { + return whileOp.emitError() + << "scf.while regions must contain one block"; + } + Block& before = whileOp.getBefore().front(); + Block& after = whileOp.getAfter().front(); + SmallVector carried(whileOp.getInits().begin(), + whileOp.getInits().end()); + while (true) { + if (failed(bindValuePairs(carried, before.getArguments(), walk, + whileOp)) || + failed(walkBlock(before, walk, state))) { + return failure(); + } + auto condition = dyn_cast(before.getTerminator()); + if (!condition) { + return whileOp.emitError() + << "scf.while before region missing scf.condition"; + } + auto value = + lookupBool(condition.getCondition(), *walk.classical, whileOp); + if (failed(value)) { + return failure(); + } + if (!*value) { + return bindValuePairs(condition.getArgs(), whileOp.getResults(), + walk, whileOp); + } + if (walk.remainingExecutionSteps == 0) { + return whileOp.emitError( + "QCO DD execution exceeds the limit of 10000 control-flow " + "steps"); + } + --walk.remainingExecutionSteps; + if (failed(bindValuePairs(condition.getArgs(), after.getArguments(), + walk, whileOp)) || + failed(walkBlock(after, walk, state))) { + return failure(); + } + auto yield = dyn_cast(after.getTerminator()); + if (!yield) { + return whileOp.emitError() + << "scf.while after region missing scf.yield"; + } + carried.assign(yield.getOperands().begin(), + yield.getOperands().end()); } }) .template Case([&](func::CallOp call) -> LogicalResult { auto callee = SymbolTable::lookupNearestSymbolFrom( call, call.getCalleeAttr()); - if (!callee.getBody().hasOneBlock()) { - return call.emitError() - << "func.call callee must have a single-block body"; + if (!callee) { + return call.emitError() << "func.call callee '" << call.getCallee() + << "' could not be resolved"; + } + if (callee.isDeclaration()) { + return call.emitError() << "func.call callee must have a body"; } - auto returnOp = - cast(callee.getBody().front().getTerminator()); Operation* calleeOp = callee.getOperation(); if (!walk.activeCalls.insert(calleeOp).second) { return call.emitError() @@ -855,16 +2205,19 @@ static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { return failure(); } - if (failed(walkBlock(callee.getBody().front(), walk, state))) { + auto returnOp = walkFunctionBody(callee, walk, state); + if (failed(returnOp)) { return failure(); } - return bindValuePairs(returnOp.getOperands(), call.getResults(), walk, + + // Map callee return operands onto call results via the return op. + return bindValuePairs(returnOp->getOperands(), call.getResults(), walk, call); }) .template Case([&](CtrlOp ctrlOp) -> LogicalResult { if (auto inner = mqt::getSoleBodyUnitary( *ctrlOp.getBody())) { - auto decoded = decodeStandardGate(inner); + auto decoded = decodeStandardGate(inner, *walk.classical); if (failed(decoded)) { return failure(); } @@ -886,7 +2239,7 @@ static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { }) .template Case( [&](UnitaryOpInterface unitary) -> LogicalResult { - auto decoded = decodeStandardGate(unitary); + auto decoded = decodeStandardGate(unitary, *walk.classical); if (failed(decoded)) { return failure(); } @@ -907,59 +2260,259 @@ static LogicalResult applyOp(Operation& op, WalkState& walk, StateDD& state) { } template -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, walkState, state))) { - return failure(); - } +static FailureOr +walkFunctionBody(func::FuncOp func, WalkState& walk, StateDD& state) { + auto terminator = walkConcreteCFG( + func.getBody().front(), walk, state, + [](Operation* op) { return isa(op); }, "function"); + if (failed(terminator)) { + return failure(); } - return success(); + return cast(*terminator); } -static FailureOr prepare(func::FuncOp func, const dd::Package& dd) { - if (!func.getBody().hasOneBlock()) { - return func.emitError() - << "QCO DD construction expects a single-block function body"; +template +static LogicalResult walkFunction(func::FuncOp func, WalkState& walk, + StateDD& state) { + auto returnOp = walkFunctionBody(func, walk, state); + if (failed(returnOp)) { + return failure(); } + return validateReturn(*returnOp, *walk.qubits, *walk.tensors); +} +namespace { +struct PreparedState { QubitMap qubits; - for (StaticOp staticOp : func.getBody().front().getOps()) { - const auto q = static_cast(staticOp.getIndex()); - qubits.bind(staticOp.getQubit(), q); - qubits.numQubits = std::max(qubits.numQubits, static_cast(q) + 1); + TensorMap tensors; +}; +} // namespace + +static FailureOr +prepare(func::FuncOp func, const dd::Package& dd, const DDBindings& bindings, + const bool bindEntryAllocations = false) { + if (func.isDeclaration()) { + return func.emitError() << "QCO DD construction requires a function body"; + } + + PreparedState prepared; + QubitMap& qubits = prepared.qubits; + for (Block& block : func.getBody()) { + for (StaticOp staticOp : block.getOps()) { + const auto index = static_cast(staticOp.getIndex()); + if (index >= dd::Package::MAX_POSSIBLE_QUBITS) { + return staticOp.emitError() + << "static qubit index exceeds the supported qubit range"; + } + const auto q = static_cast(index); + qubits.bind(staticOp.getQubit(), q); + qubits.numQubits = std::max(qubits.numQubits, static_cast(q) + 1); + } } if (qubits.numQubits == 0) { - qc::Qubit next = 0; + size_t next = 0; for (Value arg : func.getArguments()) { - if (!isa(arg.getType())) { - continue; + if (isa(arg.getType())) { + if (next >= dd::Package::MAX_POSSIBLE_QUBITS) { + return func.emitError() + << "QCO function exceeds the supported qubit range"; + } + qubits.bind(arg, static_cast(next++)); + } else if (isQTensorType(arg.getType())) { + const auto type = cast(arg.getType()); + int64_t size = type.getDimSize(0); + if (type.isDynamicDim(0)) { + const auto binding = bindings.find(arg); + if (binding == bindings.end() || !isa(binding->second)) { + return func.emitError() + << "dynamic qtensor arguments require an integer extent"; + } + size = cast(binding->second).getInt(); + if (size < 0) { + return func.emitError() + << "dynamic qtensor extent must be non-negative"; + } + } + const auto count = static_cast(size); + if (count > dd::Package::MAX_POSSIBLE_QUBITS - next) { + return func.emitError() + << "QCO function exceeds the supported qubit range"; + } + TensorSlots slots; + slots.reserve(count); + for (size_t i = 0; i < count; ++i) { + slots.emplace_back(static_cast(next++)); + } + prepared.tensors.bind(arg, std::move(slots)); } - qubits.bind(arg, next++); } qubits.numQubits = next; } - for (AllocOp alloc : func.getBody().front().getOps()) { - qubits.bind(alloc.getResult(), static_cast(qubits.numQubits++)); + if (bindEntryAllocations) { + for (AllocOp alloc : func.getBody().front().getOps()) { + qubits.bind(alloc.getResult(), + static_cast(qubits.numQubits++)); + } } if (dd.qubits() < qubits.numQubits) { return func.emitError() << "DD package has " << dd.qubits() << " qubits but function uses " << qubits.numQubits; } - return qubits; + return prepared; +} + +static bool mayAllocateQubits(func::FuncOp func, DenseSet& active) { + if (!active.insert(func).second) { + return true; + } + const auto activeGuard = llvm::make_scope_exit([&] { active.erase(func); }); + + bool allocation = false; + func.walk([&](Operation* op) { + if (allocation || isa(op)) { + allocation = true; + return; + } + auto call = dyn_cast(op); + if (!call) { + return; + } + auto callee = SymbolTable::lookupNearestSymbolFrom( + call, call.getCalleeAttr()); + allocation = + !callee || callee.isDeclaration() || mayAllocateQubits(callee, active); + }); + return allocation; } -FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd) { - auto qubitsOr = prepare(func, dd); - if (failed(qubitsOr)) { +FailureOr getNumQubits(func::FuncOp func) { + if (func.isDeclaration()) { + return func.emitError() << "QCO DD capacity requires a function body"; + } + + constexpr size_t maxQubits = dd::Package::MAX_POSSIBLE_QUBITS; + size_t required = 0; + LogicalResult result = success(); + const auto addQubits = [&](const size_t count, Operation* op) { + if (failed(result)) { + return; + } + if (count > maxQubits - required) { + result = op->emitError() + << "QCO function exceeds the supported qubit range"; + return; + } + required += count; + }; + + func.walk([&](StaticOp staticOp) { + if (failed(result)) { + return; + } + const auto index = static_cast(staticOp.getIndex()); + if (index >= maxQubits) { + result = staticOp.emitError() + << "static qubit index exceeds the supported qubit range"; + return; + } + required = std::max(required, index + 1U); + }); + + if (required == 0) { + for (Value arg : func.getArguments()) { + if (isa(arg.getType())) { + addQubits(1, func); + continue; + } + if (!isQTensorType(arg.getType())) { + continue; + } + const auto type = cast(arg.getType()); + if (type.isDynamicDim(0)) { + return func.emitError() + << "dynamic qtensor arguments prevent static DD capacity " + "calculation"; + } + addQubits(static_cast(type.getDimSize(0)), func); + } + } + + Block* entry = &func.getBody().front(); + func.walk([&](AllocOp alloc) { + if (alloc->getBlock() != entry) { + result = alloc.emitError() + << "quantum allocations outside the entry block prevent " + "static DD capacity calculation"; + return; + } + addQubits(1, alloc); + }); + func.walk([&](qtensor::AllocOp alloc) { + if (failed(result)) { + return; + } + if (alloc->getBlock() != entry) { + result = alloc.emitError() + << "quantum allocations outside the entry block prevent " + "static DD capacity calculation"; + return; + } + const auto attr = mqt::valueToConstantAttr(alloc.getSize()); + const auto integer = attr ? dyn_cast(*attr) : IntegerAttr{}; + if (!integer) { + result = alloc.emitError() + << "dynamic qtensor allocation prevents static DD capacity " + "calculation"; + return; + } + const int64_t count = integer.getInt(); + if (count <= 0) { + result = alloc.emitError() + << "qtensor allocation size must be a positive integer"; + return; + } + addQubits(static_cast(count), alloc); + }); + + DenseSet active; + func.walk([&](func::CallOp call) { + if (failed(result)) { + return; + } + auto callee = SymbolTable::lookupNearestSymbolFrom( + call, call.getCalleeAttr()); + if (!callee || callee.isDeclaration() || + mayAllocateQubits(callee, active)) { + result = call.emitError() + << "calls that may allocate qubits prevent static DD capacity " + "calculation"; + } + }); + + if (failed(result)) { return failure(); } - QubitMap qubits = std::move(*qubitsOr); + return required; +} + +FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd, + const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings, /*bindEntryAllocations=*/true); + if (failed(prepared)) { + return failure(); + } + QubitMap qubits = std::move(prepared->qubits); + TensorMap tensors = std::move(prepared->tensors); ClassicalEnv classical; - WalkState walkState{ - .qubits = &qubits, .classical = &classical, .dd = &dd, .rng = nullptr}; + if (failed(applyBindings(func, bindings, classical))) { + return failure(); + } + WalkState walkState{.qubits = &qubits, + .tensors = &tensors, + .classical = &classical, + .dd = &dd, + .rng = nullptr}; + walkState.activeCalls.insert(func.getOperation()); dd::MatrixDD state = qubits.numQubits == 0 @@ -976,27 +2529,42 @@ FailureOr buildFunctionality(func::FuncOp func, dd::Package& dd) { static FailureOr simulateImpl(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, - const QubitMap& preparedQubits, std::mt19937_64* rng, + const PreparedState& prepared, std::mt19937_64* rng, + const DDBindings& bindings, const DenseSet* deferredMeasurements = nullptr, - ClassicalEnv* finalClassical = nullptr) { + ClassicalEnv* finalClassical = nullptr, + const DeallocationMode deallocationMode = DeallocationMode::Apply, + const bool validateQuantumReturn = true) { const size_t inputQubits = in.isTerminal() ? 0U : static_cast(in.p->v) + 1U; - if (inputQubits < preparedQubits.numQubits) { + if (inputQubits < prepared.qubits.numQubits) { dd.decRef(in); return func.emitError() << "input state has " << inputQubits << " qubits but function uses " - << preparedQubits.numQubits; + << prepared.qubits.numQubits; } - QubitMap qubits = preparedQubits; + QubitMap qubits = prepared.qubits; + qubits.numQubits = inputQubits; + TensorMap tensors = prepared.tensors; ClassicalEnv classical; + if (failed(applyBindings(func, bindings, classical))) { + dd.decRef(in); + return failure(); + } WalkState walkState{.qubits = &qubits, + .tensors = &tensors, .classical = &classical, .dd = &dd, .rng = rng, - .deferredMeasurements = deferredMeasurements}; + .deferredMeasurements = deferredMeasurements, + .deallocationMode = deallocationMode}; + walkState.activeCalls.insert(func.getOperation()); dd::VectorDD state = in; - if (failed(walkFunction(func, walkState, state))) { + auto returnOp = walkFunctionBody(func, walkState, state); + if (failed(returnOp) || + (validateQuantumReturn && + failed(validateReturn(*returnOp, qubits, tensors)))) { dd.decRef(state); return failure(); } @@ -1007,13 +2575,133 @@ simulateImpl(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, } FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, - dd::Package& dd, std::mt19937_64& rng) { - auto qubits = prepare(func, dd); - if (failed(qubits)) { + dd::Package& dd, const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { dd.decRef(in); return failure(); } - return simulateImpl(func, in, dd, *qubits, &rng); + return simulateImpl(func, in, dd, *prepared, nullptr, bindings); +} + +FailureOr simulate(func::FuncOp func, const dd::VectorDD& in, + dd::Package& dd, std::mt19937_64& rng, + const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + dd.decRef(in); + return failure(); + } + return simulateImpl(func, in, dd, *prepared, &rng, bindings); +} + +dd::MatrixDD makeDensityMatrix(const dd::VectorDD& state, + const size_t numQubits, dd::Package& dd) { + if (numQubits > dd.qubits()) { + throw std::invalid_argument( + "numQubits exceeds the capacity of the DD package"); + } + if (!state.isTerminal() && std::cmp_greater_equal(state.p->v, numQubits)) { + throw std::invalid_argument( + "numQubits does not cover the input state's highest qubit"); + } + std::map, dd::mCachedEdge> cache; + const auto root = buildDensityMatrix( + state, state, static_cast(numQubits) - 1, dd, cache); + const auto density = dd::MatrixDD{.p = root.p, .w = dd.cn.lookup(root.w)}; + dd.incRef(density); + return density; +} + +namespace { +struct DensitySimulationResult { + dd::MatrixDD matrix; + size_t numQubits; +}; +} // namespace + +static LogicalResult validateDensityInput(func::FuncOp func, + const dd::MatrixDD& in, + const PreparedState& prepared) { + if (in.isTerminal() || + static_cast(in.p->v) < prepared.qubits.numQubits) { + return success(); + } + return func.emitError() << "input density matrix has " + << static_cast(in.p->v) + 1U + << " qubits but function uses " + << prepared.qubits.numQubits; +} + +static FailureOr simulateDensityImpl( + func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd, + const PreparedState& prepared, std::mt19937_64* rng, + const DDBindings& bindings, + const DenseSet* deferredMeasurements = nullptr, + ClassicalEnv* finalClassical = nullptr, + const DeallocationMode deallocationMode = DeallocationMode::Apply) { + if (failed(validateDensityInput(func, in, prepared))) { + dd.decRef(in); + return failure(); + } + + QubitMap qubits = prepared.qubits; + TensorMap tensors = prepared.tensors; + ClassicalEnv classical; + if (failed(applyBindings(func, bindings, classical))) { + dd.decRef(in); + return failure(); + } + WalkState walkState{.qubits = &qubits, + .tensors = &tensors, + .classical = &classical, + .dd = &dd, + .rng = rng, + .deferredMeasurements = deferredMeasurements, + .deallocationMode = deallocationMode}; + walkState.activeCalls.insert(func.getOperation()); + + DensityState state{in}; + if (failed(walkFunction(func, walkState, state))) { + dd.decRef(state.matrix); + return failure(); + } + if (finalClassical != nullptr) { + *finalClassical = std::move(classical); + } + return DensitySimulationResult{.matrix = state.matrix, + .numQubits = qubits.numQubits}; +} + +FailureOr simulateDensity(func::FuncOp func, + const dd::MatrixDD& in, dd::Package& dd, + const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + dd.decRef(in); + return failure(); + } + auto result = simulateDensityImpl(func, in, dd, *prepared, nullptr, bindings); + if (failed(result)) { + return failure(); + } + return result->matrix; +} + +FailureOr simulateDensity(func::FuncOp func, + const dd::MatrixDD& in, dd::Package& dd, + std::mt19937_64& rng, + const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + dd.decRef(in); + return failure(); + } + auto result = simulateDensityImpl(func, in, dd, *prepared, &rng, bindings); + if (failed(result)) { + return failure(); + } + return result->matrix; } static bool isOutputOnlyRegister(Value reg, ArrayRef outputs) { @@ -1026,22 +2714,102 @@ static bool isOutputOnlyRegister(Value reg, ArrayRef outputs) { }); } -static bool isDeferrableMeasurement(MeasureOp measure, Block* entry, - ArrayRef outputs) { - return measure->getBlock() == entry && - llvm::all_of(measure.getQubitOut().getUses(), - [](const OpOperand& use) { - return isa(use.getOwner()); - }) && - llvm::all_of(measure.getResult().getUses(), [&](const OpOperand& use) { - auto store = dyn_cast(use.getOwner()); - return store && isOutputOnlyRegister(store.getReg(), outputs); - }); +static bool hasOutputOnlyMeasurementResult(MeasureOp measure, + ArrayRef outputs, + const bool allowUnusedResult) { + Value result = measure.getResult(); + if (result.use_empty()) { + return allowUnusedResult; + } + return llvm::all_of(result.getUses(), [&](const OpOperand& use) { + auto store = dyn_cast(use.getOwner()); + return store && isOutputOnlyRegister(store.getReg(), outputs); + }); +} + +static std::optional getConstantTensorIndex(Value value) { + const auto attr = mqt::valueToConstantAttr(value); + const auto integer = attr ? dyn_cast(*attr) : IntegerAttr{}; + if (!integer || integer.getInt() < 0) { + return std::nullopt; + } + return static_cast(integer.getInt()); } -static void analyzeSampling(func::FuncOp func, Block* entry, - ArrayRef outputs, - DenseSet& active, SamplingPlan& plan) { +static bool hasOnlyTerminalQuantumUses(Value value, + std::optional tensorSlot, + func::FuncOp func, + ArrayRef outputs, + const bool allowUnusedResult, + DenseSet& visited) { + if (!visited.insert(value).second) { + return true; + } + return llvm::all_of(value.getUses(), [&](OpOperand& use) { + Operation* owner = use.getOwner(); + if (isa(owner)) { + return true; + } + if (auto measure = dyn_cast(owner)) { + return !tensorSlot && measure->getParentOfType() == func && + hasOutputOnlyMeasurementResult(measure, outputs, + allowUnusedResult) && + hasOnlyTerminalQuantumUses(measure.getQubitOut(), std::nullopt, + func, outputs, allowUnusedResult, + visited); + } + if (auto fromElements = dyn_cast(owner)) { + return !tensorSlot && + hasOnlyTerminalQuantumUses(fromElements.getResult(), + use.getOperandNumber(), func, outputs, + allowUnusedResult, visited); + } + if (auto insert = dyn_cast(owner)) { + const auto index = getConstantTensorIndex(insert.getIndex()); + if (!index) { + return false; + } + if (use.get() == insert.getScalar()) { + return !tensorSlot && + hasOnlyTerminalQuantumUses(insert.getResult(), index, func, + outputs, allowUnusedResult, visited); + } + return tensorSlot && *tensorSlot != *index && + hasOnlyTerminalQuantumUses(insert.getResult(), tensorSlot, func, + outputs, allowUnusedResult, visited); + } + if (auto extract = dyn_cast(owner)) { + const auto index = getConstantTensorIndex(extract.getIndex()); + if (!tensorSlot || !index) { + return false; + } + if (*tensorSlot == *index) { + return hasOnlyTerminalQuantumUses(extract.getResult(), std::nullopt, + func, outputs, allowUnusedResult, + visited); + } + return hasOnlyTerminalQuantumUses(extract.getOutTensor(), tensorSlot, + func, outputs, allowUnusedResult, + visited); + } + return false; + }); +} + +static bool isDeferrableMeasurement(MeasureOp measure, func::FuncOp func, + ArrayRef outputs, + const bool allowUnusedResult) { + if (!hasOutputOnlyMeasurementResult(measure, outputs, allowUnusedResult)) { + return false; + } + DenseSet visited; + return hasOnlyTerminalQuantumUses(measure.getQubitOut(), std::nullopt, func, + outputs, allowUnusedResult, visited); +} + +static void analyzeSampling(func::FuncOp func, ArrayRef outputs, + DenseSet& active, SamplingPlan& plan, + const bool allowUnusedMeasurementResults) { Operation* funcOp = func.getOperation(); if (!active.insert(funcOp).second) { plan.dynamic = true; @@ -1051,7 +2819,8 @@ static void analyzeSampling(func::FuncOp func, Block* entry, if (isa(op)) { plan.dynamic = true; } else if (auto measure = dyn_cast(op)) { - if (isDeferrableMeasurement(measure, entry, outputs)) { + if (isDeferrableMeasurement(measure, func, outputs, + allowUnusedMeasurementResults)) { plan.deferredMeasurements.insert(op); } else { plan.dynamic = true; @@ -1059,42 +2828,95 @@ static void analyzeSampling(func::FuncOp func, Block* entry, } else if (auto call = dyn_cast(op)) { auto callee = SymbolTable::lookupNearestSymbolFrom( call, call.getCalleeAttr()); - if (!callee.getBody().hasOneBlock()) { + if (!callee || callee.isDeclaration() || + !callee.getBody().hasOneBlock()) { plan.dynamic = true; } else { - analyzeSampling(callee, entry, outputs, active, plan); + // A callee return is not terminal for the entry function. Without + // following call-result dataflow, unused measurement results in + // callees must remain collapsing. + analyzeSampling(callee, outputs, active, plan, + /*allowUnusedMeasurementResults=*/false); } } }); active.erase(funcOp); } -static FailureOr getSamplingPlan(func::FuncOp func) { +static FailureOr +getSamplingPlan(func::FuncOp func, const bool validateHistogramOutputs = true, + const bool allowUnusedMeasurementResults = false) { Block& entry = func.getBody().front(); - auto returnOp = cast(entry.getTerminator()); SamplingPlan plan; - bool hasOther = false; - for (Value value : returnOp.getOperands()) { - if (isa(value.getType())) { - plan.outputs.push_back(value); - } else { - hasOther = true; + if (validateHistogramOutputs && !func.getBody().hasOneBlock()) { + if (llvm::any_of(func.getFunctionType().getResults(), + [](Type type) { return isa(type); })) { + return func.emitError() + << "QCO DD sampling does not support CBit results from " + "multi-block functions"; } + plan.dynamic = true; + return plan; + } + + if (validateHistogramOutputs && !isa(entry.getTerminator())) { + return func.emitError() + << "single-block QCO DD sampling requires func.return"; } - if (!plan.outputs.empty() && hasOther) { - return returnOp.emitError() + + bool hasOther = false; + func.getBody().walk([&](func::ReturnOp returnOp) { + for (Value value : returnOp.getOperands()) { + if (isa(value.getType())) { + plan.outputs.push_back(value); + } else { + hasOther = true; + } + } + }); + if (validateHistogramOutputs && !plan.outputs.empty() && hasOther) { + return entry.getTerminator()->emitError() << "QCO DD sampling does not support mixed CBit and non-CBit " "results"; } DenseSet active; - analyzeSampling(func, &entry, plan.outputs, active, plan); + analyzeSampling(func, plan.outputs, active, plan, + allowUnusedMeasurementResults); return plan; } +FailureOr simulateStatevector(func::FuncOp func, dd::Package& dd, + std::mt19937_64& rng) { + const auto requiredQubits = getNumQubits(func); + if (failed(requiredQubits)) { + return failure(); + } + if (dd.qubits() < *requiredQubits) { + return func.emitError() + << "DD package has " << dd.qubits() + << " qubits but function requires " << *requiredQubits; + } + const DDBindings bindings; + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + return failure(); + } + auto plan = getSamplingPlan(func, /*validateHistogramOutputs=*/false, + /*allowUnusedMeasurementResults=*/true); + if (failed(plan)) { + return failure(); + } + return simulateImpl(func, dd::makeZeroState(prepared->qubits.numQubits, dd), + dd, *prepared, &rng, bindings, + &plan->deferredMeasurements, nullptr, + DeallocationMode::PreserveAll, + /*validateQuantumReturn=*/false); +} + static FailureOr encodeOutcome(ArrayRef outputs, const ClassicalEnv& classical, - StringRef basis, size_t numQubits) { + StringRef basis) { if (outputs.empty()) { return basis.str(); } @@ -1110,9 +2932,8 @@ static FailureOr encodeOutcome(ArrayRef outputs, const auto& cell = (*reg->second)[index]; if (cell.value) { outcome.push_back(*cell.value ? '1' : '0'); - } else if (cell.deferredWire && basis.size() == numQubits && - *cell.deferredWire < numQubits) { - outcome.push_back(basis[numQubits - 1 - *cell.deferredWire]); + } else if (cell.deferredWire && *cell.deferredWire < basis.size()) { + outcome.push_back(basis[basis.size() - 1 - *cell.deferredWire]); } else { return emitError(value.getLoc()) << "returned CBit register element " << index << " is undefined"; @@ -1122,10 +2943,13 @@ static FailureOr encodeOutcome(ArrayRef outputs, return outcome; } -FailureOr> -sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng) { - auto qubits = prepare(func, dd); - if (failed(qubits)) { +static FailureOr> +sampleImpl(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, + const size_t shots, std::mt19937_64& rng, const DDBindings& bindings, + const bool preserveDeallocatedQubits = false) { + const auto inputGuard = llvm::make_scope_exit([&] { dd.decRef(in); }); + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { return failure(); } auto plan = getSamplingPlan(func); @@ -1133,15 +2957,22 @@ sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng) { return failure(); } + const size_t inputQubits = + in.isTerminal() ? 0U : static_cast(in.p->v) + 1U; + if (inputQubits < prepared->qubits.numQubits) { + return func.emitError() + << "input state has " << inputQubits << " qubits but function uses " + << prepared->qubits.numQubits; + } + std::map counts; if (shots == 0) { return counts; } - const size_t numQubits = qubits->numQubits; const auto record = [&](const ClassicalEnv& classical, StringRef basis) -> LogicalResult { - auto outcome = encodeOutcome(plan->outputs, classical, basis, numQubits); + auto outcome = encodeOutcome(plan->outputs, classical, basis); if (failed(outcome)) { return failure(); } @@ -1151,9 +2982,15 @@ sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng) { if (!plan->dynamic) { ClassicalEnv classical; + dd.incRef(in); + const auto deallocationMode = + preserveDeallocatedQubits || !plan->outputs.empty() + ? DeallocationMode::PreserveAll + : DeallocationMode::PreserveDeferred; auto state = - simulateImpl(func, dd::makeZeroState(numQubits, dd), dd, *qubits, - nullptr, &plan->deferredMeasurements, &classical); + simulateImpl(func, in, dd, *prepared, nullptr, bindings, + &plan->deferredMeasurements, &classical, deallocationMode, + /*validateQuantumReturn=*/!preserveDeallocatedQubits); if (failed(state)) { return failure(); } @@ -1168,8 +3005,13 @@ sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng) { for (size_t i = 0; i < shots; ++i) { ClassicalEnv classical; - auto state = simulateImpl(func, dd::makeZeroState(numQubits, dd), dd, - *qubits, &rng, nullptr, &classical); + dd.incRef(in); + auto state = simulateImpl( + func, in, dd, *prepared, &rng, bindings, nullptr, &classical, + preserveDeallocatedQubits || !plan->outputs.empty() + ? DeallocationMode::PreserveAll + : DeallocationMode::Apply, + /*validateQuantumReturn=*/!preserveDeallocatedQubits); if (failed(state)) { return failure(); } @@ -1184,4 +3026,125 @@ sample(func::FuncOp func, dd::Package& dd, size_t shots, std::mt19937_64& rng) { return counts; } +FailureOr> +sampleDensity(func::FuncOp func, const dd::MatrixDD& in, dd::Package& dd, + const size_t shots, std::mt19937_64& rng, + const DDBindings& bindings) { + const auto inputGuard = llvm::make_scope_exit([&] { dd.decRef(in); }); + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + return failure(); + } + if (failed(validateDensityInput(func, in, *prepared))) { + return failure(); + } + auto plan = getSamplingPlan(func); + if (failed(plan)) { + return failure(); + } + + std::map counts; + if (shots == 0) { + return counts; + } + + const auto record = [&](const ClassicalEnv& classical, + const StringRef basis) -> LogicalResult { + auto outcome = encodeOutcome(plan->outputs, classical, basis); + if (failed(outcome)) { + return failure(); + } + ++counts[*outcome]; + return success(); + }; + + if (!plan->dynamic) { + ClassicalEnv classical; + dd.incRef(in); + const auto deallocationMode = plan->outputs.empty() + ? DeallocationMode::PreserveDeferred + : DeallocationMode::PreserveAll; + auto simulated = simulateDensityImpl(func, in, dd, *prepared, nullptr, + bindings, &plan->deferredMeasurements, + &classical, deallocationMode); + if (failed(simulated)) { + return failure(); + } + const auto simulatedGuard = + llvm::make_scope_exit([&] { dd.decRef(simulated->matrix); }); + for (size_t i = 0; i < shots; ++i) { + dd.incRef(simulated->matrix); + DensityState sampleState{simulated->matrix}; + const auto sampleGuard = + llvm::make_scope_exit([&] { dd.decRef(sampleState.matrix); }); + auto outcome = measureAllDensity(sampleState, simulated->numQubits, dd, + rng, func.getOperation()); + if (failed(outcome)) { + return failure(); + } + if (failed(record(classical, *outcome))) { + return failure(); + } + } + return counts; + } + + for (size_t i = 0; i < shots; ++i) { + ClassicalEnv classical; + dd.incRef(in); + auto simulated = simulateDensityImpl( + func, in, dd, *prepared, &rng, bindings, nullptr, &classical, + plan->outputs.empty() ? DeallocationMode::Apply + : DeallocationMode::PreserveAll); + if (failed(simulated)) { + return failure(); + } + DensityState sampleState{simulated->matrix}; + const auto sampleGuard = + llvm::make_scope_exit([&] { dd.decRef(sampleState.matrix); }); + std::string basis; + if (plan->outputs.empty()) { + auto measured = measureAllDensity(sampleState, simulated->numQubits, dd, + rng, func.getOperation()); + if (failed(measured)) { + return failure(); + } + basis = std::move(*measured); + } + if (failed(record(classical, basis))) { + return failure(); + } + } + return counts; +} + +FailureOr> +sample(func::FuncOp func, const dd::VectorDD& in, dd::Package& dd, + const size_t shots, std::mt19937_64& rng, const DDBindings& bindings) { + return sampleImpl(func, in, dd, shots, rng, bindings); +} + +FailureOr> +sample(func::FuncOp func, dd::Package& dd, const size_t shots, + std::mt19937_64& rng, const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + return failure(); + } + return sampleImpl(func, dd::makeZeroState(prepared->qubits.numQubits, dd), dd, + shots, rng, bindings); +} + +FailureOr> +sampleAllQubits(func::FuncOp func, dd::Package& dd, const size_t shots, + std::mt19937_64& rng, const DDBindings& bindings) { + auto prepared = prepare(func, dd, bindings); + if (failed(prepared)) { + return failure(); + } + return sampleImpl(func, dd::makeZeroState(prepared->qubits.numQubits, dd), dd, + shots, rng, bindings, + /*preserveDeallocatedQubits=*/true); +} + } // namespace mlir::qco diff --git a/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp b/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp index 72bda4816f..61babadd21 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_dd_functionality.cpp @@ -8,6 +8,8 @@ * Licensed under the MIT License */ +#include "dd/ComplexValue.hpp" +#include "dd/DDDefinitions.hpp" #include "dd/FunctionalityConstruction.hpp" #include "dd/GateMatrixDefinitions.hpp" #include "dd/Node.hpp" @@ -21,11 +23,15 @@ #include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h" #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/Utils/DDFunctionality.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include #include #include +#include #include +#include +#include #include #include #include @@ -45,6 +51,7 @@ #include #include #include +#include #include #include #include @@ -62,8 +69,10 @@ class QCODDFunctionalityTest : public testing::Test { void SetUp() override { DialectRegistry registry; - registry.insert(); + registry + .insert(); context = std::make_unique(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); @@ -134,6 +143,13 @@ class QCODDFunctionalityTest : public testing::Test { EXPECT_TRUE(failed(sample(func, *dd, 1, rng))); } + static dd::MatrixDD makeZeroDensity(dd::Package& dd, const size_t numQubits) { + auto zero = dd::makeZeroState(numQubits, dd); + auto density = makeDensityMatrix(zero, numQubits, dd); + dd.decRef(zero); + return density; + } + void expectMlirSimulationFails(size_t numQubits, StringRef mlirCode) { auto mod = parseSourceString(mlirCode, context.get()); ASSERT_TRUE(mod); @@ -513,6 +529,26 @@ TEST_F(QCODDFunctionalityTest, RejectsUnmappedReturnedQubit) { failed(simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd, rng))); } +TEST_F(QCODDFunctionalityTest, RejectsStaticQubitBeyondDDRange) { + auto boundary = buildModule([](QCOProgramBuilder& b) { + b.sink(b.staticQubit(dd::Package::MAX_POSSIBLE_QUBITS - 1U)); + return b.intConstant(0); + }); + const auto numQubits = getNumQubits(mainFunc(*boundary)); + ASSERT_TRUE(succeeded(numQubits)); + EXPECT_EQ(*numQubits, dd::Package::MAX_POSSIBLE_QUBITS); + + auto mod = buildModule([](QCOProgramBuilder& b) { + b.sink(b.staticQubit(dd::Package::MAX_POSSIBLE_QUBITS)); + return b.intConstant(0); + }); + + auto dd = std::make_unique(1); + EXPECT_TRUE(failed(buildFunctionality(mainFunc(*mod), *dd))); + EXPECT_TRUE( + failed(simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd, rng))); +} + TEST_F(QCODDFunctionalityTest, SimulationConsumesInputReference) { auto valid = buildModule([](QCOProgramBuilder& b) { auto q = b.x(b.staticQubit(0)); @@ -564,6 +600,207 @@ TEST_F(QCODDFunctionalityTest, SimulationConsumesInputReference) { EXPECT_TRUE(zeroQubitDd->getRootSet().empty()); } +TEST_F(QCODDFunctionalityTest, + SimulationPreservesWiderInputAcrossRuntimeAllocation) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%q0: !qco.qubit) { + %q1 = qco.alloc : !qco.qubit + %q2 = qco.x %q1 : !qco.qubit -> !qco.qubit + qco.sink %q0 : !qco.qubit + qco.sink %q2 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(3); + auto input = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 1), + dd::makeZeroState(2, *dd)); + const auto output = simulate(mainFunc(*mod), input, *dd, rng); + ASSERT_TRUE(succeeded(output)); + + auto expected = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 1), + dd::makeZeroState(3, *dd)); + expected = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 2), + expected); + EXPECT_EQ(output->getVector(), expected.getVector()); + dd->decRef(*output); + dd->decRef(expected); +} + +TEST_F(QCODDFunctionalityTest, ConstructsDensityMatrix) { + auto dd = std::make_unique(1); + auto plus = dd::makeBasisState( + 1, std::vector{dd::BasisStates::plus}, *dd); + EXPECT_THROW(static_cast(makeDensityMatrix(plus, 0, *dd)), + std::invalid_argument); + EXPECT_THROW(static_cast(makeDensityMatrix(plus, 2, *dd)), + std::invalid_argument); + EXPECT_TRUE(dd->getRootSet().empty()); + const auto density = makeDensityMatrix(plus, 1, *dd); + + const auto matrix = density.getMatrix(1); + for (const auto& row : matrix) { + for (const auto& entry : row) { + EXPECT_NEAR(entry.real(), 0.5, 1e-12); + EXPECT_NEAR(entry.imag(), 0.0, 1e-12); + } + } + EXPECT_EQ(dd->getRootSet().at(plus), 1U); + EXPECT_EQ(dd->getRootSet().at(density), 1U); + + dd->decRef(density); + dd->decRef(plus); + EXPECT_TRUE(dd->getRootSet().empty()); + EXPECT_TRUE(dd->getRootSet().empty()); +} + +TEST_F(QCODDFunctionalityTest, ConstructsDensityMatrixWithComplexPhases) { + auto dd = std::make_unique(1); + const auto right = dd::makeBasisState( + 1, std::vector{dd::BasisStates::right}, *dd); + const auto density = makeDensityMatrix(right, 1, *dd); + + const auto matrix = density.getMatrix(1); + EXPECT_NEAR(matrix[0][0].real(), 0.5, 1e-12); + EXPECT_NEAR(matrix[0][1].imag(), -0.5, 1e-12); + EXPECT_NEAR(matrix[1][0].imag(), 0.5, 1e-12); + EXPECT_NEAR(matrix[1][1].real(), 0.5, 1e-12); + + dd->decRef(density); + dd->decRef(right); +} + +TEST_F(QCODDFunctionalityTest, DensityUsesFunctionExtentForSkippedLevels) { + auto mod = buildModule([](QCOProgramBuilder& b) { + b.sink(b.staticQubit(0)); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + auto maximallyMixed = dd::MatrixDD::one(); + maximallyMixed.w = dd->cn.lookup(dd::ComplexValue{0.5, 0.0}); + dd->incRef(maximallyMixed); + const auto result = simulateDensity(mainFunc(*mod), maximallyMixed, *dd); + ASSERT_TRUE(succeeded(result)); + + const auto matrix = result->getMatrix(1); + EXPECT_NEAR(matrix[0][0].real(), 0.5, 1e-12); + EXPECT_NEAR(matrix[0][1].real(), 0.0, 1e-12); + EXPECT_NEAR(matrix[1][0].real(), 0.0, 1e-12); + EXPECT_NEAR(matrix[1][1].real(), 0.5, 1e-12); + dd->decRef(*result); +} + +TEST_F(QCODDFunctionalityTest, DensityIgnoresUnboundGlobalPhase) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%theta: f64) { + %q = qco.static 0 : !qco.qubit + %plus = qco.h %q : !qco.qubit -> !qco.qubit + qco.gphase(%theta) + qco.sink %plus : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + const auto result = + simulateDensity(mainFunc(*mod), makeZeroDensity(*dd, 1), *dd); + ASSERT_TRUE(succeeded(result)); + const auto matrix = result->getMatrix(1); + for (const auto& row : matrix) { + for (const auto& entry : row) { + EXPECT_NEAR(entry.real(), 0.5, 1e-12); + EXPECT_NEAR(entry.imag(), 0.0, 1e-12); + } + } + dd->decRef(*result); +} + +TEST_F(QCODDFunctionalityTest, DensityReferencesAreConsumedOnAllPaths) { + auto valid = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + b.sink(q); + return b.intConstant(0); + }); + auto tooWide = buildModule([](QCOProgramBuilder& b) { + b.sink(b.staticQubit(0)); + b.sink(b.staticQubit(1)); + return b.intConstant(0); + }); + auto measured = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + std::tie(q, std::ignore) = b.measure(q); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(valid); + ASSERT_TRUE(tooWide); + ASSERT_TRUE(measured); + + auto dd = std::make_unique(1); + auto& roots = dd->getRootSet(); + auto simulated = + simulateDensity(mainFunc(*valid), makeZeroDensity(*dd, 1), *dd); + ASSERT_TRUE(succeeded(simulated)); + EXPECT_EQ(roots.size(), 1U); + EXPECT_EQ(roots.at(*simulated), 1U); + dd->decRef(*simulated); + EXPECT_TRUE(roots.empty()); + + EXPECT_TRUE(failed( + simulateDensity(mainFunc(*tooWide), makeZeroDensity(*dd, 1), *dd))); + EXPECT_TRUE(roots.empty()); + + std::mt19937_64 rng(9); + auto overWideDd = std::make_unique(2); + auto& overWideRoots = overWideDd->getRootSet(); + EXPECT_TRUE(failed(simulateDensity( + mainFunc(*valid), makeZeroDensity(*overWideDd, 2), *overWideDd))); + EXPECT_TRUE(overWideRoots.empty()); + EXPECT_TRUE(failed(sampleDensity( + mainFunc(*valid), makeZeroDensity(*overWideDd, 2), *overWideDd, 1, rng))); + EXPECT_TRUE(overWideRoots.empty()); + EXPECT_TRUE(failed(sampleDensity( + mainFunc(*valid), makeZeroDensity(*overWideDd, 2), *overWideDd, 0, rng))); + EXPECT_TRUE(overWideRoots.empty()); + + auto histogram = + sampleDensity(mainFunc(*valid), makeZeroDensity(*dd, 1), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 4}})); + EXPECT_TRUE(roots.empty()); + + EXPECT_TRUE(failed( + sampleDensity(mainFunc(*tooWide), makeZeroDensity(*dd, 1), *dd, 1, rng))); + EXPECT_TRUE(roots.empty()); + + histogram = + sampleDensity(mainFunc(*valid), makeZeroDensity(*dd, 1), *dd, 0, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_TRUE(histogram->empty()); + EXPECT_TRUE(roots.empty()); + + auto normalized = makeZeroDensity(*dd, 1); + auto invalid = normalized; + invalid.w = dd->cn.lookup(static_cast(invalid.w) * 0.5); + dd->incRef(invalid); + dd->decRef(normalized); + EXPECT_TRUE(failed(simulateDensity(mainFunc(*measured), invalid, *dd, rng))); + EXPECT_TRUE(roots.empty()); +} + TEST_F(QCODDFunctionalityTest, SimulateMeasureCollapsesLikePackage) { auto mod = buildModule([](QCOProgramBuilder& b) { auto q = b.h(b.staticQubit(0)); @@ -632,29 +869,12 @@ TEST_F(QCODDFunctionalityTest, SimulateIfConstantBranches) { ASSERT_TRUE(thenMod); ASSERT_TRUE(elseMod); - auto dd = std::make_unique(1); - EXPECT_TRUE(failed(buildFunctionality(mainFunc(*thenMod), *dd))); - EXPECT_TRUE(failed(buildFunctionality(mainFunc(*elseMod), *dd))); - std::mt19937_64 rng(0); - auto zero = dd::makeZeroState(1, *dd); - auto one = dd->applyOperation( - dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 0), - dd::makeZeroState(1, *dd)); - - const auto thenOut = - simulate(mainFunc(*thenMod), dd::makeZeroState(1, *dd), *dd, rng); - ASSERT_TRUE(succeeded(thenOut)); - EXPECT_EQ(thenOut->getVector(), one.getVector()); + qc::QuantumComputation thenQc(1); + thenQc.x(0); + expectEqualToQc(mainFunc(*thenMod), thenQc); - const auto elseOut = - simulate(mainFunc(*elseMod), dd::makeZeroState(1, *dd), *dd, rng); - ASSERT_TRUE(succeeded(elseOut)); - EXPECT_EQ(elseOut->getVector(), zero.getVector()); - - dd->decRef(*thenOut); - dd->decRef(*elseOut); - dd->decRef(zero); - dd->decRef(one); + const qc::QuantumComputation elseQc(1); + expectEqualToQc(mainFunc(*elseMod), elseQc); } TEST_F(QCODDFunctionalityTest, SimulateIndexSwitchBranches) { @@ -681,29 +901,12 @@ TEST_F(QCODDFunctionalityTest, SimulateIndexSwitchBranches) { ASSERT_TRUE(caseMod); ASSERT_TRUE(defaultMod); - auto dd = std::make_unique(1); - EXPECT_TRUE(failed(buildFunctionality(mainFunc(*caseMod), *dd))); - EXPECT_TRUE(failed(buildFunctionality(mainFunc(*defaultMod), *dd))); - std::mt19937_64 rng(0); - auto zero = dd::makeZeroState(1, *dd); - auto one = dd->applyOperation( - dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 0), - dd::makeZeroState(1, *dd)); - - const auto caseOut = - simulate(mainFunc(*caseMod), dd::makeZeroState(1, *dd), *dd, rng); - ASSERT_TRUE(succeeded(caseOut)); - EXPECT_EQ(caseOut->getVector(), one.getVector()); - - const auto defaultOut = - simulate(mainFunc(*defaultMod), dd::makeZeroState(1, *dd), *dd, rng); - ASSERT_TRUE(succeeded(defaultOut)); - EXPECT_EQ(defaultOut->getVector(), zero.getVector()); + qc::QuantumComputation caseQc(1); + caseQc.x(0); + expectEqualToQc(mainFunc(*caseMod), caseQc); - dd->decRef(*caseOut); - dd->decRef(*defaultOut); - dd->decRef(zero); - dd->decRef(one); + const qc::QuantumComputation defaultQc(1); + expectEqualToQc(mainFunc(*defaultMod), defaultQc); } TEST_F(QCODDFunctionalityTest, SimulateMeasureFeedsIf) { @@ -1014,6 +1217,23 @@ TEST_F(QCODDFunctionalityTest, SampleHadamardApproximatelyBalanced) { EXPECT_NEAR(static_cast(hist->at("0")), shots / 2.0, 150.0); } +TEST_F(QCODDFunctionalityTest, DensitySamplingDefersTerminalMeasurement) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.x(b.staticQubit(0)); + std::tie(q, std::ignore) = b.measure(q); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(7); + const auto histogram = + sampleDensity(mainFunc(*mod), makeZeroDensity(*dd, 1), *dd, 16, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 16}})); +} + TEST_F(QCODDFunctionalityTest, SampleResetUsesDynamicSampling) { auto mod = buildModule([](QCOProgramBuilder& b) { auto q = b.reset(b.x(b.staticQubit(0))); @@ -1066,6 +1286,8 @@ TEST_F(QCODDFunctionalityTest, SampleHandlesZeroShotsAndSimulationFailure) { const auto empty = sample(mainFunc(*unitary), *dd, 0, rng); ASSERT_TRUE(succeeded(empty)); EXPECT_TRUE(empty->empty()); + EXPECT_TRUE( + failed(sample(mainFunc(*unitary), dd::VectorDD::one(), *dd, 0, rng))); auto dynamic = parseSourceString(R"mlir( module { @@ -1137,6 +1359,64 @@ TEST_F(QCODDFunctionalityTest, EmbedsWideLocalMatrixWithoutRegisterLimit) { expectEqualToQc(mainFunc(*mod), qc); } +TEST_F(QCODDFunctionalityTest, RejectsUnsupportedOrUnboundClassicalOperations) { + for (const StringRef source : { + R"mlir(module { + func.func @main(%c: i1) { + %q = qco.static 0 : !qco.qubit + %bad = arith.index_castui %c : i1 to index + qco.sink %q : !qco.qubit + return + } + })mlir", + R"mlir(module { + func.func @main(%unmapped: i1) { + %q = qco.static 0 : !qco.qubit + %true = arith.constant true + %bad = arith.andi %unmapped, %true : i1 + qco.sink %q : !qco.qubit + return + } + })mlir", + R"mlir(module { + func.func @main(%unmapped: index) { + %q = qco.static 0 : !qco.qubit + %one = arith.constant 1 : index + %bad = arith.ori %unmapped, %one : index + qco.sink %q : !qco.qubit + return + } + })mlir", + R"mlir(module { + func.func @main() { + %bad = arith.constant 1.0 : f32 + return + } + })mlir", + R"mlir(module { + func.func @main() { + %one = arith.constant 1 : i32 + %bad = arith.sitofp %one : i32 to f32 + return + } + })mlir", + R"mlir(module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %one = arith.constant 1.0 : f64 + %bad = math.erf %one : f64 + qco.sink %q : !qco.qubit + return + } + })mlir"}) { + auto mod = parseSourceString(source, context.get()); + ASSERT_TRUE(mod); + auto dd = std::make_unique(1); + std::mt19937_64 rng(1); + EXPECT_TRUE( + failed(simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd, rng))); + } +} TEST_F(QCODDFunctionalityTest, RejectsUnmappedClassicalControl) { for (const StringRef source : {R"mlir( module { @@ -1254,6 +1534,45 @@ TEST_F(QCODDFunctionalityTest, BindsClassicalIndexResults) { dd->decRef(expected); } +TEST_F(QCODDFunctionalityTest, RejectsUnboundClassicalRegionResults) { + for (const StringRef source : { + R"mlir(module { + func.func @main(%unmapped: i1) { + %q = qco.static 0 : !qco.qubit + %true = arith.constant true + %result, %out = qco.if %true args(%arg = %q) + -> (i1, !qco.qubit) { + qco.yield %unmapped, %arg : i1, !qco.qubit + } else args(%arg = %q) { + qco.yield %true, %arg : i1, !qco.qubit + } + qco.sink %out : !qco.qubit + return + } + })mlir", + R"mlir(module { + func.func @main(%unmapped: index) { + %q = qco.static 0 : !qco.qubit + %true = arith.constant true + %zero = arith.constant 0 : index + %result, %out = qco.if %true args(%arg = %q) + -> (index, !qco.qubit) { + qco.yield %unmapped, %arg : index, !qco.qubit + } else args(%arg = %q) { + qco.yield %zero, %arg : index, !qco.qubit + } + qco.sink %out : !qco.qubit + return + } + })mlir"}) { + auto mod = parseSourceString(source, context.get()); + ASSERT_TRUE(mod); + auto dd = std::make_unique(1); + std::mt19937_64 rng(1); + EXPECT_TRUE( + failed(simulate(mainFunc(*mod), dd::makeZeroState(1, *dd), *dd, rng))); + } +} TEST_F(QCODDFunctionalityTest, Rejects) { { auto mod = buildModule([](QCOProgramBuilder& b) { @@ -1390,7 +1709,6 @@ TEST_F(QCODDFunctionalityTest, Rejects) { } } )mlir"); - for (const bool composed : {false, true}) { auto wideModifier = buildModule([composed](QCOProgramBuilder& b) { SmallVector qubits; @@ -1414,21 +1732,6 @@ TEST_F(QCODDFunctionalityTest, Rejects) { auto wideDD = std::make_unique(11); EXPECT_TRUE(failed(buildFunctionality(mainFunc(*wideModifier), *wideDD))); } - - OwningOpRef multi = - ModuleOp::create(UnknownLoc::get(context.get())); - OpBuilder builder(context.get()); - builder.setInsertionPointToStart(multi->getBody()); - auto func = func::FuncOp::create(builder, multi->getLoc(), "main", - builder.getFunctionType({}, {})); - auto* entry = func.addEntryBlock(); - auto* second = func.addBlock(); - builder.setInsertionPointToStart(entry); - func::ReturnOp::create(builder, func.getLoc()); - builder.setInsertionPointToStart(second); - func::ReturnOp::create(builder, func.getLoc()); - auto dd = std::make_unique(0); - EXPECT_TRUE(failed(buildFunctionality(func, *dd))); } TEST_F(QCODDFunctionalityTest, SimulateScfForAndFuncCallWithClassicalValues) { @@ -1595,7 +1898,159 @@ TEST_F(QCODDFunctionalityTest, ScfForSnapshotsYieldedInductionValue) { expectSimulatesFromZero(mainFunc(*mod), false); } +TEST_F(QCODDFunctionalityTest, InterpretsMultiBlockConcreteCFG) { + auto mod = parseSourceString(R"mlir( + module { + func.func @flip(%q: !qco.qubit, %condition: i1) -> !qco.qubit { + cf.cond_br %condition, ^then(%q : !qco.qubit), + ^else(%q : !qco.qubit) + ^then(%arg: !qco.qubit): + %x = qco.x %arg : !qco.qubit -> !qco.qubit + cf.br ^merge(%x : !qco.qubit) + ^else(%else_arg: !qco.qubit): + cf.br ^merge(%else_arg : !qco.qubit) + ^merge(%result: !qco.qubit): + return %result : !qco.qubit + } + func.func @main() { + %q = qco.static 0 : !qco.qubit + %true = arith.constant true + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c3 = arith.constant 3 : index + cf.br ^loop(%q, %c0 : !qco.qubit, index) + ^loop(%carried: !qco.qubit, %i: index): + %continue = arith.cmpi slt, %i, %c3 : index + cf.cond_br %continue, + ^body(%carried, %i : !qco.qubit, index), + ^run(%carried, %true : !qco.qubit, i1) + ^body(%body_arg: !qco.qubit, %body_i: index): + %loop_x = qco.x %body_arg : !qco.qubit -> !qco.qubit + %next = arith.addi %body_i, %c1 : index + cf.br ^loop(%loop_x, %next : !qco.qubit, index) + ^run(%arg: !qco.qubit, %condition: i1): + %called = func.call @flip(%arg, %condition) + : (!qco.qubit, i1) -> !qco.qubit + %result = scf.execute_region -> !qco.qubit { + %selected = arith.constant 2 : i32 + %missing = arith.constant 7 : i32 + cf.switch %selected : i32, [ + default: ^unexpected(%called : !qco.qubit), + 2: ^matched(%called : !qco.qubit) + ] + ^matched(%matched_arg: !qco.qubit): + cf.switch %missing : i32, [ + default: ^done(%matched_arg : !qco.qubit), + 1: ^unexpected(%matched_arg : !qco.qubit) + ] + ^unexpected(%unexpected_arg: !qco.qubit): + scf.yield %unexpected_arg : !qco.qubit + ^done(%done_arg: !qco.qubit): + %done_x = qco.x %done_arg : !qco.qubit -> !qco.qubit + scf.yield %done_x : !qco.qubit + } + qco.sink %result : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + qc::QuantumComputation qc(1); + qc.x(0); + qc.x(0); + qc.x(0); + qc.x(0); + qc.x(0); + expectEqualToQc(mainFunc(*mod), qc); +} + +TEST_F(QCODDFunctionalityTest, SamplesMultiBlockFallbackBasis) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + cf.br ^next(%q : !qco.qubit) + ^next(%arg: !qco.qubit): + %result = qco.x %arg : !qco.qubit -> !qco.qubit + qco.sink %result : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(3); + const auto histogram = sample(mainFunc(*mod), *dd, 8, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 8}})); +} + +TEST_F(QCODDFunctionalityTest, RejectsMultiBlockCBitSampling) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main() -> !cbit.reg<1> { + %reg = cbit.alloc(#cbit.init) : !cbit.reg<1> + cf.br ^return(%reg : !cbit.reg<1>) + ^return(%result: !cbit.reg<1>): + return %result : !cbit.reg<1> + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(0); + std::mt19937_64 rng(3); + EXPECT_TRUE(failed(sample(mainFunc(*mod), *dd, 1, rng))); +} + +TEST_F(QCODDFunctionalityTest, RejectsUnboundedConcreteCFG) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main() { + cf.br ^loop + ^loop: + cf.br ^loop + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(0); + EXPECT_TRUE(failed(buildFunctionality(mainFunc(*mod), *dd))); + EXPECT_TRUE(failed(simulate(mainFunc(*mod), dd::VectorDD::one(), *dd))); +} + TEST_F(QCODDFunctionalityTest, RejectsUnsupportedFuncCalls) { + auto selfRecursive = parseSourceString(R"mlir( + module { + func.func @main(%recurse: i1) { + scf.if %recurse { + %false = arith.constant false + func.call @main(%false) : (i1) -> () + } + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(selfRecursive); + auto selfRecursiveFunc = mainFunc(*selfRecursive); + DDBindings bindings; + bindings[selfRecursiveFunc.getArgument(0)] = + BoolAttr::get(context.get(), true); + auto zeroQubitDd = std::make_unique(0); + EXPECT_TRUE(failed(simulate(selfRecursiveFunc, dd::VectorDD::one(), + *zeroQubitDd, rng, bindings))); + EXPECT_TRUE(failed(simulateDensity(selfRecursiveFunc, + makeZeroDensity(*zeroQubitDd, 0), + *zeroQubitDd, bindings))); + auto recursive = parseSourceString(R"mlir( module { func.func @rec(%q: !qco.qubit) -> !qco.qubit { @@ -1805,6 +2260,77 @@ TEST_F(QCODDFunctionalityTest, SampleReturnsCBitRegistersInDeclaredOrder) { EXPECT_TRUE(dd->getRootSet().empty()); } +TEST_F(QCODDFunctionalityTest, InterpretsMinMaxAndCommonMathOperations) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.staticQubit(0); + Value minusOne = arith::ConstantIntOp::create(b, -1, 8); + Value two = arith::ConstantIntOp::create(b, 2, 8); + Value three = arith::ConstantIntOp::create(b, 3, 8); + SmallVector checks{ + arith::CmpIOp::create(b, arith::CmpIPredicate::eq, + arith::MaxSIOp::create(b, minusOne, three), + three), + arith::CmpIOp::create(b, arith::CmpIPredicate::eq, + arith::MinSIOp::create(b, minusOne, three), + minusOne), + arith::CmpIOp::create(b, arith::CmpIPredicate::eq, + arith::MaxUIOp::create(b, minusOne, two), + minusOne), + arith::CmpIOp::create(b, arith::CmpIPredicate::eq, + arith::MinUIOp::create(b, minusOne, two), two)}; + + const auto f64 = b.getF64Type(); + const auto constant = [&](const double value) -> Value { + return arith::ConstantFloatOp::create(b, f64, llvm::APFloat(value)); + }; + Value fm2 = constant(-2.0); + Value f0 = constant(0.0); + Value f1 = constant(1.0); + Value f12 = constant(1.2); + Value f18 = constant(1.8); + Value f2 = constant(2.0); + Value f4 = constant(4.0); + Value nan = constant(std::numeric_limits::quiet_NaN()); + const auto checkFloat = [&](Value actual, Value expected) { + checks.emplace_back(arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, + actual, expected)); + }; + checkFloat(arith::MaxNumFOp::create(b, nan, f2), f2); + checkFloat(arith::MinNumFOp::create(b, nan, f2), f2); + Value maximum = arith::MaximumFOp::create(b, nan, f2); + Value minimum = arith::MinimumFOp::create(b, nan, f2); + checks.emplace_back( + arith::CmpFOp::create(b, arith::CmpFPredicate::UNO, maximum, maximum)); + checks.emplace_back( + arith::CmpFOp::create(b, arith::CmpFPredicate::UNO, minimum, minimum)); + checkFloat(math::AbsFOp::create(b, fm2), f2); + checkFloat(math::CeilOp::create(b, f12), f2); + checkFloat(math::CosOp::create(b, f0), f1); + checkFloat(math::ExpOp::create(b, f0), f1); + checkFloat(math::FloorOp::create(b, f18), f1); + checkFloat(math::LogOp::create(b, f1), f0); + checkFloat(math::SinOp::create(b, f0), f0); + checkFloat(math::SqrtOp::create(b, f4), f2); + checkFloat(math::TanOp::create(b, f0), f0); + checkFloat(math::PowFOp::create(b, f2, f2), f4); + + Value all = checks.front(); + for (Value check : ArrayRef(checks).drop_front()) { + all = arith::AndIOp::create(b, all, check); + } + q = b.qcoIf( + all, q, [&](Value arg) { return b.x(arg); }, + [&](Value arg) { return arg; }); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + qc::QuantumComputation qc(1); + qc.x(0); + expectEqualToQc(mainFunc(*mod), qc); +} + TEST_F(QCODDFunctionalityTest, SampleRejectsUndefinedAndMixedResults) { auto undefined = buildModule([](QCOProgramBuilder& b) { return b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); @@ -1853,18 +2379,190 @@ TEST_F(QCODDFunctionalityTest, EXPECT_TRUE(dd->getRootSet().empty()); } -TEST_F(QCODDFunctionalityTest, SampleExecutesControlMeasurementPerShot) { +TEST_F(QCODDFunctionalityTest, DefersTensorMeasurementDespiteLaterUnrelatedOp) { auto mod = buildModule([](QCOProgramBuilder& b) { auto reg = b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); - auto q = b.x(b.staticQubit(0)); - Value bit; - std::tie(q, bit) = b.measure(q, reg, 0); - q = b.qcoIf( - bit, q, [&](Value arg) { return arg; }, [&](Value arg) { return arg; }); - q = b.x(q); - b.sink(q); - return reg; + auto q0 = b.h(b.allocQubit()); + auto q1 = b.allocQubit(); + std::tie(q0, std::ignore) = b.measure(q0, reg, 0); + auto tensor = b.qtensorFromElements({q0, q1}); + std::tie(tensor, q0) = b.qtensorExtract(tensor, 0); + tensor = b.qtensorInsert(q0, tensor, 0); + std::tie(tensor, q1) = b.qtensorExtract(tensor, 1); + q1 = b.x(q1); + tensor = b.qtensorInsert(q1, tensor, 1); + b.qtensorDealloc(tensor); + return reg; + }); + ASSERT_TRUE(mod); + + std::mt19937_64 rng(11); + auto singleDD = std::make_unique(2); + ASSERT_TRUE(succeeded(sample(mainFunc(*mod), *singleDD, 1, rng))); + const auto singleEvolutionLookups = + singleDD->matrixVectorMultiplication.getStats().lookups; + + auto dd = std::make_unique(2); + const auto histogram = sample(mainFunc(*mod), *dd, 64, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(histogram->at("0") + histogram->at("1"), 64U); + EXPECT_EQ(dd->matrixVectorMultiplication.getStats().lookups, + singleEvolutionLookups); + EXPECT_TRUE(dd->getRootSet().empty()); + + auto stateDD = std::make_unique(2); + auto expected = dd::makeZeroState(2, *stateDD); + expected = stateDD->applyOperation( + stateDD->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::H), 0), + expected); + expected = stateDD->applyOperation( + stateDD->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 1), + expected); + auto state = simulateStatevector(mainFunc(*mod), *stateDD, rng); + ASSERT_TRUE(succeeded(state)); + EXPECT_EQ(state->getVector(), expected.getVector()); + stateDD->decRef(*state); + stateDD->decRef(expected); +} + +TEST_F(QCODDFunctionalityTest, + StatevectorAllowsMixedClassicalAndNonClassicalResults) { + auto mod = buildModule([](QCOProgramBuilder& b) -> SmallVector { + auto reg = + b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); + auto q = b.h(b.staticQubit(0)); + std::tie(q, std::ignore) = b.measure(q, reg, 0); + b.sink(q); + return {reg, b.intConstant(0)}; + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + auto expected = dd::makeZeroState(1, *dd); + expected = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::H), 0), + expected); + std::mt19937_64 rng(19); + auto state = simulateStatevector(mainFunc(*mod), *dd, rng); + ASSERT_TRUE(succeeded(state)); + EXPECT_EQ(state->getVector(), expected.getVector()); + dd->decRef(*state); + dd->decRef(expected); +} + +TEST_F(QCODDFunctionalityTest, + StatevectorPreservesDeallocatedWiresBelowQuantumResults) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto low = b.qtensorAlloc(1); + auto high = b.x(b.allocQubit()); + b.qtensorDealloc(low); + return high; + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(23); + auto state = simulateStatevector(mainFunc(*mod), *dd, rng); + ASSERT_TRUE(succeeded(state)); + auto expected = dd::makeZeroState(2, *dd); + expected = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::X), 1), + expected); + EXPECT_EQ(state->getVector(), expected.getVector()); + dd->decRef(*state); + dd->decRef(expected); +} + +TEST_F(QCODDFunctionalityTest, + SampleAllQubitsPreservesDeallocatedWiresBelowQuantumResults) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto low = b.qtensorAlloc(1); + auto high = b.x(b.allocQubit()); + b.qtensorDealloc(low); + return high; + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(27); + const auto histogram = sampleAllQubits(mainFunc(*mod), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"10", 4}})); +} + +TEST_F(QCODDFunctionalityTest, UnusedMeasurementDoesNotRetainDeallocatedWire) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto tensor = b.qtensorAlloc(1); + Value q; + std::tie(tensor, q) = b.qtensorExtract(tensor, 0); + std::tie(q, std::ignore) = b.measure(q); + tensor = b.qtensorInsert(q, tensor, 0); + b.qtensorDealloc(tensor); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(29); + const auto histogram = sample(mainFunc(*mod), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"", 4}})); +} + +TEST_F(QCODDFunctionalityTest, + StatevectorDefersTerminalMeasurementWithUnusedResult) { + 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); + + auto dd = std::make_unique(1); + auto expected = dd::makeZeroState(1, *dd); + expected = dd->applyOperation( + dd->makeGateDD(dd::opToSingleQubitGateMatrix(qc::OpType::H), 0), + expected); + std::mt19937_64 rng(30); + auto state = simulateStatevector(mainFunc(*mod), *dd, rng); + ASSERT_TRUE(succeeded(state)); + EXPECT_EQ(state->getVector(), expected.getVector()); + dd->decRef(*state); + dd->decRef(expected); +} + +TEST_F(QCODDFunctionalityTest, SampleDefersAllocatedQubitMeasurement) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto reg = + b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); + auto q = b.x(b.allocQubit()); + std::tie(q, std::ignore) = b.measure(q, reg, 0); + b.sink(q); + return reg; + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(11); + const auto histogram = sample(mainFunc(*mod), *dd, 8, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 8}})); +} + +TEST_F(QCODDFunctionalityTest, SampleExecutesControlMeasurementPerShot) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto reg = + b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); + auto q = b.x(b.staticQubit(0)); + Value bit; + std::tie(q, bit) = b.measure(q, reg, 0); + q = b.qcoIf( + bit, q, [&](Value arg) { return arg; }, [&](Value arg) { return arg; }); + q = b.x(q); + b.sink(q); + return reg; }); ASSERT_TRUE(mod); @@ -1915,6 +2613,47 @@ TEST_F(QCODDFunctionalityTest, FuncCallSharesClassicalCBitStorage) { expectSimulatesFromZero(mainFunc(*mod), true); } +TEST_F(QCODDFunctionalityTest, RejectsUnsupportedClassicalMemRefs) { + for (const StringRef source : { + R"mlir(module { + func.func @main(%reg: memref) { + %value = memref.load %reg[] : memref + return + } + })mlir", + R"mlir(module { + func.func @main(%reg: memref) { + %value = arith.constant true + memref.store %value, %reg[] : memref + return + } + })mlir", + R"mlir(module { + func.func @main(%n: index) { + %reg = memref.alloc(%n) : memref + memref.dealloc %reg : memref + return + } + })mlir", + R"mlir(module { + func.func @main() { + %reg = memref.alloc() : memref<1xf32> + memref.dealloc %reg : memref<1xf32> + return + } + })mlir", + R"mlir(module { + func.func @main() { + %reg = memref.alloc() : memref<1xi1> + %value = arith.constant true + %i2 = arith.constant 2 : index + memref.store %value, %reg[%i2] : memref<1xi1> + return + } + })mlir"}) { + expectMlirSimulationFails(0, source); + } +} TEST_F(QCODDFunctionalityTest, SampleExecutesCalleeMeasurementBeforeCallerGate) { auto mod = parseSourceString(R"mlir( @@ -1950,6 +2689,22 @@ TEST_F(QCODDFunctionalityTest, EXPECT_EQ(dd->matrixVectorMultiplication.getStats().lookups, perShotLookups * 128U); EXPECT_TRUE(dd->getRootSet().empty()); + + auto stateDD = std::make_unique(1); + auto state = simulateStatevector(mainFunc(*mod), *stateDD, rng); + ASSERT_TRUE(succeeded(state)); + const auto vector = state->getVector(); + ASSERT_EQ(vector.size(), 2U); + EXPECT_NEAR(std::norm(vector[0]), 0.5, 1e-12); + EXPECT_NEAR(std::norm(vector[1]), 0.5, 1e-12); + stateDD->decRef(*state); + + const auto densityHistogram = + sampleDensity(mainFunc(*mod), makeZeroDensity(*dd, 1), *dd, 128, rng); + ASSERT_TRUE(succeeded(densityHistogram)); + ASSERT_EQ(densityHistogram->size(), 2U); + EXPECT_EQ(densityHistogram->at("0") + densityHistogram->at("1"), 128U); + EXPECT_TRUE(dd->getRootSet().empty()); } TEST_F(QCODDFunctionalityTest, SampleExecutesNestedMeasurementPerShot) { @@ -1993,4 +2748,901 @@ TEST_F(QCODDFunctionalityTest, SampleExecutesNestedMeasurementPerShot) { EXPECT_TRUE(dd->getRootSet().empty()); } +TEST_F(QCODDFunctionalityTest, SymbolicParametersUseBindings) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%theta: f64) { + %q = qco.static 0 : !qco.qubit + %twice = arith.addf %theta, %theta : f64 + %q1 = qco.rx(%twice) %q : !qco.qubit -> !qco.qubit + qco.gphase(%theta) + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + auto concrete = buildModule([](QCOProgramBuilder& b) { + auto q = b.rx(std::numbers::pi, b.staticQubit(0)); + b.gphase(std::numbers::pi / 2.0); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + ASSERT_TRUE(concrete); + + auto func = mainFunc(*mod); + DDBindings bindings; + bindings[func.getArgument(0)] = FloatAttr::get( + cast(func.getArgument(0).getType()), std::numbers::pi / 2.0); + + auto dd = std::make_unique(1); + auto actual = buildFunctionality(func, *dd, bindings); + auto expected = buildFunctionality(mainFunc(*concrete), *dd); + ASSERT_TRUE(succeeded(actual)); + ASSERT_TRUE(succeeded(expected)); + EXPECT_EQ(actual->getMatrix(1), expected->getMatrix(1)); + dd->decRef(*actual); + dd->decRef(*expected); + + std::mt19937_64 rng(5); + const auto histogram = sample(func, *dd, 8, rng, bindings); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 8}})); + + auto density = simulateDensity(func, makeZeroDensity(*dd, 1), *dd, bindings); + ASSERT_TRUE(succeeded(density)); + const auto densityMatrix = density->getMatrix(1); + EXPECT_NEAR(densityMatrix[0][0].real(), 0.0, 1e-12); + EXPECT_NEAR(densityMatrix[1][1].real(), 1.0, 1e-12); + dd->decRef(*density); + + EXPECT_TRUE(failed(buildFunctionality(func, *dd))); + bindings[func.getArgument(0)] = + IntegerAttr::get(IntegerType::get(context.get(), 64), 1); + EXPECT_TRUE(failed(buildFunctionality(func, *dd, bindings))); + EXPECT_TRUE( + failed(simulateDensity(func, makeZeroDensity(*dd, 1), *dd, bindings))); + EXPECT_TRUE(dd->getRootSet().empty()); + + bindings[func.getArgument(0)] = + FloatAttr::get(Float32Type::get(context.get()), 1.0); + EXPECT_TRUE(failed(buildFunctionality(func, *dd, bindings))); +} + +TEST_F(QCODDFunctionalityTest, BuildsThroughConcreteControlFlow) { + 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; }); + q = b.qcoIndexSwitch(1, q, ArrayRef{0, 1}, + SmallVector>{ + [&](Value arg) { return b.h(arg); }, + [&](Value arg) { return b.z(arg); }}, + [&](Value arg) { return arg; }); + q = b.scfFor(0, 2, 1, ValueRange{q.value}, + [&](Value /*index*/, ValueRange args) -> SmallVector { + return {b.h(args[0])}; + })[0]; + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + qc::QuantumComputation qc(1); + qc.x(0); + qc.z(0); + qc.h(0); + qc.h(0); + expectEqualToQc(mainFunc(*mod), qc); +} + +TEST_F(QCODDFunctionalityTest, StandardScfRegionsAndWhileCarryValues) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main() { + %q = qco.static 0 : !qco.qubit + %q1 = scf.execute_region -> !qco.qubit { + %out = qco.x %q : !qco.qubit -> !qco.qubit + scf.yield %out : !qco.qubit + } + %true = arith.constant true + %selector = scf.if %true -> index { + %one = arith.constant 1 : index + scf.yield %one : index + } else { + %zero = arith.constant 0 : index + scf.yield %zero : index + } + %apply_z = scf.index_switch %selector -> i1 + case 1 { + %yes = arith.constant true + scf.yield %yes : i1 + } + default { + %no = arith.constant false + scf.yield %no : i1 + } + %q2 = qco.if %apply_z args(%qarg = %q1) -> (!qco.qubit) { + %out = qco.z %qarg : !qco.qubit -> !qco.qubit + qco.yield %out : !qco.qubit + } else args(%qarg = %q1) { + qco.yield %qarg : !qco.qubit + } + %zero = arith.constant 0 : index + %result:2 = scf.while (%qarg = %q2, %i = %zero) + : (!qco.qubit, index) -> (!qco.qubit, index) { + %one = arith.constant 1 : index + %condition = arith.cmpi slt, %i, %one : index + scf.condition(%condition) %qarg, %i : !qco.qubit, index + } do { + ^bb0(%qarg: !qco.qubit, %i: index): + %out = qco.x %qarg : !qco.qubit -> !qco.qubit + %one = arith.constant 1 : index + %next = arith.addi %i, %one : index + scf.yield %out, %next : !qco.qubit, index + } + %false = arith.constant false + %final = scf.while (%qarg = %result#0) + : (!qco.qubit) -> !qco.qubit { + scf.condition(%false) %qarg : !qco.qubit + } do { + ^bb0(%qarg: !qco.qubit): + %unreachable = qco.h %qarg : !qco.qubit -> !qco.qubit + scf.yield %unreachable : !qco.qubit + } + qco.sink %final : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + qc::QuantumComputation qc(1); + qc.x(0); + qc.z(0); + qc.x(0); + expectEqualToQc(mainFunc(*mod), qc); +} + +TEST_F(QCODDFunctionalityTest, DeallocationRemovesSeparableQTensorWires) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%live: !qco.qubit) -> !qco.qubit { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c40 = arith.constant 40 : index + %tensor = qtensor.alloc(%c40) : tensor<40x!qco.qubit> + %remaining0, %q0 = qtensor.extract %tensor[%c0] + : tensor<40x!qco.qubit> + %remaining1, %q1 = qtensor.extract %remaining0[%c1] + : tensor<40x!qco.qubit> + %plus = qco.h %q0 : !qco.qubit -> !qco.qubit + %one = qco.x %q1 : !qco.qubit -> !qco.qubit + %restored0 = qtensor.insert %plus into %remaining1[%c0] + : tensor<40x!qco.qubit> + %restored1 = qtensor.insert %one into %restored0[%c1] + : tensor<40x!qco.qubit> + qtensor.dealloc %restored1 : tensor<40x!qco.qubit> + %result = qco.x %live : !qco.qubit -> !qco.qubit + return %result : !qco.qubit + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + // Expanding this 41-qubit statevector would require 2^41 amplitudes. + auto dd = std::make_unique(41); + std::mt19937_64 rng(3); + const auto histogram = sample(mainFunc(*mod), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 4}})); +} + +TEST_F(QCODDFunctionalityTest, RejectsEntangledQTensorDeallocation) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto tensor = b.qtensorAlloc(2); + Value q0; + Value q1; + std::tie(tensor, q0) = b.qtensorExtract(tensor, 0); + std::tie(tensor, q1) = b.qtensorExtract(tensor, 1); + q0 = b.h(q0); + std::tie(q0, q1) = b.cx(q0, q1); + tensor = b.qtensorInsert(q0, tensor, 0); + tensor = b.qtensorInsert(q1, tensor, 1); + b.qtensorDealloc(tensor); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(3); + EXPECT_TRUE(failed(sample(mainFunc(*mod), *dd, 1, rng))); +} + +TEST_F(QCODDFunctionalityTest, + SampleAllQubitsPreservesDeallocatedEntangledWires) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto tensor = b.qtensorAlloc(2); + Value q0; + Value q1; + std::tie(tensor, q0) = b.qtensorExtract(tensor, 0); + std::tie(tensor, q1) = b.qtensorExtract(tensor, 1); + q0 = b.h(q0); + std::tie(q0, q1) = b.cx(q0, q1); + tensor = b.qtensorInsert(q0, tensor, 0); + tensor = b.qtensorInsert(q1, tensor, 1); + b.qtensorDealloc(tensor); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(3); + const auto histogram = sampleAllQubits(mainFunc(*mod), *dd, 64, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(histogram->at("00") + histogram->at("11"), 64U); +} + +TEST_F(QCODDFunctionalityTest, + ClassicalOnlySamplingTreatsDeallocationAsLifetimeMarker) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto tensor = b.qtensorAlloc(2); + Value q0; + Value q1; + std::tie(tensor, q0) = b.qtensorExtract(tensor, 0); + std::tie(tensor, q1) = b.qtensorExtract(tensor, 1); + q0 = b.h(q0); + std::tie(q0, q1) = b.cx(q0, q1); + tensor = b.qtensorInsert(q0, tensor, 0); + tensor = b.qtensorInsert(q1, tensor, 1); + b.qtensorDealloc(tensor); + return b.allocClassicalBitRegister(1); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(3); + const auto histogram = sample(mainFunc(*mod), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"0", 4}})); +} + +TEST_F(QCODDFunctionalityTest, + CBitSamplingPreservesUnrelatedEntangledDeallocation) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto reg = + b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); + auto measured = b.allocQubit(); + std::tie(measured, std::ignore) = b.measure(measured, reg, 0); + + auto tensor = b.qtensorAlloc(2); + Value q0; + Value q1; + std::tie(tensor, q0) = b.qtensorExtract(tensor, 0); + std::tie(tensor, q1) = b.qtensorExtract(tensor, 1); + q0 = b.h(q0); + std::tie(q0, q1) = b.cx(q0, q1); + tensor = b.qtensorInsert(q0, tensor, 0); + tensor = b.qtensorInsert(q1, tensor, 1); + b.qtensorDealloc(tensor); + b.sink(measured); + return reg; + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(3); + std::mt19937_64 rng(31); + const auto histogram = sample(mainFunc(*mod), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"0", 4}})); +} + +TEST_F(QCODDFunctionalityTest, DeferredMeasurementsTrackDeallocatedLowerWires) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto low = b.qtensorAlloc(1); + auto high = b.qtensorAlloc(1); + auto reg = + b.allocClassicalBitRegister(1, {}, cbit::Initialization::Undefined); + Value measured; + std::tie(high, measured) = b.qtensorExtract(high, 0); + measured = b.x(measured); + std::tie(measured, std::ignore) = b.measure(measured, reg, 0); + high = b.qtensorInsert(measured, high, 0); + b.qtensorDealloc(low); + b.qtensorDealloc(high); + return reg; + }); + ASSERT_TRUE(mod); + + std::mt19937_64 rng(3); + auto vectorDD = std::make_unique(2); + const auto vectorHistogram = sample(mainFunc(*mod), *vectorDD, 4, rng); + ASSERT_TRUE(succeeded(vectorHistogram)); + EXPECT_EQ(*vectorHistogram, (std::map{{"1", 4}})); + + auto densityDD = std::make_unique(2); + const auto densityHistogram = sampleDensity( + mainFunc(*mod), makeZeroDensity(*densityDD, 0), *densityDD, 4, rng); + ASSERT_TRUE(succeeded(densityHistogram)); + EXPECT_EQ(*densityHistogram, (std::map{{"1", 4}})); +} + +TEST_F(QCODDFunctionalityTest, RejectsUnboundedStatevectorCapacity) { + auto nested = parseSourceString(R"mlir( + module { + func.func @main() { + %condition = arith.constant true + scf.if %condition { + %q = qco.alloc : !qco.qubit + qco.sink %q : !qco.qubit + } + return + } + } + )mlir", + context.get()); + auto called = parseSourceString(R"mlir( + module { + func.func @allocate() { + %q = qco.alloc : !qco.qubit + qco.sink %q : !qco.qubit + return + } + func.func @main() { + func.call @allocate() : () -> () + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(nested); + ASSERT_TRUE(called); + + EXPECT_TRUE(failed(getNumQubits(mainFunc(*nested)))); + EXPECT_TRUE(failed(getNumQubits(mainFunc(*called)))); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(37); + EXPECT_TRUE(failed(simulateStatevector(mainFunc(*nested), *dd, rng))); + EXPECT_TRUE(failed(simulateStatevector(mainFunc(*called), *dd, rng))); +} + +TEST_F(QCODDFunctionalityTest, + DensityAllocationAndPartialTraceProduceMixedState) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto live = b.h(b.allocQubit()); + auto tensor = b.qtensorAlloc(1); + Value allocated; + std::tie(tensor, allocated) = b.qtensorExtract(tensor, 0); + b.qtensorDealloc(tensor); + std::tie(live, allocated) = b.cx(live, allocated); + b.qtensorDealloc(b.qtensorFromElements({allocated})); + b.sink(live); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + auto result = simulateDensity(mainFunc(*mod), makeZeroDensity(*dd, 0), *dd); + ASSERT_TRUE(succeeded(result)); + const auto matrix = result->getMatrix(1); + EXPECT_NEAR(matrix[0][0].real(), 0.5, 1e-12); + EXPECT_NEAR(matrix[0][1].real(), 0.0, 1e-12); + EXPECT_NEAR(matrix[1][0].real(), 0.0, 1e-12); + EXPECT_NEAR(matrix[1][1].real(), 0.5, 1e-12); + dd->decRef(*result); + + std::mt19937_64 rng(7); + constexpr size_t shots = 1000; + const auto histogram = + sampleDensity(mainFunc(*mod), makeZeroDensity(*dd, 0), *dd, shots, rng); + ASSERT_TRUE(succeeded(histogram)); + ASSERT_EQ(histogram->size(), 2U); + EXPECT_EQ(histogram->at("0") + histogram->at("1"), shots); + EXPECT_NEAR(static_cast(histogram->at("0")), shots / 2.0, 100.0); + EXPECT_TRUE(dd->getRootSet().empty()); +} + +TEST_F(QCODDFunctionalityTest, DensityMeasurementFeedsClassicalControl) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.h(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); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(11); + const auto result = + simulateDensity(mainFunc(*mod), makeZeroDensity(*dd, 1), *dd, rng); + ASSERT_TRUE(succeeded(result)); + const auto matrix = result->getMatrix(1); + EXPECT_NEAR(matrix[0][0].real(), 1.0, 1e-12); + EXPECT_NEAR(matrix[1][1].real(), 0.0, 1e-12); + dd->decRef(*result); +} + +TEST_F(QCODDFunctionalityTest, DensityResetForcesZero) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q = b.reset(b.x(b.staticQubit(0))); + b.sink(q); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(3); + const auto result = + simulateDensity(mainFunc(*mod), makeZeroDensity(*dd, 1), *dd, rng); + ASSERT_TRUE(succeeded(result)); + const auto matrix = result->getMatrix(1); + EXPECT_NEAR(matrix[0][0].real(), 1.0, 1e-12); + EXPECT_NEAR(matrix[1][1].real(), 0.0, 1e-12); + dd->decRef(*result); +} + +TEST_F(QCODDFunctionalityTest, DynamicAllocationsAndQTensorBookkeeping) { + auto mod = buildModule([](QCOProgramBuilder& b) { + auto q0 = b.x(b.allocQubit()); + auto one = arith::ConstantIndexOp::create(b, 1).getResult(); + auto tensor = b.qtensorAlloc(one); + Value remaining; + Value q1; + std::tie(remaining, q1) = b.qtensorExtract(tensor, 0); + auto output = b.qtensorFromElements({q0, b.x(q1)}); + b.qtensorDealloc(remaining); + std::tie(output, q0) = b.qtensorExtract(output, 0); + std::tie(output, q1) = b.qtensorExtract(output, 1); + b.qtensorDealloc(output); + b.sink(q0); + b.sink(q1); + return b.intConstant(0); + }); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(3); + const auto histogram = sample(mainFunc(*mod), *dd, 8, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"11", 8}})); + + auto smallDd = std::make_unique(1); + EXPECT_TRUE(failed(sample(mainFunc(*mod), *smallDd, 1, rng))); + + auto invalidIndex = parseSourceString(R"mlir( + module { + func.func @main() { + %one = arith.constant 1 : index + %tensor = qtensor.alloc(%one) : tensor + %remaining, %q = qtensor.extract %tensor[%one] + : tensor + qco.sink %q : !qco.qubit + qtensor.dealloc %remaining : tensor + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(invalidIndex); + auto oneQubitDd = std::make_unique(1); + EXPECT_TRUE(failed(sample(mainFunc(*invalidIndex), *oneQubitDd, 1, rng))); +} + +TEST_F(QCODDFunctionalityTest, DynamicQTensorArgumentUsesBoundExtent) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%arg0: tensor) + -> tensor { + %one = arith.constant 1 : index + %remaining, %q = qtensor.extract %arg0[%one] + : tensor + %q1 = qco.x %q : !qco.qubit -> !qco.qubit + %result = qtensor.insert %q1 into %remaining[%one] + : tensor + return %result : tensor + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + auto func = mainFunc(*mod); + DDBindings bindings; + bindings[func.getArgument(0)] = + IntegerAttr::get(IndexType::get(context.get()), 2); + + auto dd = std::make_unique(2); + std::mt19937_64 rng(7); + const auto histogram = sample(func, *dd, 4, rng, bindings); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"10", 4}})); + + EXPECT_TRUE(failed(buildFunctionality(func, *dd))); + bindings[func.getArgument(0)] = + IntegerAttr::get(IndexType::get(context.get()), -1); + EXPECT_TRUE(failed(buildFunctionality(func, *dd, bindings))); +} + +TEST_F(QCODDFunctionalityTest, RejectsQTensorBeyondQubitRange) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%qubits: tensor<65537x!qco.qubit>) { + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + const auto func = mainFunc(*mod); + EXPECT_TRUE(failed(getNumQubits(func))); + EXPECT_TRUE(failed(buildFunctionality(func, *dd))); + + std::mt19937_64 rng(7); + EXPECT_TRUE(failed(simulateStatevector(func, *dd, rng))); +} + +TEST_F(QCODDFunctionalityTest, QTensorFlowsThroughLoopAndCall) { + auto mod = parseSourceString(R"mlir( + module { + func.func @flip(%arg: tensor<1x!qco.qubit>) + -> tensor<1x!qco.qubit> { + %zero = arith.constant 0 : index + %remaining, %q = qtensor.extract %arg[%zero] + : tensor<1x!qco.qubit> + %q1 = qco.x %q : !qco.qubit -> !qco.qubit + %result = qtensor.insert %q1 into %remaining[%zero] + : tensor<1x!qco.qubit> + return %result : tensor<1x!qco.qubit> + } + func.func @main() { + %zero = arith.constant 0 : index + %one = arith.constant 1 : index + %tensor = qtensor.alloc(%one) : tensor<1x!qco.qubit> + %result = scf.for %i = %zero to %one step %one + iter_args(%arg = %tensor) -> tensor<1x!qco.qubit> { + %next = func.call @flip(%arg) + : (tensor<1x!qco.qubit>) -> tensor<1x!qco.qubit> + scf.yield %next : tensor<1x!qco.qubit> + } + %remaining, %q = qtensor.extract %result[%zero] + : tensor<1x!qco.qubit> + qtensor.dealloc %remaining : tensor<1x!qco.qubit> + qco.sink %q : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto dd = std::make_unique(1); + std::mt19937_64 rng(13); + const auto histogram = sample(mainFunc(*mod), *dd, 4, rng); + ASSERT_TRUE(succeeded(histogram)); + EXPECT_EQ(*histogram, (std::map{{"1", 4}})); +} + +TEST_F(QCODDFunctionalityTest, WiderMemRefCallsShareStorage) { + auto mod = parseSourceString(R"mlir( + module { + func.func @set(%reg: memref, %value: i16) { + %zero = arith.constant 0 : index + memref.store %value, %reg[%zero] : memref + return + } + func.func @main() { + %one = arith.constant 1 : index + %reg = memref.alloc(%one) : memref + %three = arith.constant 3 : i16 + %four = arith.constant 4 : i16 + %seven = arith.addi %three, %four : i16 + %two = arith.constant 2 : i16 + %fourteen = arith.muli %seven, %two : i16 + %quotient = arith.divsi %fourteen, %two : i16 + %remainder = arith.remui %quotient, %two : i16 + %shifted = arith.shli %remainder, %two : i16 + %restored = arith.shrui %shifted, %two : i16 + %wide = arith.extui %restored : i16 to i32 + %narrow = arith.trunci %wide : i32 to i16 + %as_float = arith.sitofp %narrow : i16 to f64 + %back = arith.fptosi %as_float : f64 to i16 + func.call @set(%reg, %quotient) : (memref, i16) -> () + %zero = arith.constant 0 : index + %stored = memref.load %reg[%zero] : memref + %expected = arith.constant 7 : i16 + %integer_ok = arith.cmpi eq, %stored, %expected : i16 + %casts_ok = arith.cmpi eq, %back, %remainder : i16 + %one_float = arith.constant 1.0 : f64 + %two_float = arith.addf %one_float, %one_float : f64 + %four_float = arith.addf %two_float, %two_float : f64 + %half = arith.divf %four_float, %two_float : f64 + %float_remainder = arith.remf %half, %one_float : f64 + %zero_float = arith.constant 0.0 : f64 + %float_ok = arith.cmpf oeq, %float_remainder, %zero_float : f64 + %integer_and_casts = arith.andi %integer_ok, %casts_ok : i1 + %condition = arith.andi %integer_and_casts, %float_ok : i1 + %q = qco.static 0 : !qco.qubit + %q1 = qco.if %condition args(%qin = %q) -> (!qco.qubit) { + %out = qco.x %qin : !qco.qubit -> !qco.qubit + qco.yield %out : !qco.qubit + } else args(%qin = %q) { + qco.yield %qin : !qco.qubit + } + memref.dealloc %reg : memref + qco.sink %q1 : !qco.qubit + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + expectSimulatesFromZero(mainFunc(*mod), true); +} + +TEST_F(QCODDFunctionalityTest, + SupportsAdditionalClassicalOperationsAndBindings) { + auto mod = parseSourceString(R"mlir( + module { + func.func @main(%idx: index, %word: i16, %flag: i1) { + %zero = arith.constant 0 : i8 + %one = arith.constant 1 : i8 + %two = arith.constant 2 : i8 + %four = arith.constant 4 : i8 + %negative = arith.constant -5 : i8 + %sle = arith.cmpi sle, %one, %two : i8 + %sgt = arith.cmpi sgt, %two, %one : i8 + %sge = arith.cmpi sge, %two, %two : i8 + %ult = arith.cmpi ult, %one, %two : i8 + %ule = arith.cmpi ule, %two, %two : i8 + %ugt = arith.cmpi ugt, %two, %one : i8 + %uge = arith.cmpi uge, %two, %two : i8 + %quotient = arith.divui %four, %two : i8 + %remainder = arith.remsi %negative, %two : i8 + %shifted = arith.shrsi %negative, %one : i8 + %extended = arith.extsi %negative : i8 to i16 + %as_index = arith.index_cast %word : i16 to index + %selected = arith.select %flag, %one, %zero : i8 + %one_float = arith.constant 1.0 : f64 + %two_float = arith.constant 2.0 : f64 + %difference = arith.subf %two_float, %one_float : f64 + %product = arith.mulf %difference, %two_float : f64 + %negated = arith.negf %product : f64 + %zero_index = arith.constant 0 : index + %indices = memref.alloc() : memref<1xindex> + %loaded_index = memref.load %indices[%zero_index] : memref<1xindex> + memref.dealloc %indices : memref<1xindex> + %floats = memref.alloc() : memref<1xf64> + %loaded_float = memref.load %floats[%zero_index] : memref<1xf64> + memref.dealloc %floats : memref<1xf64> + %false = arith.constant false + scf.if %false { + } + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + + auto func = mainFunc(*mod); + DDBindings bindings; + bindings[func.getArgument(0)] = + IntegerAttr::get(IndexType::get(context.get()), 3); + bindings[func.getArgument(1)] = + IntegerAttr::get(IntegerType::get(context.get(), 16), -2); + bindings[func.getArgument(2)] = + IntegerAttr::get(IntegerType::get(context.get(), 1), 1); + auto dd = std::make_unique(0); + const auto output = simulate(func, dd::VectorDD::one(), *dd, rng, bindings); + ASSERT_TRUE(succeeded(output)); + EXPECT_TRUE(output->isTerminal()); + dd->decRef(*output); +} + +TEST_F(QCODDFunctionalityTest, RejectsClassicalRuntimeErrors) { + for (const StringRef source : { + R"mlir(module { + func.func @main(%unbound: i16) { + %true = arith.constant true + %zero = arith.constant 0 : i16 + %selected = arith.select %true, %unbound, %zero : i16 + return + } + })mlir", + R"mlir(module { + func.func @main(%reg: memref<1xi16>) { + %zero = arith.constant 0 : index + %value = memref.load %reg[%zero] : memref<1xi16> + return + } + })mlir", + R"mlir(module { + func.func @main(%index: index) { + %reg = memref.alloc() : memref<1xi16> + %value = memref.load %reg[%index] : memref<1xi16> + return + } + })mlir", + R"mlir(module { + func.func @main(%value: i16) { + %zero = arith.constant 0 : index + %reg = memref.alloc() : memref<1xi16> + memref.store %value, %reg[%zero] : memref<1xi16> + return + } + })mlir", + R"mlir(module { + func.func @main() { + %negative = arith.constant -1 : index + %reg = memref.alloc(%negative) : memref + return + } + })mlir", + R"mlir(module { + func.func @main() { + %zero = arith.constant 0 : i8 + %one = arith.constant 1 : i8 + %invalid = arith.divui %one, %zero : i8 + return + } + })mlir", + R"mlir(module { + func.func @main() { + %huge = arith.constant 1.0e+300 : f64 + %invalid = arith.fptosi %huge : f64 to i8 + return + } + })mlir", + R"mlir(module { + func.func @main(%rhs: i8) { + %one = arith.constant 1 : i8 + %invalid = arith.divui %one, %rhs : i8 + return + } + })mlir", + R"mlir(module { + func.func @main(%lhs: i8) { + %one = arith.constant 1 : i8 + %invalid = arith.divui %lhs, %one : i8 + return + } + })mlir", + R"mlir(module { + func.func @main(%amount: i8) { + %one = arith.constant 1 : i8 + %invalid = arith.shli %one, %amount : i8 + return + } + })mlir", + R"mlir(module { + func.func @main(%lhs: f64) { + %zero = arith.constant 0.0 : f64 + %invalid = arith.cmpf oeq, %lhs, %zero : f64 + return + } + })mlir", + R"mlir(module { + func.func @main(%value: i8) { + %invalid = arith.sitofp %value : i8 to f64 + return + } + })mlir", + R"mlir(module { + func.func @main(%value: f64) { + %invalid = arith.fptosi %value : f64 to i8 + return + } + })mlir", + R"mlir(module { + func.func @consume(%reg: memref<1xi16>) { + return + } + func.func @main(%reg: memref<1xi16>) { + func.call @consume(%reg) : (memref<1xi16>) -> () + return + } + })mlir", + R"mlir(module { + func.func @main(%condition: i1) { + scf.if %condition { + } + return + } + })mlir", + R"mlir(module { + func.func @main(%selector: index) { + scf.index_switch %selector + default { + } + return + } + })mlir", + R"mlir(module { + func.func @main(%size: index) { + %tensor = qtensor.alloc(%size) : tensor + qtensor.dealloc %tensor : tensor + return + } + })mlir"}) { + expectMlirSimulationFails(0, source); + } + + auto mod = parseSourceString(R"mlir( + module { + func.func @main() { + %zero = arith.constant 0 : index + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(mod); + auto func = mainFunc(*mod); + auto constant = *func.getBody().front().getOps().begin(); + DDBindings bindings; + bindings[constant.getResult()] = + IntegerAttr::get(IndexType::get(context.get()), 0); + auto dd = std::make_unique(0); + EXPECT_TRUE(failed(simulate(func, dd::VectorDD::one(), *dd, rng, bindings))); + EXPECT_TRUE(dd->getRootSet().empty()); +} + +TEST_F(QCODDFunctionalityTest, BuildFunctionalityRestrictsRuntimeAllocations) { + auto topLevel = parseSourceString(R"mlir( + module { + func.func @main() { + %q = qco.alloc : !qco.qubit + %out = qco.x %q : !qco.qubit -> !qco.qubit + qco.sink %out : !qco.qubit + return + } + } + )mlir", + context.get()); + auto nested = parseSourceString(R"mlir( + module { + func.func @main() { + %true = arith.constant true + scf.if %true { + %q = qco.alloc : !qco.qubit + qco.sink %q : !qco.qubit + } + return + } + } + )mlir", + context.get()); + auto tensor = parseSourceString(R"mlir( + module { + func.func @main() { + %one = arith.constant 1 : index + %tensor = qtensor.alloc(%one) : tensor + qtensor.dealloc %tensor : tensor + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(topLevel); + ASSERT_TRUE(nested); + ASSERT_TRUE(tensor); + + auto dd = std::make_unique(1); + const auto functionality = buildFunctionality(mainFunc(*topLevel), *dd); + ASSERT_TRUE(succeeded(functionality)); + dd->decRef(*functionality); + EXPECT_TRUE(failed(buildFunctionality(mainFunc(*nested), *dd))); + EXPECT_TRUE(failed(buildFunctionality(mainFunc(*tensor), *dd))); +} + } // namespace diff --git a/python/mqt/core/mlir.pyi b/python/mqt/core/mlir.pyi index 9b37fbee84..d70adbce18 100644 --- a/python/mqt/core/mlir.pyi +++ b/python/mqt/core/mlir.pyi @@ -10,7 +10,7 @@ import enum import os -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Literal, Unpack, overload import qiskit.circuit @@ -472,11 +472,14 @@ class QCOProgram(Program): Set ``copy=True`` to preserve it. """ - def build_functionality(self, dd_package: mqt.core.dd.DDPackage) -> mqt.core.dd.MatrixDD: + def build_functionality( + self, dd_package: mqt.core.dd.DDPackage, *, bindings: Mapping[int, bool | int | float] = {} + ) -> mqt.core.dd.MatrixDD: """Build a matrix DD for a static unitary QCO program. Args: dd_package: DD package with enough qubits for the program. + bindings: Concrete entry-argument values keyed by zero-based argument index. Returns: Matrix DD of the program functionality. @@ -486,7 +489,12 @@ class QCOProgram(Program): """ def simulate( - self, initial_state: mqt.core.dd.VectorDD, dd_package: mqt.core.dd.DDPackage, seed: int = 0 + self, + initial_state: mqt.core.dd.VectorDD, + dd_package: mqt.core.dd.DDPackage, + seed: int = 0, + *, + bindings: Mapping[int, bool | int | float] = {}, ) -> mqt.core.dd.VectorDD: """Simulate a QCO program on a DD state. @@ -497,6 +505,7 @@ class QCOProgram(Program): dd_package: DD package with enough qubits for the program. seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other value produces reproducible measurement and reset results. + bindings: Concrete entry-argument values keyed by zero-based argument index. Returns: Output state DD. @@ -506,7 +515,43 @@ class QCOProgram(Program): has too few qubits, or the program is unsupported for simulation. """ - def sample(self, dd_package: mqt.core.dd.DDPackage, shots: int = 1024, seed: int = 0) -> dict[str, int]: + def simulate_density( + self, + initial_state: mqt.core.dd.MatrixDD, + dd_package: mqt.core.dd.DDPackage, + seed: int = 0, + *, + bindings: Mapping[int, bool | int | float] = {}, + ) -> mqt.core.dd.MatrixDD: + """Simulate a QCO program on a density-matrix DD. + + Args: + initial_state: Input density-matrix DD with a live reference in + ``dd_package``. It represents exactly the program's inferred initial + quantum register; skipped DD levels denote identity factors within + that register. A valid input reference is consumed. + dd_package: DD package with enough qubits for the program. + seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other + value produces reproducible measurement and reset results. + bindings: Concrete entry-argument values keyed by zero-based argument index. + + Returns: + Output density-matrix DD. + + Raises: + ValueError: When ``initial_state`` has no live reference in ``dd_package`` + or the program is unsupported for simulation. + """ + + def sample( + self, + dd_package: mqt.core.dd.DDPackage, + shots: int = 1024, + seed: int = 0, + *, + initial_state: mqt.core.dd.VectorDD | None = None, + bindings: Mapping[int, bool | int | float] = {}, + ) -> dict[str, int]: """Sample the declared outputs of a QCO program. Args: @@ -514,13 +559,48 @@ class QCOProgram(Program): shots: Number of shots (default 1024). seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other value produces reproducible results. + initial_state: Optional input state with a live reference in ``dd_package``. + A valid input reference is consumed. + bindings: Concrete entry-argument values keyed by zero-based argument index. Returns: Histogram of returned CBit registers in return order, each MSB first. If no CBit result exists, final ``measureAll`` bitstrings instead. Raises: - ValueError: When the program is unsupported for sampling. + ValueError: When ``initial_state`` has no live reference in ``dd_package`` + or the program is unsupported for sampling. + """ + + def sample_density( + self, + initial_state: mqt.core.dd.MatrixDD, + dd_package: mqt.core.dd.DDPackage, + shots: int = 1024, + seed: int = 0, + *, + bindings: Mapping[int, bool | int | float] = {}, + ) -> dict[str, int]: + """Sample the declared outputs of a QCO program from a density-matrix DD. + + Args: + initial_state: Input density-matrix DD with a live reference in + ``dd_package``. It represents exactly the program's inferred initial + quantum register; skipped DD levels denote identity factors within + that register. A valid input reference is consumed. + dd_package: DD package with enough qubits for the program. + shots: Number of shots (default 1024). + seed: RNG seed. ``0`` (default) selects nondeterministic seeding. Any other + value produces reproducible results. + bindings: Concrete entry-argument values keyed by zero-based argument index. + + Returns: + Histogram of returned CBit registers in return order, each MSB first. If + no CBit result exists, final ``measureAll`` bitstrings instead. + + Raises: + ValueError: When ``initial_state`` has no live reference in ``dd_package`` + or the program is unsupported for sampling. """ class JeffProgram(Program): @@ -593,6 +673,28 @@ class QIRProgram(Program): def write_bitcode(self, path: str | os.PathLike) -> None: """Write this program as LLVM bitcode.""" +def make_density_matrix( + state: mqt.core.dd.VectorDD, num_qubits: int, dd_package: mqt.core.dd.DDPackage +) -> mqt.core.dd.MatrixDD: + """Construct ``|psi>& eliminate) { - auto r = trace(a, eliminate, eliminate.size()); + if (!a.isTerminal() && static_cast(a.p->v) >= eliminate.size()) { + throw std::invalid_argument( + "Elimination mask does not cover the matrix decision diagram."); + } + + std::vector keptBefore(eliminate.size() + 1U, 0U); + for (std::size_t q = 0; q < eliminate.size(); ++q) { + keptBefore[q + 1U] = keptBefore[q] + (eliminate[q] ? 0U : 1U); + } + + auto r = trace(a, eliminate, keptBefore); return {.p = r.p, .w = cn.lookup(r.w)}; } ComplexValue Package::trace(const mEdge& a, const std::size_t numQubits) { if (a.isIdentity()) { return static_cast(a.w); } + if (!a.isTerminal() && static_cast(a.p->v) >= numQubits) { + throw std::invalid_argument( + "Qubit count does not cover the matrix decision diagram."); + } + const auto eliminate = std::vector(numQubits, true); - return trace(a, eliminate, numQubits).w; + const auto keptBefore = std::vector(numQubits + 1U, 0U); + return trace(a, eliminate, keptBefore).w; } bool Package::isCloseToIdentity(const mEdge& m, const fp tol, const std::vector& garbage, @@ -910,41 +926,43 @@ bool Package::isCloseToIdentity(const mEdge& m, const fp tol, return isCloseToIdentityRecursive(m, visited, tol, garbage, checkCloseToOne); } mCachedEdge Package::trace(const mEdge& a, const std::vector& eliminate, - std::size_t level, std::size_t alreadyEliminated) { + const std::vector& keptBefore) { const auto aWeight = static_cast(a.w); if (aWeight.approximatelyZero()) { return mCachedEdge::zero(); } - // If `a` is the identity matrix or there is nothing left to eliminate, - // then simply return `a` - if (a.isIdentity() || - std::none_of(eliminate.begin(), - eliminate.begin() + - static_cast::difference_type>(level), - [](bool v) { return v; })) { + if (a.isTerminal()) { + return mCachedEdge{a.p, aWeight}; + } + + const auto v = static_cast(a.p->v); + assert(v < eliminate.size()); + + // Eliminated identity levels above this node do not affect its value or + // numbering. If every logical level at and below it is kept, return it. + if (keptBefore[v + 1U] == v + 1U) { return mCachedEdge{a.p, aWeight}; } - const auto v = a.p->v; + const auto lowerKept = keptBefore[v]; if (eliminate[v]) { // Lookup nodes marked for elimination in the compute table if all // lower-level qubits are eliminated as well: if the trace has already // been computed, return the result - const auto eliminateAll = - std::all_of(eliminate.begin(), - eliminate.begin() + - static_cast::difference_type>(level), - [](bool e) { return e; }); - if (eliminateAll) { + const auto fullSubtrace = keptBefore[v + 1U] == 0U; + if (fullSubtrace) { if (const auto* r = getTraceComputeTable().lookup(a.p); r != nullptr) { return {r->p, r->w * aWeight}; } } - const auto elims = alreadyEliminated + 1; - auto r = add2(trace(a.p->e[0], eliminate, level - 1, elims), - trace(a.p->e[3], eliminate, level - 1, elims), v - 1); + auto low = trace(a.p->e[0], eliminate, keptBefore); + auto high = trace(a.p->e[3], eliminate, keptBefore); + assert(lowerKept != 0U || (low.isTerminal() && high.isTerminal())); + const auto addLevel = + lowerKept == 0U ? Qubit{0} : static_cast(lowerKept - 1U); + auto r = add2(low, high, addLevel); // The resulting weight is continuously normalized to the range [0,1] for // matrix nodes @@ -952,7 +970,7 @@ mCachedEdge Package::trace(const mEdge& a, const std::vector& eliminate, // Insert result into compute table if all lower-level qubits are // eliminated as well - if (eliminateAll) { + if (fullSubtrace) { getTraceComputeTable().insert(a.p, r); } r.w = r.w * aWeight; @@ -960,17 +978,12 @@ mCachedEdge Package::trace(const mEdge& a, const std::vector& eliminate, } std::array edge{}; - std::ranges::transform(std::as_const(a.p->e), edge.begin(), - [this, &eliminate, &alreadyEliminated, - &level](const mEdge& e) -> mCachedEdge { - return trace(e, eliminate, level - 1, - alreadyEliminated); - }); - const auto adjustedV = static_cast( - static_cast(a.p->v) - - (static_cast(std::ranges::count(eliminate, true)) - - alreadyEliminated)); - auto r = makeDDNode(adjustedV, edge); + std::ranges::transform( + std::as_const(a.p->e), edge.begin(), + [this, &eliminate, &keptBefore](const mEdge& e) -> mCachedEdge { + return trace(e, eliminate, keptBefore); + }); + auto r = makeDDNode(static_cast(lowerKept), edge); r.w = r.w * aWeight; return r; } diff --git a/test/dd/test_package.cpp b/test/dd/test_package.cpp index 338a95e7b2..258468b51f 100644 --- a/test/dd/test_package.cpp +++ b/test/dd/test_package.cpp @@ -346,6 +346,45 @@ TEST(DDPackageTest, PartialIdentityTrace) { EXPECT_EQ(RealNumber::val(mul.w.r), 1.); } +TEST(DDPackageTest, PartialTraceSkippedIdentityAboveRoot) { + auto dd = std::make_unique(2); + const auto input = getDD(qc::StandardOperation(0, qc::X), *dd); + + ASSERT_FALSE(input.isTerminal()); + ASSERT_EQ(input.p->v, 0); + + const auto reduced = dd->partialTrace(input, {false, true}); + + EXPECT_EQ(reduced, input); +} + +TEST(DDPackageTest, PartialTraceSkippedIdentityInsideDiagram) { + auto dd = std::make_unique(4); + const auto input = getDD(qc::StandardOperation(3_pc, 1, qc::X), *dd); + const auto expected = getDD(qc::StandardOperation(2_pc, 1, qc::X), *dd); + + ASSERT_FALSE(input.isTerminal()); + ASSERT_EQ(input.p->v, 3); + ASSERT_FALSE(input.p->e[3].isTerminal()); + ASSERT_EQ(input.p->e[3].p->v, 1); + + const auto reduced = dd->partialTrace(input, {false, false, true, false}); + + EXPECT_EQ(reduced, expected); +} + +TEST(DDPackageTest, PartialTraceRenumbersDiagonalBlockSum) { + auto dd = std::make_unique(3); + const auto input = getDD(qc::StandardOperation(2_pc, 1, qc::X), *dd); + const auto x = getDD(qc::StandardOperation(0, qc::X), *dd); + auto expected = dd->add(Package::makeIdent(), x); + expected.w = dd->cn.lookup(static_cast(expected.w) / 2.0); + + const auto reduced = dd->partialTrace(input, {true, false, true}); + + EXPECT_EQ(reduced, expected); +} + TEST(DDPackageTest, PartialSWapMatTrace) { auto dd = std::make_unique(2); auto swapGate = diff --git a/test/python/test_qco_dd.py b/test/python/test_qco_dd.py index 2aa732115e..eae8e41571 100644 --- a/test/python/test_qco_dd.py +++ b/test/python/test_qco_dd.py @@ -14,7 +14,7 @@ import pytest from mqt.core.dd import DDPackage -from mqt.core.mlir import OutputFormat, QCOProgram, compile_program +from mqt.core.mlir import OutputFormat, QCOProgram, compile_program, make_density_matrix def _x_program() -> QCOProgram: @@ -60,6 +60,34 @@ def _measure_program() -> QCOProgram: """) +def _bound_qtensor_program() -> QCOProgram: + """Construct a bound RX program over a dynamic QTensor argument. + + Returns: + The constructed QCO program. + """ + return QCOProgram.from_mlir_str(""" +module { + func.func @main(%apply: i1, %theta: f64, + %input: tensor) -> tensor + attributes {mqt.entry_point} { + %c1 = arith.constant 1 : index + %remaining, %q = qtensor.extract %input[%c1] + : tensor + %q1 = qco.if %apply args(%q_in = %q) -> (!qco.qubit) { + %rotated = qco.rx(%theta) %q_in : !qco.qubit -> !qco.qubit + qco.yield %rotated : !qco.qubit + } else args(%q_in = %q) { + qco.yield %q_in : !qco.qubit + } + %result = qtensor.insert %q1 into %remaining[%c1] + : tensor + return %result : tensor + } +} +""") + + def test_unitary_x_build_simulate_and_sample() -> None: """X on |0>: unitary matrix, simulation to |1>, deterministic sampling.""" program = _x_program() @@ -97,8 +125,8 @@ def test_simulate_measure_uses_default_or_explicit_seed() -> None: package.dec_ref_vec(expected) -def test_simulate_rejects_state_from_different_package() -> None: - """Simulation rejects a state owned by a different DD package.""" +def test_dd_apis_reject_state_from_different_package_without_consuming_it() -> None: + """DD APIs reject a foreign state without consuming its live reference.""" program = _x_program() source_package = DDPackage(1) target_package = DDPackage(1) @@ -110,7 +138,13 @@ def test_simulate_rejects_state_from_different_package() -> None: with pytest.raises(ValueError, match=r"live reference in dd_package"): program.simulate(zero, target_package, seed=7) - source_package.dec_ref_vec(zero) + with pytest.raises(ValueError, match=r"live reference in dd_package"): + make_density_matrix(zero, 1, target_package) + with pytest.raises(ValueError, match=r"live reference in dd_package"): + program.sample(target_package, initial_state=zero) + + out = program.simulate(zero, source_package) + source_package.dec_ref_vec(out) target_package.dec_ref_vec(target_zero) @@ -175,3 +209,143 @@ def test_compiler_to_sampler_outputs(source: str, num_qubits: int, expected: set assert set(counts) == expected assert sum(counts.values()) == shots + + # Compiled OpenQASM programs allocate their qubits, so their initial + # quantum register is empty even though the package needs full capacity. + zero = package.zero_state(0) + density = make_density_matrix(zero, 0, package) + package.dec_ref_vec(zero) + density_counts = program.sample_density(density, package, shots=shots, seed=17) + + assert set(density_counts) == expected + assert sum(density_counts.values()) == shots + + +def test_symbolic_bindings_across_dd_apis() -> None: + """All DD APIs accept concrete scalar and dynamic QTensor bindings.""" + program = _bound_qtensor_program() + package = DDPackage(2) + bindings = {0: True, 1: float(np.pi), 2: 2} + + matrix = program.build_functionality(package, bindings=bindings) + package.dec_ref_mat(matrix) + + zero = package.zero_state(2) + out = program.simulate(zero, package, bindings=bindings) + expected = package.computational_basis_state(2, [False, True]) + assert np.allclose(np.abs(out.get_vector()), np.abs(expected.get_vector())) + package.dec_ref_vec(out) + + zero = package.zero_state(2) + density = make_density_matrix(zero, 2, package) + package.dec_ref_vec(zero) + density_out = program.simulate_density(density, package, bindings=bindings) + expected_density = make_density_matrix(expected, 2, package) + assert np.allclose(density_out.get_matrix(2), expected_density.get_matrix(2)) + package.dec_ref_mat(density_out) + package.dec_ref_mat(expected_density) + package.dec_ref_vec(expected) + + assert program.sample(package, shots=8, seed=4, bindings=bindings) == {"10": 8} + zero = package.zero_state(2) + density = make_density_matrix(zero, 2, package) + package.dec_ref_vec(zero) + assert program.sample_density(density, package, shots=8, seed=4, bindings=bindings) == {"10": 8} + + +@pytest.mark.parametrize( + "bindings", + [ + pytest.param({3: 0}, id="argument-index"), + pytest.param({0: 1}, id="boolean-type"), + pytest.param({1: 1}, id="float-type"), + pytest.param({2: -1}, id="negative-extent"), + ], +) +def test_python_dd_bindings_reject_invalid_values(bindings: dict[int, bool | int | float]) -> None: + """Binding indices and values must match entry argument types.""" + program = _bound_qtensor_program() + package = DDPackage(2) + with pytest.raises(ValueError, match=r"out of range|does not match"): + program.build_functionality(package, bindings=bindings) + + +def test_sample_from_supplied_initial_state() -> None: + """Sampling consumes a valid caller-supplied input state.""" + program = _x_program() + package = DDPackage(1) + + one = package.computational_basis_state(1, [True]) + assert program.sample(package, shots=8, seed=6, initial_state=one) == {"0": 8} + + +def test_density_simulation_and_sampling() -> None: + """Construct, simulate, and sample pure-state density matrices.""" + package = DDPackage(1) + + zero = package.zero_state(1) + density = make_density_matrix(zero, 1, package) + package.dec_ref_vec(zero) + out = _x_program().simulate_density(density, package) + assert np.allclose(out.get_matrix(1), np.asarray([[0.0, 0.0], [0.0, 1.0]])) + package.dec_ref_mat(out) + + zero = package.zero_state(1) + density = make_density_matrix(zero, 1, package) + package.dec_ref_vec(zero) + out = _measure_program().simulate_density(density, package) + assert np.allclose(out.get_matrix(1), np.asarray([[1.0, 0.0], [0.0, 0.0]])) + package.dec_ref_mat(out) + + zero = package.zero_state(1) + density = make_density_matrix(zero, 1, package) + package.dec_ref_vec(zero) + assert _x_program().sample_density(density, package, shots=16, seed=8) == {"1": 16} + + +def test_make_density_matrix_rejects_invalid_input() -> None: + """Density construction requires a live state and sufficient capacity.""" + package = DDPackage(1) + zero = package.zero_state(1) + with pytest.raises(ValueError, match=r"does not cover"): + make_density_matrix(zero, 0, package) + with pytest.raises(ValueError, match=r"exceeds the capacity"): + make_density_matrix(zero, 2, package) + density = make_density_matrix(zero, 1, package) + package.dec_ref_mat(density) + package.dec_ref_vec(zero) + + unrooted = package.zero_state(1) + package.dec_ref_vec(unrooted) + with pytest.raises(ValueError, match=r"live reference in dd_package"): + make_density_matrix(unrooted, 1, package) + + +def test_density_apis_reject_foreign_or_unrooted_state_without_consuming_it() -> None: + """Density execution checks matrix ownership before consuming a reference.""" + program = _x_program() + source_package = DDPackage(1) + target_package = DDPackage(1) + zero = source_package.zero_state(1) + density = make_density_matrix(zero, 1, source_package) + source_package.dec_ref_vec(zero) + + for call in ( + lambda: program.simulate_density(density, target_package), + lambda: program.simulate_density(density, target_package, seed=7), + lambda: program.sample_density(density, target_package, shots=1, seed=7), + ): + with pytest.raises(ValueError, match=r"live reference in dd_package"): + call() + + out = program.simulate_density(density, source_package) + source_package.dec_ref_mat(out) + + zero = source_package.zero_state(1) + unrooted = make_density_matrix(zero, 1, source_package) + source_package.dec_ref_vec(zero) + source_package.dec_ref_mat(unrooted) + with pytest.raises(ValueError, match=r"live reference in dd_package"): + program.simulate_density(unrooted, source_package) + with pytest.raises(ValueError, match=r"live reference in dd_package"): + program.sample_density(unrooted, source_package, shots=1)