Add a state preparation pass before synthesis

include/cudaq/Optimizer/Transforms/Passes.h

@@ -40,6 +40,8 @@ std::unique_ptr<mlir::Pass> createLowerToCFGPass();
 std::unique_ptr<mlir::Pass> createObserveAnsatzPass(std::vector<bool> &);
 std::unique_ptr<mlir::Pass> createQuakeAddMetadata();
 std::unique_ptr<mlir::Pass> createQuakeAddDeallocs();
+std::unique_ptr<mlir::Pass> createStatePreparation();
+std::unique_ptr<mlir::Pass> createStatePreparation(std::string_view, void *);
 std::unique_ptr<mlir::Pass> createQuakeSynthesizer();
 std::unique_ptr<mlir::Pass> createQuakeSynthesizer(std::string_view, void *);
 std::unique_ptr<mlir::Pass> createRaiseToAffinePass();

include/cudaq/Optimizer/Transforms/Passes.td

@@ -512,6 +512,17 @@ def PruneCtrlRelations : Pass<"pruned-ctrl-form", "mlir::func::FuncOp"> {
  }];
 }
 
+def PrepareState : Pass<"state-prep", "mlir::ModuleOp"> {
+ let summary =
+ "Convert state vector data into gates";
+ let description = [{
+ Convert quake representation that includes qubit initialization
+ from data into qubit initialization using gates.
+ }];
+
+ let constructor = "cudaq::opt::createStatePreparation()";
+}
+
 def QuakeSynthesize : Pass<"quake-synth", "mlir::ModuleOp"> {
  let summary =
  "Synthesize concrete quantum program from Quake code plus runtime values.";

lib/Optimizer/Transforms/CMakeLists.txt

@@ -42,6 +42,8 @@ add_cudaq_library(OptTransforms
  QuakeSynthesizer.cpp
  RefToVeqAlloc.cpp
  RegToMem.cpp
+ StateDecomposer.cpp
+ StatePreparation.cpp
  PySynthCallableBlockArgs.cpp
 
  DEPENDS

lib/Optimizer/Transforms/GenKernelExecution.cpp

@@ -1220,7 +1220,6 @@ class GenerateKernelExecution
  }
  continue;
  }
-
  stVal = builder.create<cudaq::cc::InsertValueOp>(loc, stVal.getType(),
  stVal, arg, idx);
  }

lib/Optimizer/Transforms/StateDecomposer.cpp

