Skip to content

Commit

Permalink
* Saving progress on using ApplyOp
Browse files Browse the repository at this point in the history
  • Loading branch information
khalatepradnya committed Sep 2, 2024
1 parent c871ca3 commit 1cdfebf
Show file tree
Hide file tree
Showing 5 changed files with 102 additions and 63 deletions.
116 changes: 68 additions & 48 deletions lib/Optimizer/Transforms/UnitarySynthesis.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -67,14 +67,16 @@ struct BasisZYZ {

std::array<std::complex<double>, 4> matrix;
EulerAngles angles;
double phase;
/// Global phase is ignored
// double phase;

/// 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() {
auto det = (matrix[0] * matrix[3]) - (matrix[1] * matrix[2]);
phase = 0.5 * std::atan2(det.imag(), det.real());

// auto det = (matrix[0] * matrix[3]) - (matrix[1] * matrix[2]);
// phase = 0.5 * std::atan2(det.imag(), det.real());

auto abs_00 = std::abs(matrix[0]);
auto abs_01 = std::abs(matrix[1]);
Expand Down Expand Up @@ -118,60 +120,78 @@ class CustomUnitaryPattern
LogicalResult matchAndRewrite(quake::CustomUnitarySymbolOp customOp,
PatternRewriter &rewriter) const override {

// Fetch the unitary matrix generator for this custom operation
auto parentModule = customOp->getParentOfType<ModuleOp>();
auto sref = customOp.getGenerator();
StringRef generatorName = sref.getRootReference();
auto globalOp =
parentModule.lookupSymbol<cudaq::cc::GlobalOp>(generatorName);
auto unitary = readGlobalConstantArray(globalOp);

if (unitary.size() != 4)
return customOp.emitError(
"Decomposition of only single qubit custom operations supported.");

/// TODO: Maintain a cache of decomposed custom operations

// Use Euler angle decomposition for single qubit operation
auto zyz = BasisZYZ(unitary);
zyz.decompose();

// op info
Location loc = customOp->getLoc();
auto targets = customOp.getTargets();
auto controls = customOp.getControls();

/// TODO: Handle adjoint case

auto floatTy = cast<FloatType>(rewriter.getF64Type());

if (0. != zyz.angles.alpha) {
auto alpha = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.alpha, floatTy);
rewriter.create<quake::RzOp>(loc, alpha, controls, targets);
}

if (0. != zyz.angles.beta) {
auto beta = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.beta, floatTy);
rewriter.create<quake::RyOp>(loc, beta, controls, targets);
}
/// Get the global constant holding the concrete matrix corresponding to
/// this custom operation invocation
StringRef generatorName = customOp.getGenerator().getRootReference();
auto globalOp =
parentModule.lookupSymbol<cudaq::cc::GlobalOp>(generatorName);

if (0. != zyz.angles.gamma) {
auto gamma = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.gamma, floatTy);
rewriter.create<quake::RzOp>(loc, gamma, controls, targets);
/// The decomposed sequence of quantum operations are in a function
auto pair = generatorName.split(".rodata");
std::string funcName = pair.first.str() + ".kernel" + pair.second.str();

/// If the replacement function doesn't exist, create it here
if (!parentModule.lookupSymbol<LLVM::LLVMFuncOp>(funcName)) {

auto unitary = readGlobalConstantArray(globalOp);

/// TODO: Expand the logic to decompose upto 4-qubit operations
if (unitary.size() != 4)
return customOp.emitError(
"Decomposition of only single qubit custom operations supported.");

/// Controls are handled via apply specialization
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);
rewriter.restoreInsertionPoint(insPt);

insPt = rewriter.saveInsertionPoint();
auto *block = func.addEntryBlock();
rewriter.setInsertionPointToStart(block);

/// Use Euler angle decomposition for single qubit operation
auto zyz = BasisZYZ(unitary);
zyz.decompose();

/// For 1-qubit operation, apply on 'all' the targets
auto arguments = func.getArguments();
auto floatTy = cast<FloatType>(rewriter.getF64Type());

if (0. != zyz.angles.alpha) {
auto alpha = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.alpha, floatTy);
rewriter.create<quake::RzOp>(loc, alpha, ValueRange{}, arguments);
}

if (0. != zyz.angles.beta) {
auto beta = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.beta, floatTy);
rewriter.create<quake::RyOp>(loc, beta, ValueRange{}, arguments);
}

