Skip to content

Commit

Permalink
Pass a list of controls to builder rotation gates
Browse files Browse the repository at this point in the history
Signed-off-by: A.M. Santana <[email protected]>

Allows for a list of individual controls to rotation gates in C++ and python builders
New C++ unit tests for controlled rotations
New python unit tests for controlled rotations
New python compiler test for controlled rotations
  • Loading branch information
anthony-santana authored Nov 7, 2023
1 parent 6eb43e0 commit 339fefe
Show file tree
Hide file tree
Showing 10 changed files with 968 additions and 85 deletions.
2 changes: 2 additions & 0 deletions include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ constexpr static const char QIRReadResultBody[] =

constexpr static const char NVQIRInvokeWithControlBits[] =
"invokeWithControlQubits";
constexpr static const char NVQIRInvokeRotationWithControlBits[] =
"invokeRotationWithControlQubits";
constexpr static const char NVQIRPackSingleQubitInArray[] =
"packSingleQubitInArray";
constexpr static const char NVQIRReleasePackedQubitArray[] =
Expand Down
164 changes: 110 additions & 54 deletions lib/Optimizer/CodeGen/LowerToQIR.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -502,13 +502,7 @@ class OneTargetOneParamRewrite : public ConvertOpToLLVMPattern<OP> {
ConversionPatternRewriter &rewriter) const override {
auto instName = instOp->getName().stripDialect().str();
auto numControls = instOp.getControls().size();

// TODO: handle general control ops. For now, only allow Rn with 1
// control
if (numControls > 1)
return instOp.emitError("unsupported controlled op " + instName +
" with " + std::to_string(numControls) +
" ctrl qubits");
auto instOperands = adaptor.getOperands();

auto loc = instOp.getLoc();
ModuleOp parentModule = instOp->template getParentOfType<ModuleOp>();
Expand All @@ -527,11 +521,10 @@ class OneTargetOneParamRewrite : public ConvertOpToLLVMPattern<OP> {
v = rewriter.create<arith::ExtFOp>(loc, rewriter.getF64Type(), v);
return v;
};
Value v =
instOp.getIsAdj()
? rewriter.create<arith::NegFOp>(loc, adaptor.getOperands()[0])
: adaptor.getOperands()[0];
funcArgs.push_back(castToDouble(v));
Value val = instOp.getIsAdj()
? rewriter.create<arith::NegFOp>(loc, instOperands[0])
: instOperands[0];
funcArgs.push_back(castToDouble(val));

// If no controls, then this is easy
if (numControls == 0) {
Expand All @@ -549,75 +542,138 @@ class OneTargetOneParamRewrite : public ConvertOpToLLVMPattern<OP> {
return success();
}

// We have 1 control, is it a veq or a ref?
auto control = *instOp.getControls().begin();
qirFunctionName += "__ctl";
auto negateFunctionName = qirFunctionName + "x";
auto negatedQubitCtrls = instOp.getNegatedQubitControls();

// All signatures will take an Array* for the controls
tmpArgTypes.push_back(qubitArrayType);

// If type is a VeqType, then we're good, just forward to the call op
// We have >= 1 control, is the first a veq or a ref?
auto control = *instOp.getControls().begin();
Type type = control.getType();
if (type.isa<quake::VeqType>()) {

// Add the control array to the args.
funcArgs.push_back(adaptor.getControls().front());

// This is a single target op, add that type
tmpArgTypes.push_back(qubitIndexType);
// Return an error if the type is not a veq or ref.
if (!type.isa<quake::VeqType, quake::RefType>()) {
std::string typeName;
llvm::raw_string_ostream typeNameStream(typeName);
typeNameStream << type;
return instOp.emitError("unsupported control type for " + instName +
": " + typeName);
}

FlatSymbolRefAttr instSymbolRef =
// If the control is a qubit, we need to pack it into an Array*
if (numControls == 1 && type.isa<quake::RefType>()) {
FlatSymbolRefAttr packingSymbolRef =
cudaq::opt::factory::createLLVMFunctionSymbol(
qirFunctionName, /*return type=*/LLVM::LLVMVoidType::get(context),
std::move(tmpArgTypes), parentModule);
cudaq::opt::NVQIRPackSingleQubitInArray,
/*return type=*/qubitArrayType, {qubitIndexType}, parentModule);

// Pack the qubit into the array
Value result =
rewriter
.create<LLVM::CallOp>(loc, qubitArrayType, packingSymbolRef,
adaptor.getControls().front())
.getResult();
// The array result is what we want for the function args
funcArgs.push_back(result);
}

// Add the target op
funcArgs.push_back(adaptor.getTargets().front());
// If type is a VeqType, then we're good, just forward to the call op
if (numControls == 1 && type.isa<quake::VeqType>()) {
if (negatedQubitCtrls)
return instOp.emitError("unsupported controlled op " + instName +
" with vector of ctrl qubits");

// Here we already have and Array*, Qubit*
rewriter.replaceOpWithNewOp<LLVM::CallOp>(instOp, TypeRange{},
instSymbolRef, funcArgs);
return success();
// Add the control array to the args.
funcArgs.push_back(adaptor.getControls().front());
}

// If the control is a qubit, we need to pack it into an Array*
FlatSymbolRefAttr packingSymbolRef =
cudaq::opt::factory::createLLVMFunctionSymbol(
cudaq::opt::NVQIRPackSingleQubitInArray,
/*return type=*/qubitArrayType, {qubitIndexType}, parentModule);
// Pack the qubit into the array
Value result =
rewriter
.create<LLVM::CallOp>(loc, qubitArrayType, packingSymbolRef,
adaptor.getControls().front())
.getResult();
// The array result is what we want for the function args
funcArgs.push_back(result);

// This is a single target op, add that type
tmpArgTypes.push_back(qubitIndexType);

// __quantum__qis__NAME__ctl(double, Array*, Qubit*) Type
auto instOpQISFunctionType = LLVM::LLVMFunctionType::get(
LLVM::LLVMVoidType::get(context), tmpArgTypes);

// Get function pointer to ctrl operation
FlatSymbolRefAttr instSymbolRef =
cudaq::opt::factory::createLLVMFunctionSymbol(
qirFunctionName, /*return type=*/LLVM::LLVMVoidType::get(context),
std::move(tmpArgTypes), parentModule);

// Now we know we have the proper veq/ref type, we can
// handle multiple controls.
if (numControls > 1) {

// Get symbol for
// void invokeWithControlQubits(const std::size_t nControls, void
// (*QISFunction)(double, Array*, Qubit*), Qubit*, ...);
auto applyMultiControlFunction =
cudaq::opt::factory::createLLVMFunctionSymbol(
cudaq::opt::NVQIRInvokeRotationWithControlBits,
LLVM::LLVMVoidType::get(context),
{rewriter.getF64Type(), rewriter.getI64Type(),
LLVM::LLVMPointerType::get(instOpQISFunctionType)},
parentModule, true);

Value ctrlOpPointer = rewriter.create<LLVM::AddressOfOp>(
loc, LLVM::LLVMPointerType::get(instOpQISFunctionType),
instSymbolRef);

auto arraySize = rewriter.create<LLVM::ConstantOp>(
loc, rewriter.getI64Type(), numControls);

// This will need the numControls, function pointer, and all Qubit*
// operands
SmallVector<Value> args = {castToDouble(val), arraySize, ctrlOpPointer};
FlatSymbolRefAttr negateFuncRef;
if (negatedQubitCtrls) {
negateFuncRef = cudaq::opt::factory::createLLVMFunctionSymbol(
negateFunctionName,
/*return type=*/LLVM::LLVMVoidType::get(context),
{cudaq::opt::getQubitType(context)}, parentModule);
for (auto v : llvm::enumerate(instOperands)) {
if ((v.index() < numControls) && (*negatedQubitCtrls)[v.index()])
rewriter.create<LLVM::CallOp>(loc, TypeRange{}, negateFuncRef,
v.value());
args.push_back(v.value());
}
} else {
args.append(instOperands.begin(), instOperands.end());
}

// Call our utility function.
rewriter.replaceOpWithNewOp<LLVM::CallOp>(
instOp, TypeRange{}, applyMultiControlFunction, args);

if (negatedQubitCtrls) {
for (auto v : llvm::enumerate(instOperands))
if ((v.index() < numControls) && (*negatedQubitCtrls)[v.index()])
rewriter.create<LLVM::CallOp>(loc, TypeRange{}, negateFuncRef,
v.value());
}

return success();
}

// Add the target op
funcArgs.push_back(adaptor.getTargets().front());

// Here we already have and Array*, Qubit*
rewriter.replaceOpWithNewOp<LLVM::CallOp>(instOp, TypeRange{},
instSymbolRef, funcArgs);

// We need to release the control Array.
FlatSymbolRefAttr releaseSymbolRef =
cudaq::opt::factory::createLLVMFunctionSymbol(
cudaq::opt::NVQIRReleasePackedQubitArray,
/*return type=*/LLVM::LLVMVoidType::get(context), {qubitArrayType},
parentModule);
Value ctrlArray = funcArgs[1];
rewriter.create<LLVM::CallOp>(loc, TypeRange{}, releaseSymbolRef,
ctrlArray);
// If the control was a ref, we need to release the control Array.
if (type.isa<quake::RefType>()) {
FlatSymbolRefAttr releaseSymbolRef =
cudaq::opt::factory::createLLVMFunctionSymbol(
cudaq::opt::NVQIRReleasePackedQubitArray,
/*return type=*/LLVM::LLVMVoidType::get(context),
{qubitArrayType}, parentModule);
Value ctrlArray = funcArgs[1];
rewriter.create<LLVM::CallOp>(loc, TypeRange{}, releaseSymbolRef,
ctrlArray);
}

return success();
}
Expand Down
54 changes: 52 additions & 2 deletions python/runtime/cudaq/builder/py_kernel_builder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ and rotations and return a valid, callable, CUDA Quantum kernel.
" to the given target qubits, with the provided list of control " \
"qubits.\n" \
"\nArgs:\n" \
" controls (:class:`QuakeValue`): The list of qubits to use as " \
" controls (list[QuakeValue]): The list of qubits to use as " \
"controls for the operation.\n" \
" target (:class:`QuakeValue`): The target qubit of the " \
"operation.\n" \
Expand Down Expand Up @@ -345,7 +345,57 @@ and rotations and return a valid, callable, CUDA Quantum kernel.
" target = kernel.qalloc()\n" \
" # Apply a controlled-" #NAME " between the two qubits.\n" \
" kernel.c" #NAME \
"(parameter=3.14, control=control, target=target)\n")
"(parameter=3.14, control=control, target=target)\n") \
.def( \
"c" #NAME, \
[](kernel_builder<> &self, QuakeValue &parameter, \
std::vector<QuakeValue> &controls, \
QuakeValue &target) { self.NAME(parameter, controls, target); }, \
py::arg("parameter"), py::arg("controls"), py::arg("target"), \
"Apply a controlled-" #NAME " operation" \
" to the given target qubit, with the provided list of control " \
"qubits.\n" \
"\nArgs:\n" \
" parameter (float): The kernel argument to " \
"parameterize the " #NAME " gate over.\n" \
" controls (list[QuakeValue]): The control qubits for the " \
"operation.\n" \
" target (:class:`QuakeValue`): The target qubit of the " \
"operation.\n" \
"\n.. code-block:: python\n\n" \
" # Example:\n" \
" kernel, angle = cudaq.make_kernel(float)\n" \
" c1 = kernel.qalloc()\n" \
" c2 = kernel.qalloc()\n" \
" target = kernel.qalloc()\n" \
" # Apply a controlled-" #NAME " between the qubits.\n" \
" kernel.c" #NAME \
"(parameter=angle, controls=[c1, c2], target=target)\n") \
.def( \
"c" #NAME, \
[](kernel_builder<> &self, double &parameter, \
std::vector<QuakeValue> &controls, \
QuakeValue &target) { self.NAME(parameter, controls, target); }, \
py::arg("parameter"), py::arg("controls"), py::arg("target"), \
"Apply a controlled-" #NAME " operation" \
" to the given target qubit, with the provided list of control " \
"qubits.\n" \
"\nArgs:\n" \
" parameter (float): The float value to " \
"parameterize the " #NAME " gate over.\n" \
" controls (list[QuakeValue]): The control qubits for the " \
"operation.\n" \
" target (:class:`QuakeValue`): The target qubit of the " \
"operation.\n" \
"\n.. code-block:: python\n\n" \
" # Example:\n" \
" kernel = cudaq.make_kernel()\n" \
" c1 = kernel.qalloc()\n" \
" c2 = kernel.qalloc()\n" \
" target = kernel.qalloc()\n" \
" # Apply a controlled-" #NAME " between the qubits.\n" \
" kernel.c" #NAME \
"(parameter=3.14, controls=[c1, c2], target=target)\n")

#define ADD_BUILDER_PARAM_TWO_QUBIT_LIB_GATE(NAME, CUDAQ_FUNC) \
.def( \
Expand Down
Loading

0 comments on commit 339fefe

Please sign in to comment.