@@ -0,0 +1,128 @@
+/****************************************************************-*- C++ -*-****
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. *
+ * All rights reserved. *
+ * *
+ * This source code and the accompanying materials are made available under *
+ * the terms of the Apache License 2.0 which accompanies this distribution. *
+ ******************************************************************************/
+
+#include "StateDecomposer.h"
+
+namespace cudaq::details {
+
+std::vector<std::size_t> grayCode(std::size_t numBits) {
+ std::vector<std::size_t> result(1ULL << numBits);
+ for (std::size_t i = 0; i < (1ULL << numBits); ++i)
+ result[i] = ((i >> 1) ^ i);
+ return result;
+}
+
+std::vector<std::size_t> getControlIndices(std::size_t numBits) {
+ auto code = grayCode(numBits);
+ std::vector<std::size_t> indices;
+ for (auto i = 0u; i < code.size(); ++i) {
+ // The position of the control in the lth CNOT gate is set to match
+ // the position where the lth and (l + 1)th bit strings g[l] and g[l+1] of
+ // the binary reflected Gray code differ.
+ auto position = std::log2(code[i] ^ code[(i + 1) % code.size()]);
+ // N.B: In CUDA Quantum we write the least significant bit (LSb) on the left
+ //
+ // lsb -v
+ // 001
+ // ^- msb
+ //
+ // Meaning that the bitstring 001 represents the number four instead of one.
+ // The above position calculation uses the 'normal' convention of writing
+ // numbers with the LSb on the left.
+ //
+ // Now, what we need to find out is the position of the 1 in the bitstring.
+ // If we take LSb as being position 0, then for the normal convention its
+ // position will be 0. Using CUDA Quantum convention it will be 2. Hence,
+ // we need to convert the position we find using:
+ //
+ // numBits - position - 1
+ //
+ // The extra -1 is to account for indices starting at 0. Using the above
+ // examples:
+ //
+ // bitstring: 001
+ // numBits: 3
+ // position: 0
+ //
+ // We have the converted position: 2, which is what we need.
+ indices.emplace_back(numBits - position - 1);
+ }
+ return indices;
+}
+
+std::vector<double> convertAngles(const std::span<double> alphas) {
+ // Implements Eq. (3) from https://arxiv.org/pdf/quant-ph/0407010.pdf
+ //
+ // N.B: The paper does fails to explicitly define what is the dot operator in
+ // the exponent of -1. Ref. 3 solves the mystery: its the bitwise inner
+ // product.
+ auto bitwiseInnerProduct = [](std::size_t a, std::size_t b) {
+ auto product = a & b;
+ auto sumOfProducts = 0;
+ while (product) {
+ sumOfProducts += product & 0b1 ? 1 : 0;
+ product = product >> 1;
+ }
+ return sumOfProducts;
+ };
+ std::vector<double> thetas(alphas.size(), 0);
+ for (std::size_t i = 0u; i < alphas.size(); ++i) {
+ for (std::size_t j = 0u; j < alphas.size(); ++j)
+ thetas[i] +=
+ bitwiseInnerProduct(j, ((i >> 1) ^ i)) & 0b1 ? -alphas[j] : alphas[j];
+ thetas[i] /= alphas.size();
+ }
+ return thetas;
+}
+
+std::vector<double> getAlphaZ(const std::span<double> data,
+ std::size_t numQubits, std::size_t k) {
+ // Implements Eq. (5) from https://arxiv.org/pdf/quant-ph/0407010.pdf
+ std::vector<double> angles;
+ double divisor = static_cast<double>(1ULL << (k - 1));
+ for (std::size_t j = 1; j <= (1ULL << (numQubits - k)); ++j) {
+ double angle = 0.0;
+ for (std::size_t l = 1; l <= (1ULL << (k - 1)); ++l)
+ // N.B: There is an extra '-1' on these indices computations to account
+ // for the fact that our indices start at 0.
+ angle += data[(2 * j - 1) * (1 << (k - 1)) + l - 1] -
+ data[(2 * j - 2) * (1 << (k - 1)) + l - 1];
+ angles.push_back(angle / divisor);
+ }
+ return angles;
+}
+
+std::vector<double> getAlphaY(const std::span<double> data,
+ std::size_t numQubits, std::size_t k) {
+ // Implements Eq. (8) from https://arxiv.org/pdf/quant-ph/0407010.pdf
+ // N.B: There is an extra '-1' on these indices computations to account for
+ // the fact that our indices start at 0.
+ std::vector<double> angles;
+ for (std::size_t j = 1; j <= (1ULL << (numQubits - k)); ++j) {
+ double numerator = 0;
+ for (std::size_t l = 1; l <= (1ULL << (k - 1)); ++l) {
+ numerator +=
+ std::pow(std::abs(data[(2 * j - 1) * (1 << (k - 1)) + l - 1]), 2);
+ }
+
+ double denominator = 0;
+ for (std::size_t l = 1; l <= (1ULL << k); ++l) {
+ denominator += std::pow(std::abs(data[(j - 1) * (1 << k) + l - 1]), 2);
+ }
+
+ if (denominator == 0.0) {
+ assert(numerator == 0.0 &&
+ "If the denominator is zero, the numerator must also be zero.");
+ angles.push_back(0.0);
+ continue;
+ }
+ angles.push_back(2.0 * std::asin(std::sqrt(numerator / denominator)));
+ }
+ return angles;
+}
+} // namespace cudaq::details

lib/Optimizer/Transforms/StateDecomposer.h

