Skip to content

Commit d2bed4c

Browse files
annagrinannagrin
authored
Add a few state-related cc ops (#2354)
* Add a few state-related cc ops Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Fix test_argument_conversion Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Add printing in failing tests Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Add printing in failing tests Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Fix failing tests Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Address CR comments Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Fix comment Signed-off-by: Anna Gringauze <agringauze@nvidia.com> * Address CR comments Signed-off-by: Anna Gringauze <agringauze@nvidia.com> --------- Signed-off-by: Anna Gringauze <agringauze@nvidia.com>
1 parent c5c9361 commit d2bed4c

File tree

18 files changed

+323
-211
lines changed

18 files changed

+323
-211
lines changed

‎.github/workflows/config/spelling_allowlist.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ SLURM
108108
SVD
109109
Stim
110110
Superpositions
111+
Superstaq
111112
TBI
112113
TCP
113114
TableGen

‎docs/sphinx/targets/cpp/infleqtion.cpp

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22
// ```
33
// nvq++ --target infleqtion infleqtion.cpp -o out.x && ./out.x
44
// ```
5-
// This will submit the job to the Infleqtion's ideal simulator,
6-
// cq_sqale_simulator (default). Alternatively, we can enable hardware noise
5+
// This will submit the job to the ideal simulator for Infleqtion,
6+
// `cq_sqale_simulator` (default). Alternatively, we can enable hardware noise
77
// model simulation by specifying `noise-sim` to the flag `--infleqtion-method`,
88
// e.g.,
99
// ```
@@ -17,9 +17,9 @@
1717
// nvq++ --target infleqtion --infleqtion-machine cq_sqale_qpu
1818
// --infleqtion-method dry-run infleqtion.cpp -o out.x && ./out.x
1919
// ```
20-
// Note: If targeting ideal cloud simulation, `--infleqtion-machine
21-
// cq_sqale_simulator` is optional since it is the default configuration if not
22-
// provided.
20+
// Note: If targeting ideal cloud simulation,
21+
// `--infleqtion-machine cq_sqale_simulator` is optional since it is the
22+
// default configuration if not provided.
2323

2424
#include <cudaq.h>
2525
#include <fstream>
@@ -38,7 +38,7 @@ struct ghz {
3838
};
3939

4040
int main() {
41-
// Submit to infleqtion asynchronously (e.g., continue executing
41+
// Submit to Infleqtion asynchronously (e.g., continue executing
4242
// code in the file until the job has been returned).
4343
auto future = cudaq::sample_async(ghz{});
4444
// ... classical code to execute in the meantime ...
@@ -58,7 +58,7 @@ int main() {
5858
auto async_counts = readIn.get();
5959
async_counts.dump();
6060

61-
// OR: Submit to infleqtion synchronously (e.g., wait for the job
61+
// OR: Submit to Infleqtion synchronously (e.g., wait for the job
6262
// result to be returned before proceeding).
6363
auto counts = cudaq::sample(ghz{});
6464
counts.dump();

‎include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1397,4 +1397,90 @@ def CustomUnitarySymbolOp :
13971397
}];
13981398
}
13991399