if (0. != zyz.angles.gamma) {
auto gamma = cudaq::opt::factory::createFloatConstant(
loc, rewriter, zyz.angles.gamma, floatTy);
rewriter.create<quake::RzOp>(loc, gamma, ValueRange{}, arguments);
}

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

/// ASKME: Apply global phase?
// if (0. != zyz.phase) {
// auto phase = cudaq::opt::factory::createFloatConstant(loc, rewriter,
// zyz.phase,
// floatTy);
// Value negPhase = rewriter.create<arith::NegFOp>(loc, phase);
// rewriter.create<quake::R1Op>(loc, phase, controls, targets);
// rewriter.create<quake::RzOp>(loc, negPhase, controls, targets);
// }
rewriter.create<quake::ApplyOp>(
loc, TypeRange{}, SymbolRefAttr::get(rewriter.getContext(), funcName),
customOp.isAdj(), controls, targets);

rewriter.eraseOp(customOp);
return success();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ config:
# Add the rest-qpu library to the link list
link-libs: ["-lcudaq-rest-qpu"]
# Define the lowering pipeline
platform-lowering-config: "const-prop-complex,canonicalize,cse,lift-array-value,state-prep,expand-measurements,unrolling-pipeline,func.func(unitary-synthesis),decomposition{enable-patterns=U3ToRotations},func.func(lower-to-cfg),canonicalize,func.func(multicontrol-decomposition),quantinuum-gate-set-mapping"
platform-lowering-config: "const-prop-complex,canonicalize,cse,lift-array-value,state-prep,expand-measurements,unrolling-pipeline,func.func(unitary-synthesis),canonicalize,apply-op-specialization,decomposition{enable-patterns=U3ToRotations},func.func(lower-to-cfg),canonicalize,func.func(multicontrol-decomposition),quantinuum-gate-set-mapping"
# Tell the rest-qpu that we are generating Adaptive QIR.
codegen-emission: qir-adaptive
# Library mode is only for simulators, physical backends must turn this off
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
# CHECK-DAG: LINKLIBS="${LINKLIBS} -lcudaq-rest-qpu"

# Define the lowering pipeline, here we lower to Adaptive QIR
# CHECK-DAG: PLATFORM_LOWERING_CONFIG="const-prop-complex,canonicalize,cse,lift-array-value,state-prep,expand-measurements,unrolling-pipeline,func.func(unitary-synthesis),decomposition{enable-patterns=U3ToRotations},func.func(lower-to-cfg),canonicalize,func.func(multicontrol-decomposition),quantinuum-gate-set-mapping"
# CHECK-DAG: PLATFORM_LOWERING_CONFIG="const-prop-complex,canonicalize,cse,lift-array-value,state-prep,expand-measurements,unrolling-pipeline,func.func(unitary-synthesis),canonicalize,apply-op-specialization,decomposition{enable-patterns=U3ToRotations},func.func(lower-to-cfg),canonicalize,func.func(multicontrol-decomposition),quantinuum-gate-set-mapping"

