Skip to content

Commit

Permalink
* Added failing test(s)
Browse files Browse the repository at this point in the history
* Skeleton for a new decomposer class
* Refactoring
* Saving progress
  • Loading branch information
khalatepradnya committed Sep 19, 2024
1 parent 2c05ff1 commit 4f1c808
Show file tree
Hide file tree
Showing 4 changed files with 287 additions and 67 deletions.
5 changes: 5 additions & 0 deletions lib/Optimizer/Transforms/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -65,3 +65,8 @@ add_cudaq_library(OptTransforms
OptimBuilder
QuakeDialect
)

target_include_directories(OptTransforms SYSTEM
PRIVATE ${CMAKE_SOURCE_DIR}/tpls/eigen
PRIVATE ${CMAKE_SOURCE_DIR}/runtime
)
308 changes: 244 additions & 64 deletions lib/Optimizer/Transforms/UnitarySynthesis.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
******************************************************************************/

#include "PassDetails.h"
#include "common/EigenDense.h"
#include "cudaq/Optimizer/Builder/Factory.h"
#include "cudaq/Optimizer/CodeGen/Passes.h"
#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
Expand All @@ -30,13 +31,29 @@ namespace cudaq::opt {
using namespace mlir;

namespace {

class Decomposer {
public:
virtual void decompose() = 0;
virtual void emitDecomposedFuncOp(quake::CustomUnitarySymbolOp customOp,
PatternRewriter &rewriter,
std::string funcName) = 0;
/// Ignore angles less than some threshold
bool isAboveThreshold(double value) {
const double epsilon = 1e-9;
return std::abs(value) > epsilon;
};
virtual ~Decomposer() = default;
};

struct EulerAngles {
double alpha;
double beta;
double gamma;
};

struct BasisZYZ {
struct OneQubitOpZYZ : public Decomposer {
/// TODO: Update to use Eigen matrix
std::array<std::complex<double>, 4> matrix;
EulerAngles angles;
/// Updates to the global phase
Expand All @@ -45,7 +62,7 @@ struct BasisZYZ {
/// This logic is based on https://arxiv.org/pdf/quant-ph/9503016 and its
/// corresponding explanation in https://threeplusone.com/pubs/on_gates.pdf,
/// Section 4.
void decompose() {
void decompose() override {
using namespace std::complex_literals;
/// Rescale the input unitary matrix, `u`, to be special unitary.
/// Extract a phase factor, `phase`, so that
Expand All @@ -68,10 +85,217 @@ struct BasisZYZ {
angles.gamma = sum - diff;
}

BasisZYZ(const std::vector<std::complex<double>> &vec) {
void emitDecomposedFuncOp(quake::CustomUnitarySymbolOp customOp,
PatternRewriter &rewriter,
std::string funcName) override {
auto parentModule = customOp->getParentOfType<ModuleOp>();
Location loc = customOp->getLoc();
auto targets = customOp.getTargets();
auto funcTy =
FunctionType::get(parentModule.getContext(), targets[0].getType(), {});
auto insPt = rewriter.saveInsertionPoint();
rewriter.setInsertionPointToStart(parentModule.getBody());
auto func =
rewriter.create<func::FuncOp>(parentModule->getLoc(), funcName, funcTy);
func.setPrivate();
auto *block = func.addEntryBlock();
rewriter.setInsertionPointToStart(block);
auto arguments = func.getArguments();
FloatType floatTy = rewriter.getF64Type();

/// NOTE: Operator notation is right-to-left, whereas circuit notation
/// is left-to-right. Hence, angles are applied as
/// Rz(gamma)Ry(beta)Rz(alpha)
if (isAboveThreshold(angles.gamma)) {
auto gamma = cudaq::opt::factory::createFloatConstant(
loc, rewriter, angles.gamma, floatTy);
rewriter.create<quake::RzOp>(loc, gamma, ValueRange{}, arguments);
}
if (isAboveThreshold(angles.beta)) {
auto beta = cudaq::opt::factory::createFloatConstant(
loc, rewriter, angles.beta, floatTy);
rewriter.create<quake::RyOp>(loc, beta, ValueRange{}, arguments);
}
if (isAboveThreshold(angles.alpha)) {
auto alpha = cudaq::opt::factory::createFloatConstant(
loc, rewriter, angles.alpha, floatTy);
rewriter.create<quake::RzOp>(loc, alpha, ValueRange{}, arguments);
}
/// NOTE: Typically global phase can be ignored but, if this decomposition
/// is applied in a kernel that is called with `cudaq::control`, the global
/// phase will become a local phase and give a wrong result if we don't keep
/// track of that.
/// NOTE: R1-Rz pair results in a half the applied global phase angle,
/// hence, we need to multiply the angle by 2
auto globalPhase = 2 * phase;
if (isAboveThreshold(globalPhase)) {
auto phase = cudaq::opt::factory::createFloatConstant(
loc, rewriter, globalPhase, floatTy);
Value negPhase = rewriter.create<arith::NegFOp>(loc, phase);
rewriter.create<quake::R1Op>(loc, phase, ValueRange{}, arguments[0]);
rewriter.create<quake::RzOp>(loc, negPhase, ValueRange{}, arguments[0]);
}
rewriter.create<func::ReturnOp>(loc);
rewriter.restoreInsertionPoint(insPt);
}

OneQubitOpZYZ(const std::vector<std::complex<double>> &vec) {
std::copy(vec.begin(), vec.begin() + 4, matrix.begin());
decompose();
}

OneQubitOpZYZ(const Eigen::Matrix2cd &vec) {
for (size_t r = 0; r < 2; r++)
for (size_t c = 0; c < 2; c++)
matrix[r * 2 + c] = vec(r, c);
decompose();
}
};

struct KAKComponents {
// KAK decomposition allows to express arbitrary 2-qubit unitary (U) in the
// form: U = (a1 ⊗ a0) x exp(i(xXX + yYY + zZZ)) x (b1 ⊗ b0) where, a0, a1
// (after), b0, b1 (before) are single qubit operations, and the exponential
// is specified by the 3 elements of canonical class vector x, y, z
Eigen::Matrix2cd a0;
Eigen::Matrix2cd a1;
Eigen::Matrix2cd b0;
Eigen::Matrix2cd b1;
double x;
double y;
double z;
};

const Eigen::Matrix4cd &MagicBasisMatrix() {
static Eigen::Matrix4cd MagicBasisMatrix;
constexpr std::complex<double> i{0.0, 1.0};
MagicBasisMatrix << 1.0, 0.0, 0.0, i, 0.0, i, 1.0, 0, 0, i, -1.0, 0, 1.0, 0,
0, -i;
MagicBasisMatrix = MagicBasisMatrix * M_SQRT1_2;
return MagicBasisMatrix;
}

const Eigen::Matrix4cd &MagicBasisMatrixAdj() {
static Eigen::Matrix4cd MagicBasisMatrixAdj = MagicBasisMatrix().adjoint();
return MagicBasisMatrixAdj;
}

const Eigen::Matrix4cd &GammaFactor() {
/// Gamma matrix = +1 +1 −1 +1
/// +1 +1 +1 −1
/// +1 −1 −1 −1
/// +1 −1 +1 +1
static Eigen::Matrix4cd GammaT;
GammaT << 1, 1, 1, 1, 1, 1, -1, -1, -1, 1, -1, 1, 1, -1, -1, 1;
GammaT /= 4;
return GammaT;
}

std::pair<Eigen::Matrix2cd, Eigen::Matrix2cd>
extractSU2FromSO4(const Eigen::Matrix4cd &matrix) {
Eigen::Matrix2cd part1 = matrix.block<2, 2>(0, 0);
Eigen::Matrix2cd part2 = matrix.block<2, 2>(2, 2);
return std::make_pair(part1, part2);
}

struct TwoQubitOpKAK : public Decomposer {
Eigen::Matrix4cd matrix;
KAKComponents components;
/// Updates to the global phase
std::complex<double> phase;

/// This logic is based on the Cartan's KAK decomposition.
/// Ref: https://arxiv.org/pdf/quant-ph/0507171
void decompose() override {
using namespace std::complex_literals;
/// Step0: Convert to special unitary
auto normFactor = std::pow(matrix.determinant(), 0.25);
auto specialUnitary = matrix / normFactor;
/// Step1: Convert the target matrix into magic basis
Eigen::Matrix4cd matrixMagicBasis =
MagicBasisMatrixAdj() * specialUnitary * MagicBasisMatrix();
/// Step2: Diagonalize the matrix
Eigen::JacobiSVD<Eigen::Matrix4cd> svd(
matrixMagicBasis, Eigen::ComputeFullU | Eigen::ComputeFullV);
/// Step3: Get the KAK components
auto before = extractSU2FromSO4(svd.matrixU());
auto after = extractSU2FromSO4(svd.matrixV().transpose());
components.a0 = after.first;
components.a1 = after.second;
components.b0 = before.first;
components.b1 = before.second;
/// Step4: Get the coefficients of canonical class vector
Eigen::Vector4cd diagonalPhases;
for (size_t i = 0; i < 4; i++)
diagonalPhases(i) = std::arg(svd.singularValues()(i));
auto coefficients = GammaFactor() * diagonalPhases;
components.x = coefficients(1).real();
components.y = coefficients(2).real();
components.z = coefficients(3).real();
phase = std::exp(1i * coefficients(0)) + std::arg(normFactor);
/// Step5: Canonicalize coefficients?
}

void emitDecomposedFuncOp(quake::CustomUnitarySymbolOp customOp,
PatternRewriter &rewriter,
std::string funcName) override {
auto b0 = OneQubitOpZYZ(components.b0);
b0.emitDecomposedFuncOp(customOp, rewriter, funcName + "b0");
auto b1 = OneQubitOpZYZ(components.b1);
b1.emitDecomposedFuncOp(customOp, rewriter, funcName + "b1");
auto a0 = OneQubitOpZYZ(components.a0);
a0.emitDecomposedFuncOp(customOp, rewriter, funcName + "a0");
auto a1 = OneQubitOpZYZ(components.a1);
a1.emitDecomposedFuncOp(customOp, rewriter, funcName + "a1");

auto parentModule = customOp->getParentOfType<ModuleOp>();
Location loc = customOp->getLoc();
auto targets = customOp.getTargets();
auto funcTy =
FunctionType::get(parentModule.getContext(), targets.getTypes(), {});
auto insPt = rewriter.saveInsertionPoint();
rewriter.setInsertionPointToStart(parentModule.getBody());
auto func =
rewriter.create<func::FuncOp>(parentModule->getLoc(), funcName, funcTy);
func.setPrivate();
auto *block = func.addEntryBlock();
rewriter.setInsertionPointToStart(block);
auto arguments = func.getArguments();
FloatType floatTy = rewriter.getF64Type();
rewriter.create<quake::ApplyOp>(
loc, TypeRange{},
SymbolRefAttr::get(rewriter.getContext(), funcName + "b0"), false,
ValueRange{}, ValueRange{arguments[0]});
rewriter.create<quake::ApplyOp>(
loc, TypeRange{},
SymbolRefAttr::get(rewriter.getContext(), funcName + "b1"), false,
ValueRange{}, ValueRange{arguments[1]});
/// TODO: Insert XYZ coefficient gates
rewriter.create<quake::ApplyOp>(
loc, TypeRange{},
SymbolRefAttr::get(rewriter.getContext(), funcName + "a0"), false,
ValueRange{}, ValueRange{arguments[0]});
rewriter.create<quake::ApplyOp>(
loc, TypeRange{},
SymbolRefAttr::get(rewriter.getContext(), funcName + "a1"), false,
ValueRange{}, ValueRange{arguments[1]});
auto globalPhase = 2 * std::abs(phase);
if (isAboveThreshold(globalPhase)) {
auto phase = cudaq::opt::factory::createFloatConstant(
loc, rewriter, globalPhase, floatTy);
Value negPhase = rewriter.create<arith::NegFOp>(loc, phase);
rewriter.create<quake::R1Op>(loc, phase, ValueRange{}, arguments[0]);
rewriter.create<quake::RzOp>(loc, negPhase, ValueRange{}, arguments[0]);
}

rewriter.create<func::ReturnOp>(loc);
rewriter.restoreInsertionPoint(insPt);
}

TwoQubitOpKAK(const Eigen::MatrixXcd &vec) {
matrix = vec;
decompose();
}
};

class CustomUnitaryPattern
Expand All @@ -82,9 +306,6 @@ class CustomUnitaryPattern
LogicalResult matchAndRewrite(quake::CustomUnitarySymbolOp customOp,
PatternRewriter &rewriter) const override {
auto parentModule = customOp->getParentOfType<ModuleOp>();
Location loc = customOp->getLoc();
auto targets = customOp.getTargets();
auto controls = customOp.getControls();
/// Get the global constant holding the concrete matrix corresponding to
/// this custom operation invocation
StringRef generatorName = customOp.getGenerator().getRootReference();
Expand All @@ -95,72 +316,31 @@ class CustomUnitaryPattern
std::string funcName = pair.first.str() + ".kernel" + pair.second.str();
/// If the replacement function doesn't exist, create it here
if (!parentModule.lookupSymbol<func::FuncOp>(funcName)) {
auto unitary = cudaq::opt::factory::readGlobalConstantArray(globalOp);
/// TODO: Expand the logic to decompose upto 4-qubit operations
if (unitary.size() != 4) {
auto matrix = cudaq::opt::factory::readGlobalConstantArray(globalOp);
switch (matrix.size()) {
case 4: {
auto zyz = OneQubitOpZYZ(matrix);
zyz.emitDecomposedFuncOp(customOp, rewriter, funcName);
} break;
case 16: {
auto unitary = Eigen::Map<Eigen::Matrix4cd>(matrix.data());
if (!unitary.isUnitary()) {
customOp.emitWarning("The custom operation matrix must be unitary.");
return failure();
}
auto kak = TwoQubitOpKAK(unitary);
kak.emitDecomposedFuncOp(customOp, rewriter, funcName);
} break;
default:
customOp.emitWarning(
"Decomposition of only single qubit custom operations supported.");
return failure();
}
/// Controls are handled via apply specialization, hence not included in
/// arguments
auto funcTy =
FunctionType::get(parentModule.getContext(), targets.getTypes(), {});
auto insPt = rewriter.saveInsertionPoint();
rewriter.setInsertionPointToStart(parentModule.getBody());
auto func = rewriter.create<func::FuncOp>(parentModule->getLoc(),
funcName, funcTy);
func.setPrivate();
auto *block = func.addEntryBlock();
rewriter.setInsertionPointToStart(block);
/// Use Euler angle decomposition for single qubit operation
auto zyz = BasisZYZ(unitary);
/// For 1-qubit operation, apply on 'all' the targets
auto arguments = func.getArguments();
FloatType floatTy = rewriter.getF64Type();
/// Ignore angles less than some threshold
auto isAboveThreshold = [&](auto value) {
const double epsilon = 1e-9;
return std::abs(value) > epsilon;
};
/// NOTE: Operator notation is right-to-left, whereas circuit notation is
/// left-to-right. Hence, angles are applied as Rz(gamma)Ry(beta)Rz(alpha)
if (isAboveThreshold(zyz.angles.gamma)) {
auto gamma = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.gamma, floatTy);
rewriter.create<quake::RzOp>(loc, gamma, ValueRange{}, arguments);
}
if (isAboveThreshold(zyz.angles.beta)) {
auto beta = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.beta, floatTy);
rewriter.create<quake::RyOp>(loc, beta, ValueRange{}, arguments);
}
if (isAboveThreshold(zyz.angles.alpha)) {
auto alpha = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.alpha, floatTy);
rewriter.create<quake::RzOp>(loc, alpha, ValueRange{}, arguments);
}
/// NOTE: Typically global phase can be ignored but, if this decomposition
/// is applied in a kernel that is called with `cudaq::control`, the
/// global phase will become a local phase and give a wrong result if we
/// don't keep track of that.
/// NOTE: R1-Rz pair results in a half the applied global phase angle,
/// hence, we need to multiply the angle by 2
auto globalPhase = 2 * zyz.phase;
if (isAboveThreshold(globalPhase)) {
auto phase = cudaq::opt::factory::createFloatConstant(
loc, rewriter, globalPhase, floatTy);
Value negPhase = rewriter.create<arith::NegFOp>(loc, phase);
rewriter.create<quake::R1Op>(loc, phase, ValueRange{}, arguments[0]);
rewriter.create<quake::RzOp>(loc, negPhase, ValueRange{}, arguments[0]);
}
rewriter.create<func::ReturnOp>(loc);
rewriter.restoreInsertionPoint(insPt);
}
rewriter.replaceOpWithNewOp<quake::ApplyOp>(
customOp, TypeRange{},
SymbolRefAttr::get(rewriter.getContext(), funcName), customOp.isAdj(),
controls, targets);
customOp.getControls(), customOp.getTargets());
return success();
}
};
Expand Down
Loading

0 comments on commit 4f1c808

Please sign in to comment.