1400+
//===----------------------------------------------------------------------===//
1401+
// Quantum states
1402+
//===----------------------------------------------------------------------===//
1403+
1404+
def quake_CreateStateOp : QuakeOp<"create_state", [Pure]> {
1405+
let summary = "Create state from data";
1406+
let description = [{
1407+
This operation takes a pointer to state data and creates a quantum state,
1408+
where state data is a pointer to an array of float or complex numbers.
1409+
The operation can be optimized away in DeleteStates pass, or replaced
1410+
by an intrinsic runtime call on simulators.
1411+
1412+
```mlir
1413+
%0 = quake.create_state %data %len : (!cc.ptr<!cc.array<complex<f64> x 8>>, i64) -> !cc.ptr<!cc.state>
1414+
```
1415+
}];
1416+
1417+
let arguments = (ins
1418+
cc_PointerType:$data,
1419+
AnySignlessInteger:$length
1420+
);
1421+
let results = (outs cc_PointerType:$result);
1422+
let assemblyFormat = [{
1423+
$data `,` $length `:` functional-type(operands, results) attr-dict
1424+
}];
1425+
}
1426+
1427+
def QuakeOp_DeleteStateOp : QuakeOp<"delete_state", [] > {
1428+
let summary = "Delete quantum state";
1429+
let description = [{
1430+
This operation takes a pointer to the state and deletes the state object.
1431+
The operation can be created in in DeleteStates pass, and replaced later
1432+
by an intrinsic runtime call on simulators.
1433+
1434+
```mlir
1435+
quake.delete_state %state : !cc.ptr<!cc.state>
1436+
```
1437+
}];
1438+
1439+
let arguments = (ins cc_PointerType:$state);
1440+
let results = (outs);
1441+
let assemblyFormat = [{
1442+
$state `:` type(operands) attr-dict
1443+
}];
1444+
}
1445+
1446+
def quake_GetNumberOfQubitsOp : QuakeOp<"get_number_of_qubits", [Pure] > {
1447+
let summary = "Get number of qubits from a quantum state";
1448+
let description = [{
1449+
This operation takes a pointer to the state as an argument and returns
1450+
a number of qubits in the state. The operation can be optimized away in
1451+
some passes like ReplaceStateByKernel or DeleteStates, or replaced by an
1452+
intrinsic runtime call when the target is one of the simulators.
1453+
1454+
```mlir
1455+
%0 = quake.get_number_of_qubits %state : (!cc.ptr<!cc.state>) -> i64
1456+
```
1457+
}];
1458+
1459+
let arguments = (ins cc_PointerType:$state);
1460+
let results = (outs AnySignlessInteger:$result);
1461+
let assemblyFormat = [{
1462+
$state `:` functional-type(operands, results) attr-dict
1463+
}];
1464+
}
1465+
1466+
def QuakeOp_GetStateOp : QuakeOp<"get_state", [Pure] > {
1467+
let summary = "Get state from kernel with the provided name.";
1468+
let description = [{
1469+
This operation is created by argument synthesis of state pointer arguments
1470+
for quantum devices. It takes a kernel name as ASCIIZ string literal value
1471+
and returns the kernel's quantum state. The operation is replaced by a call
1472+
to the kernel with the provided name in ReplaceStateByKernel pass.
1473+
1474+
```mlir
1475+
%0 = quake.get_state "callee" : !cc.ptr<!cc.state>
1476+
```
1477+
}];
1478+
1479+
let arguments = (ins StrAttr:$calleeName);
1480+
let results = (outs cc_PointerType:$result);
1481+
let assemblyFormat = [{
1482+
$calleeName `:` qualified(type(results)) attr-dict
1483+
}];
1484+
}
1485+
14001486
#endif // CUDAQ_OPTIMIZER_DIALECT_QUAKE_OPS

‎include/cudaq/Optimizer/Transforms/Passes.td

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -778,9 +778,8 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
778778
func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
779779
%c8_i64 = arith.constant 8 : i64
780780
%0 = cc.address_of @foo.rodata_synth_0 : !cc.ptr<!cc.array<complex<f32> x 8>>
781-
%3 = cc.cast %0 : (!cc.ptr<!cc.array<complex<f32> x 8>>) -> !cc.ptr<i8>
782-
%4 = call @__nvqpp_cudaq_state_createFromData_fp32(%3, %c8_i64) : (!cc.ptr<i8>, i64) -> !cc.ptr<!cc.state>
783-
%5 = call @__nvqpp_cudaq_state_numberOfQubits(%4) : (!cc.ptr<!cc.state>) -> i64
781+
%4 = cc.create_state %3, %c8_i64 : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
782+
%5 = cc.get_number_of_qubits %4 : (!cc.ptr<!cc.state>) -> i64
784783
%6 = quake.alloca !quake.veq<?>[%5 : i64]
785784
%7 = quake.init_state %6, %4 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
786785

‎lib/Frontend/nvqpp/ConvertExpr.cpp

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2698,19 +2698,12 @@ bool QuakeBridgeVisitor::VisitCXXConstructExpr(clang::CXXConstructExpr *x) {
26982698
initials = load.getPtrvalue();
26992699
}
27002700
if (isStateType(initials.getType())) {
2701-
IRBuilder irBuilder(builder.getContext());
2702-
auto mod =
2703-
builder.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
2704-
auto result =
2705-
irBuilder.loadIntrinsic(mod, getNumQubitsFromCudaqState);
2706-
assert(succeeded(result) && "loading intrinsic should never fail");
27072701
Value state = initials;
27082702
auto i64Ty = builder.getI64Type();
2709-
auto numQubits = builder.create<func::CallOp>(
2710-
loc, i64Ty, getNumQubitsFromCudaqState, ValueRange{state});
2703+
auto numQubits =
2704+
builder.create<quake::GetNumberOfQubitsOp>(loc, i64Ty, state);
27112705
auto veqTy = quake::VeqType::getUnsized(ctx);
2712-
Value alloc = builder.create<quake::AllocaOp>(loc, veqTy,
2713-
numQubits.getResult(0));
2706+
Value alloc = builder.create<quake::AllocaOp>(loc, veqTy, numQubits);
27142707
return pushValue(builder.create<quake::InitializeStateOp>(
27152708
loc, veqTy, alloc, state));
27162709
}

‎lib/Optimizer/CodeGen/QuakeToCodegen.cpp

Lines changed: 89 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88