@@ -0,0 +1,177 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. *
+ * All rights reserved. *
+ * *
+ * This source code and the accompanying materials are made available under *
+ * the terms of the Apache License 2.0 which accompanies this distribution. *
+ ******************************************************************************/
+
+#include "PassDetails.h"
+#include "cudaq/Optimizer/Builder/Runtime.h"
+#include "cudaq/Optimizer/CodeGen/QIRFunctionNames.h"
+#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeTypes.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Target/LLVMIR/TypeToLLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include <span>
+
+#include <iostream>
+
+namespace cudaq::details {
+
+/// @brief Converts angles of a uniformly controlled rotation to angles of
+/// non-controlled rotations.
+std::vector<double> convertAngles(const std::span<double> alphas);
+
+/// @brief Return the control indices dictated by the gray code implementation.
+///
+/// Here, numBits is the number of controls.
+std::vector<std::size_t> getControlIndices(std::size_t numBits);
+
+/// @brief Return angles required to implement a uniformly controlled z-rotation
+/// on the `kth` qubit.
+std::vector<double> getAlphaZ(const std::span<double> data,
+ std::size_t numQubits, std::size_t k);
+
+/// @brief Return angles required to implement a uniformly controlled y-rotation
+/// on the `kth` qubit.
+std::vector<double> getAlphaY(const std::span<double> data,
+ std::size_t numQubits, std::size_t k);
+} // namespace cudaq::details
+
+class StateGateBuilder {
+public:
+ StateGateBuilder(mlir::OpBuilder &b, mlir::Location &l, mlir::Value &q)
+ : builder(b), loc(l), qubits(q) {}
+
+ template <typename Op>
+ void applyRotationOp(double theta, std::size_t target) {
+ auto qubit = createQubitRef(target);
+ auto thetaValue = createAngleValue(theta);
+ builder.create<Op>(loc, thetaValue, mlir::ValueRange{}, qubit);
+ };
+
+ void applyX(std::size_t control, std::size_t target) {
+ auto qubitC = createQubitRef(control);
+ auto qubitT = createQubitRef(target);
+ builder.create<quake::XOp>(loc, qubitC, qubitT);
+ };
+
+private:
+ mlir::Value createQubitRef(std::size_t index) {
+ if (qubitRefs.contains(index)) {
+ return qubitRefs[index];
+ }
+
+ auto indexValue = builder.create<mlir::arith::ConstantIntOp>(
+ loc, index, builder.getIntegerType(64));
+ auto ref = builder.create<quake::ExtractRefOp>(loc, qubits, indexValue);
+ qubitRefs[index] = ref;
+ return ref;
+ }
+
+ mlir::Value createAngleValue(double angle) {
+ return builder.create<mlir::arith::ConstantFloatOp>(
+ loc, llvm::APFloat{angle}, builder.getF64Type());
+ }
+
+ mlir::OpBuilder &builder;
+ mlir::Location &loc;
+ mlir::Value &qubits;
+
+ std::unordered_map<std::size_t, mlir::Value> qubitRefs =
+ std::unordered_map<std::size_t, mlir::Value>();
+};
+
+class StateDecomposer {
+public:
+ StateDecomposer(StateGateBuilder &b, std::span<std::complex<double>> a)
+ : builder(b), amplitudes(a), numQubits(log2(a.size())) {}
+
+ /// @brief Decompose the input state vector data to a set of controlled
+ /// operations and rotations. This function takes as input a `OpBuilder`
+ /// and appends the operations of the decomposition to its internal
+ /// representation. This implementation follows the algorithm defined in
+ /// `https://arxiv.org/pdf/quant-ph/0407010.pdf`.
+ void decompose() {
+
+ // Decompose the state into phases and magnitudes.
+ bool needsPhaseEqualization = false;
+ std::vector<double> phases;
+ std::vector<double> magnitudes;
+ for (const auto &a : amplitudes) {
+ phases.push_back(std::arg(a));
+ magnitudes.push_back(std::abs(a));
+ // FIXME: remove magic number.
+ needsPhaseEqualization |= std::abs(phases.back()) > 1e-10;
+ }
+
+ // N.B: The algorithm, as described in the paper, creates a circuit that
+ // begins with a target state and brings it to the all zero state. Hence,
+ // this implementation do the two steps described in Section III in reverse
+ // order.
+
+ // Apply uniformly controlled y-rotations, the construction in Eq. (4).
+ for (std::size_t j = 1; j <= numQubits; ++j) {
+ auto k = numQubits - j + 1;
+ auto numControls = j - 1;
+ auto target = j - 1;
+ auto alphaYk = cudaq::details::getAlphaY(magnitudes, numQubits, k);
+ applyRotation<quake::RyOp>(alphaYk, numControls, target);
+ }
+
+ if (!needsPhaseEqualization)
+ return;
+
+ // Apply uniformly controlled z-rotations, the construction in Eq. (4).
+ for (std::size_t j = 1; j <= numQubits; ++j) {
+ auto k = numQubits - j + 1;
+ auto numControls = j - 1;
+ auto target = j - 1;
+ auto alphaZk = cudaq::details::getAlphaZ(phases, numQubits, k);
+ if (alphaZk.empty())
+ continue;
+ applyRotation<quake::RzOp>(alphaZk, numControls, target);
+ }
+ }
+
+private:
+ /// @brief Apply a uniformly controlled rotation on the target qubit.
+ template <typename Op>
+ void applyRotation(const std::span<double> alphas, std::size_t numControls,
+ std::size_t target) {
+
+ // In our model the index 1 (i.e. |01>) in quantum state data
+ // corresponds to qubits[0]=1 and qubits[1] = 0.
+ // Revert the order of qubits as the state preparation algorithm
+ // we use assumes the opposite.
+ auto qubitIndex = [&](std::size_t i) { return numQubits - i - 1; };
+
+ auto thetas = cudaq::details::convertAngles(alphas);
+ if (numControls == 0) {
+ builder.applyRotationOp<Op>(thetas[0], qubitIndex(target));
+ return;
+ }
+
+ auto controlIndices = cudaq::details::getControlIndices(numControls);
+ assert(thetas.size() == controlIndices.size());
+ for (auto [i, c] : llvm::enumerate(controlIndices)) {
+ builder.applyRotationOp<Op>(thetas[i], qubitIndex(target));
+ builder.applyX(qubitIndex(c), qubitIndex(target));
+ }
+ }
+
+ StateGateBuilder &builder;
+ std::span<std::complex<double>> amplitudes;
+ std::size_t numQubits;
+};