# Tell the rest-qpu that we are generating QIR.
# CHECK-DAG: CODEGEN_EMISSION=qir-adaptive
Expand Down
12 changes: 9 additions & 3 deletions test/Transforms/UnitarySynthesis/basic_test.qke
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,15 @@ module attributes {quake.mangled_name_map = {__nvqpp__mlirgen__k1 = "__nvqpp__ml
cc.global constant @__nvqpp__mlirgen__custom_x_generator_1.rodata (dense<[(0.000000e+00,0.000000e+00), (1.000000e+00,0.000000e+00), (1.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00)]> : tensor<4xcomplex<f64>>) : !cc.array<complex<f64> x 4>
}

// CHECK-LABEL: func.func @__nvqpp__mlirgen__custom_x_generator_1.kernel(
// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) {
// CHECK: %[[VAL_1:.*]] = arith.constant 3.1415926535897931 : f64
// CHECK: quake.ry (%[[VAL_1]]) %[[VAL_0]] : (f64, !quake.ref) -> ()
// CHECK: return
// CHECK: }

// CHECK-LABEL: func.func @__nvqpp__mlirgen__k1() attributes {"cudaq-entrypoint"} {
// CHECK: %[[VAL_0:.*]] = arith.constant 3.1415926535897931 : f64
// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.ref
// CHECK: quake.ry (%[[VAL_0]]) %[[VAL_1]] : (f64, !quake.ref) -> ()
// CHECK: %[[VAL_0:.*]] = quake.alloca !quake.ref
// CHECK: quake.apply @__nvqpp__mlirgen__custom_x_generator_1.kernel %[[VAL_0]] : (!quake.ref) -> ()
// CHECK: return
// CHECK: }
33 changes: 23 additions & 10 deletions test/Transforms/UnitarySynthesis/bell_pair.qke
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
// the terms of the Apache License 2.0 which accompanies this distribution. //
// ========================================================================== //

// RUN: cudaq-opt --unitary-synthesis %s | FileCheck %s
// RUN: cudaq-opt --unitary-synthesis --canonicalize %s | FileCheck %s

module attributes {quake.mangled_name_map = {__nvqpp__mlirgen__bell = "__nvqpp__mlirgen__bell_PyKernelEntryPointRewrite"}} {
func.func @__nvqpp__mlirgen__bell() attributes {"cudaq-entrypoint"} {
Expand All @@ -21,15 +21,28 @@ module attributes {quake.mangled_name_map = {__nvqpp__mlirgen__bell = "__nvqpp__
cc.global constant @__nvqpp__mlirgen__custom_x_generator_1.rodata (dense<[(0.000000e+00,0.000000e+00), (1.000000e+00,0.000000e+00), (1.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00)]> : tensor<4xcomplex<f64>>) : !cc.array<complex<f64> x 4>
}

// CHECK-LABEL: func.func @__nvqpp__mlirgen__bell() attributes {"cudaq-entrypoint"} {
// CHECK: %[[VAL_0:.*]] = arith.constant 3.1415926535897931 : f64
// CHECK-LABEL: func.func @__nvqpp__mlirgen__custom_h_generator_1.kernel(
// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) {
// CHECK: %[[VAL_1:.*]] = arith.constant 1.5707963267948968 : f64
// CHECK: %[[VAL_2:.*]] = quake.alloca !quake.veq<2>
// CHECK: %[[VAL_3:.*]] = quake.extract_ref %[[VAL_2]][0] : (!quake.veq<2>) -> !quake.ref
// CHECK: quake.rz (%[[VAL_0]]) %[[VAL_3]] : (f64, !quake.ref) -> ()
// CHECK: quake.ry (%[[VAL_1]]) %[[VAL_3]] : (f64, !quake.ref) -> ()
// CHECK: quake.rz (%[[VAL_0]]) %[[VAL_3]] : (f64, !quake.ref) -> ()
// CHECK: %[[VAL_4:.*]] = quake.extract_ref %[[VAL_2]][1] : (!quake.veq<2>) -> !quake.ref
// CHECK: quake.ry (%[[VAL_0]]) {{\[}}%[[VAL_3]]] %[[VAL_4]] : (f64, !quake.ref, !quake.ref) -> ()
// CHECK: %[[VAL_2:.*]] = arith.constant 3.1415926535897931 : f64
// CHECK: quake.rz (%[[VAL_2]]) %[[VAL_0]] : (f64, !quake.ref) -> ()
// CHECK: quake.ry (%[[VAL_1]]) %[[VAL_0]] : (f64, !quake.ref) -> ()
// CHECK: quake.rz (%[[VAL_2]]) %[[VAL_0]] : (f64, !quake.ref) -> ()
// CHECK: return
// CHECK: }

// CHECK-LABEL: func.func @__nvqpp__mlirgen__custom_x_generator_1.kernel(
// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) {
// CHECK: %[[VAL_1:.*]] = arith.constant 3.1415926535897931 : f64
// CHECK: quake.ry (%[[VAL_1]]) %[[VAL_0]] : (f64, !quake.ref) -> ()
// CHECK: return
// CHECK: }

// CHECK-LABEL: func.func @__nvqpp__mlirgen__bell() attributes {"cudaq-entrypoint"} {
// CHECK: %[[VAL_0:.*]] = quake.alloca !quake.veq<2>
// CHECK: %[[VAL_1:.*]] = quake.extract_ref %[[VAL_0]][0] : (!quake.veq<2>) -> !quake.ref
// CHECK: quake.apply @__nvqpp__mlirgen__custom_h_generator_1.kernel %[[VAL_1]] : (!quake.ref) -> ()
// CHECK: %[[VAL_2:.*]] = quake.extract_ref %[[VAL_0]][1] : (!quake.veq<2>) -> !quake.ref
// CHECK: quake.apply @__nvqpp__mlirgen__custom_x_generator_1.kernel {{\[}}%[[VAL_1]]] %[[VAL_2]] : (!quake.ref, !quake.ref) -> ()
// CHECK: return
// CHECK: }

0 comments on commit 1cdfebf

Please sign in to comment.