99
#include "QuakeToCodegen.h"
1010
#include "CodeGenOps.h"
11+
#include "cudaq/Optimizer/Builder/Intrinsics.h"
12+
#include "cudaq/Optimizer/CodeGen/Passes.h"
13+
#include "cudaq/Optimizer/CodeGen/QIRFunctionNames.h"
1114
#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
1215
#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
1316
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
@@ -62,10 +65,95 @@ class ExpandComplexCast : public OpRewritePattern<cudaq::cc::CastOp> {
6265
return success();
6366
}
6467
};
68+
69+
class CreateStateOpPattern : public OpRewritePattern<quake::CreateStateOp> {
70+
public:
71+
using OpRewritePattern::OpRewritePattern;
72+
73+
LogicalResult matchAndRewrite(quake::CreateStateOp createStateOp,
74+
PatternRewriter &rewriter) const override {
75+
auto module = createStateOp->getParentOfType<ModuleOp>();
76+
auto loc = createStateOp.getLoc();
77+
auto ctx = createStateOp.getContext();
78+
auto buffer = createStateOp.getOperand(0);
79+
auto size = createStateOp.getOperand(1);
80+
81+
auto bufferTy = buffer.getType();
82+
auto ptrTy = cast<cudaq::cc::PointerType>(bufferTy);
83+
auto arrTy = cast<cudaq::cc::ArrayType>(ptrTy.getElementType());
84+
auto eleTy = arrTy.getElementType();
85+
auto is64Bit = isa<Float64Type>(eleTy);
86+
87+
if (auto cTy = dyn_cast<ComplexType>(eleTy))
88+
is64Bit = isa<Float64Type>(cTy.getElementType());
89+
90+
auto createStateFunc = is64Bit ? cudaq::createCudaqStateFromDataFP64
91+
: cudaq::createCudaqStateFromDataFP32;
92+
cudaq::IRBuilder irBuilder(ctx);
93+
auto result = irBuilder.loadIntrinsic(module, createStateFunc);
94+
assert(succeeded(result) && "loading intrinsic should never fail");
95+
96+
auto stateTy = cudaq::cc::StateType::get(ctx);
97+
auto statePtrTy = cudaq::cc::PointerType::get(stateTy);
98+
auto i8PtrTy = cudaq::cc::PointerType::get(rewriter.getI8Type());
99+
auto cast = rewriter.create<cudaq::cc::CastOp>(loc, i8PtrTy, buffer);
100+
101+
rewriter.replaceOpWithNewOp<func::CallOp>(
102+
createStateOp, statePtrTy, createStateFunc, ValueRange{cast, size});
103+
return success();
104+
}
105+
};
106+
107+
class DeleteStateOpPattern : public OpRewritePattern<quake::DeleteStateOp> {
108+
public:
109+
using OpRewritePattern::OpRewritePattern;
110+
111+
LogicalResult matchAndRewrite(quake::DeleteStateOp deleteStateOp,
112+
PatternRewriter &rewriter) const override {
113+
auto module = deleteStateOp->getParentOfType<ModuleOp>();
114+
auto ctx = deleteStateOp.getContext();
115+
auto state = deleteStateOp.getOperand();
116+
117+
cudaq::IRBuilder irBuilder(ctx);
118+
auto result = irBuilder.loadIntrinsic(module, cudaq::deleteCudaqState);
119+
assert(succeeded(result) && "loading intrinsic should never fail");
120+
121+
rewriter.replaceOpWithNewOp<func::CallOp>(deleteStateOp, std::nullopt,
122+
cudaq::deleteCudaqState,
123+
mlir::ValueRange{state});
124+
return success();
125+
}
126+
};
127+
128+
class GetNumberOfQubitsOpPattern
129+
: public OpRewritePattern<quake::GetNumberOfQubitsOp> {
130+
public:
131+
using OpRewritePattern::OpRewritePattern;
132+
133+
LogicalResult matchAndRewrite(quake::GetNumberOfQubitsOp getNumQubitsOp,
134+
PatternRewriter &rewriter) const override {
135+
auto module = getNumQubitsOp->getParentOfType<ModuleOp>();
136+
auto ctx = getNumQubitsOp.getContext();
137+
auto state = getNumQubitsOp.getOperand();
138+
139+
cudaq::IRBuilder irBuilder(ctx);
140+
auto result =
141+
irBuilder.loadIntrinsic(module, cudaq::getNumQubitsFromCudaqState);
142+
assert(succeeded(result) && "loading intrinsic should never fail");
143+
144+
rewriter.replaceOpWithNewOp<func::CallOp>(
145+
getNumQubitsOp, rewriter.getI64Type(),
146+
cudaq::getNumQubitsFromCudaqState, state);
147+
return success();
148+
}
149+
};
150+
65151
} // namespace
66152

67153
void cudaq::codegen::populateQuakeToCodegenPatterns(
68154
mlir::RewritePatternSet &patterns) {
69155
auto *ctx = patterns.getContext();
70-
patterns.insert<CodeGenRAIIPattern, ExpandComplexCast>(ctx);
156+
patterns
157+
.insert<CodeGenRAIIPattern, CreateStateOpPattern, DeleteStateOpPattern,
158+
ExpandComplexCast, GetNumberOfQubitsOpPattern>(ctx);
71159
}

0 commit comments

Comments
 (0)