Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

New launch kernel for python #2706

Draft
wants to merge 102 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
102 commits
Select commit Hold shift + click to select a range
ac01dd1
DCO Remediation Commit for Ben Howe <[email protected]>
bmhowe23 Oct 11, 2024
21a87c1
State pointer synthesis for quantum hardware
annagrin Sep 17, 2024
3fc56de
Merge with main
annagrin Oct 17, 2024
7969a75
Merge with main
annagrin Oct 17, 2024
755d0d1
Fix test failure on anyon platform
annagrin Oct 17, 2024
dc5e77e
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Oct 17, 2024
382bc99
Make StateInitialization a funcOp pass
annagrin Oct 17, 2024
d3a05d4
Fix issues and tests for the rest of quantum architectures
annagrin Oct 18, 2024
ac151f2
Merge with main
annagrin Oct 18, 2024
51ef054
Fix failing quantinuum state prep tests
annagrin Oct 18, 2024
0cdf3e9
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Oct 18, 2024
5307aa4
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Oct 21, 2024
a7f5387
Address CR comments
annagrin Oct 21, 2024
eb8db13
Merge with main
annagrin Oct 21, 2024
9f0937f
Format
annagrin Oct 21, 2024
2f3a623
Fix failing test
annagrin Oct 22, 2024
b381350
Format
annagrin Oct 22, 2024
dc87ca4
Format
annagrin Oct 22, 2024
e4c7735
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Oct 22, 2024
53a34c9
Replaced getState intrinsic by cc.get_state op
annagrin Oct 22, 2024
30777f3
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Oct 22, 2024
fe6d409
Remove print
annagrin Oct 22, 2024
48704e3
Remove getCudaqState references
annagrin Oct 22, 2024
137f621
Minor updates
annagrin Oct 22, 2024
ad7c6bc
Fix failing quake test
annagrin Oct 23, 2024
83683f7
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Nov 4, 2024
78c0a44
Add a few state-related cc ops
annagrin Nov 5, 2024
6682c39
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Nov 5, 2024
102f819
Fix test_argument_conversion
annagrin Nov 5, 2024
6b2c015
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Nov 5, 2024
5ea1d97
Add printing in failing tests
annagrin Nov 5, 2024
074c60f
Add printing in failing tests
annagrin Nov 5, 2024
310f6ca
Fix failing tests
annagrin Nov 12, 2024
f0176ae
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Nov 12, 2024
d17fa6d
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Nov 12, 2024
3425182
Merge with state-ops
annagrin Nov 12, 2024
6fdccba
Add description for new algorithm for state syntesis
annagrin Nov 12, 2024
fc5e154
Merge with main
annagrin Jan 9, 2025
1dfa805
Fix tests
annagrin Jan 9, 2025
b67fc88
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Jan 9, 2025
9563371
Make intermediate IR legal by separating allocs
annagrin Jan 21, 2025
f32b066
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Jan 21, 2025
008e8c1
DCO Remediation Commit for Anna Gringauze <[email protected]>
annagrin Jan 21, 2025
84a4369
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Jan 22, 2025
1c0a4b3
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 11, 2025
f8e35eb
Address some PR comments
annagrin Feb 12, 2025
e79ad6a
Address more CR comments
annagrin Feb 12, 2025
88cd5d5
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 12, 2025
c0d9ae9
Cleanup
annagrin Feb 13, 2025
1ecd8cc
Address CR comments
annagrin Feb 13, 2025
0238a66
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 14, 2025
de387fc
Address more CR comments
annagrin Feb 14, 2025
a5150a5
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 14, 2025
7cf306a
Address more CR comments
annagrin Feb 18, 2025
16de803
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 18, 2025
1402471
Store new functions in subst module and update synthesis
annagrin Feb 20, 2025
816fdca
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 20, 2025
9a528dd
Make argument synthesis transitive
annagrin Feb 20, 2025
af2fd79
Update callers of synthesis
annagrin Feb 20, 2025
a25d468
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Feb 20, 2025
4d6f7ee
Use PointerOf in quake defs
annagrin Feb 20, 2025
e7d95d5
Addressed more CR comments
annagrin Feb 20, 2025
c4d600f
Recursive with caching
annagrin Feb 21, 2025
e58f577
StateAggregatorWithArgumentConverter
annagrin Feb 21, 2025
6c0dd7d
Make ArgumentConverter handle the state call tree
annagrin Feb 25, 2025
472ff60
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into ar…
annagrin Mar 3, 2025
5ef4c3d
Make argument converter handle kernels created from states
annagrin Mar 3, 2025
ab28edb
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Mar 3, 2025
5e43057
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into ar…
annagrin Mar 3, 2025
36d9425
Merge branch 'argument-converter-all-functions' of https://github.com…
annagrin Mar 3, 2025
7c76533
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Mar 4, 2025
2026cb3
Temp
annagrin Mar 6, 2025
9f70371
Temp
annagrin Mar 6, 2025
9a84fd1
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Mar 6, 2025
7977a2e
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Mar 6, 2025
43b5f85
Temp
annagrin Mar 6, 2025
b0e1c1b
Temp
annagrin Mar 7, 2025
a3cfd75
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Mar 10, 2025
efdbbcc
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Mar 10, 2025
6fefc27
Fix null alloc size and add tests
annagrin Mar 10, 2025
dcf565e
Fix const prop complex and numQubits bugs
annagrin Mar 11, 2025
53b7f5a
Add tests
annagrin Mar 12, 2025
ba5c77e
Format
annagrin Mar 12, 2025
e003e9c
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Mar 12, 2025
3e5ee3b
Fix state copy in python
annagrin Mar 12, 2025
fb5ad4b
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Mar 12, 2025
8590dfa
Remove temp files
annagrin Mar 12, 2025
79e8fdb
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into st…
annagrin Mar 12, 2025
7e39116
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Mar 12, 2025
55e25dd
Keep storing ops when generating numSubits func
annagrin Mar 12, 2025
745fc5d
Cleanup
annagrin Mar 12, 2025
4322ee4
Merge with main
annagrin Mar 12, 2025
b8c8c8a
Cleanup
annagrin Mar 12, 2025
98c19af
Merge with main
annagrin Mar 14, 2025
edf0224
Cleanup
annagrin Mar 14, 2025
6b5161a
Fix failing test
annagrin Mar 14, 2025
054f4b7
Merge branch 'main' of https://github.com/NVIDIA/cuda-quantum into qu…
annagrin Mar 14, 2025
4df4390
Fix failing doc build
annagrin Mar 14, 2025
95bf223
Merge with quantum-state-synthesis
annagrin Mar 14, 2025
7a2e0c4
Allow old launch kernel for tests
annagrin Mar 15, 2025
dff43a0
Cleanup
annagrin Mar 15, 2025
91af463
Cleanup
annagrin Mar 15, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/sphinx/api/languages/cpp_api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,8 @@ Common

.. doxygenclass:: cudaq::RemoteSimulationState

.. doxygenclass:: cudaq::QPUState

.. doxygenclass:: cudaq::registry::RegisteredType
:members:

Expand Down
8 changes: 4 additions & 4 deletions docs/sphinx/targets/cpp/pasqal.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@
// nvq++ --target pasqal pasqal.cpp -o out.x
// ./out.x
// ```
// Assumes a valid set of credentials (PASQAL_AUTH_TOKEN, PASQAL_PROJECT_ID)
// have been set. To set PASQAL_AUTH_TOKEN from Pasqal Cloud username and
// password, use pasqal_auth.py in this folder.
// Assumes a valid set of credentials (`PASQAL_AUTH_TOKEN`, `PASQAL_PROJECT_ID`)
// have been set. To set `PASQAL_AUTH_TOKEN` from Pasqal Cloud username and
// password, use `pasqal_auth.py` in this folder.

#include "cudaq/algorithms/evolve.h"
#include "cudaq/dynamics_integrators.h"
Expand All @@ -15,7 +15,7 @@
#include <map>
#include <vector>

// This example illustrates how to use Pasqal's EMU_MPS emulator over Pasqal's
// This example illustrates how to use Pasqal EMU_MPS emulator over Pasqal
// cloud via CUDA-Q. Contact Pasqal at [email protected] or through
// https://community.pasqal.com for assistance.

Expand Down
2 changes: 1 addition & 1 deletion docs/sphinx/targets/cpp/quera_basic.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
// Credentials must be set before running this program.
// Amazon Braket costs apply.

// This example illustrates how to use QuEra's Aquila device on Braket with
// This example illustrates how to use QuEra Aquila device on Braket with
// CUDA-Q. It is a CUDA-Q implementation of the getting started materials for
// Braket available here:
// https://docs.aws.amazon.com/braket/latest/developerguide/braket-get-started-hello-ahs.html
Expand Down
2 changes: 1 addition & 1 deletion docs/sphinx/targets/cpp/quera_intro.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
#include <map>
#include <vector>

// This example illustrates how to use QuEra's Aquila device on Braket with
// This example illustrates how to use QuEra Aquila device on Braket with
// CUDA-Q. It is a CUDA-Q implementation of the getting started materials for
// Braket available here:
// https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/analog_hamiltonian_simulation/01_Introduction_to_Aquila.ipynb
Expand Down
2 changes: 1 addition & 1 deletion docs/sphinx/targets/python/quera_intro.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from cudaq.operator import *
import numpy as np

# This example illustrates how to use QuEra's Aquila device on Braket with CUDA-Q.
# This example illustrates how to use QuEra Aquila device on Braket with CUDA-Q.
# It is a CUDA-Q implementation of the getting started materials for Braket available here:
# https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/analog_hamiltonian_simulation/01_Introduction_to_Aquila.ipynb

Expand Down
35 changes: 24 additions & 11 deletions include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -1595,7 +1595,7 @@ def quake_CreateStateOp : QuakeOp<"create_state", [Pure]> {
cc_PointerType:$data,
AnySignlessInteger:$length
);
let results = (outs cc_PointerType:$result);
let results = (outs PointerOf<[cc_StateType]>:$result);
let assemblyFormat = [{
$data `,` $length `:` functional-type(operands, results) attr-dict
}];
Expand All @@ -1613,7 +1613,7 @@ def QuakeOp_DeleteStateOp : QuakeOp<"delete_state", [] > {
```
}];

let arguments = (ins cc_PointerType:$state);
let arguments = (ins PointerOf<[cc_StateType]>:$state);
let results = (outs);
let assemblyFormat = [{
$state `:` type(operands) attr-dict
Expand All @@ -1633,30 +1633,43 @@ def quake_GetNumberOfQubitsOp : QuakeOp<"get_number_of_qubits", [Pure] > {
```
}];

let arguments = (ins cc_PointerType:$state);
let arguments = (ins PointerOf<[cc_StateType]>:$state);
let results = (outs AnySignlessInteger:$result);
let assemblyFormat = [{
$state `:` functional-type(operands, results) attr-dict
}];
}

def QuakeOp_GetStateOp : QuakeOp<"get_state", [Pure] > {
def QuakeOp_MaterializeStateOp : QuakeOp<"materialize_state", [Pure] > {
let summary = "Get state from kernel with the provided name.";
let description = [{
This operation is created by argument synthesis of state pointer arguments
for quantum devices. It takes a kernel name as ASCIIZ string literal value
and returns the kernel's quantum state. The operation is replaced by a call
to the kernel with the provided name in ReplaceStateByKernel pass.
for quantum devices.

It takes two kernel names as symbol references:
- @num_qubits for determining the size of the allocation to initialize
- @init for initializing the state the same way as the original kernel
passed to `cudaq::get_state`.

This operation will return of the original kernel passed to
`cudaq::get_state`. `cudaq::get_state`.

The operation may be replaced by calls to the @num_qubits and @init calls,
which will reproduce the specified state in the `ReplaceStateByKernel`
pass.

```mlir
%0 = quake.get_state "callee" : !cc.ptr<!cc.state>
%0 = quake.materialize_state @num_qubits @init : !cc.ptr<!cc.state>
```
}];

let arguments = (ins StrAttr:$calleeName);
let results = (outs cc_PointerType:$result);
let arguments = (ins
FlatSymbolRefAttr:$numQubitsFunc,
FlatSymbolRefAttr:$initFunc
);
let results = (outs PointerOf<[cc_StateType]>:$result);
let assemblyFormat = [{
$calleeName `:` qualified(type(results)) attr-dict
$numQubitsFunc $initFunc `:` qualified(type(results)) attr-dict
}];
}

Expand Down
80 changes: 75 additions & 5 deletions include/cudaq/Optimizer/Transforms/Passes.td
Original file line number Diff line number Diff line change
Expand Up @@ -814,8 +814,8 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
%c8_i64 = arith.constant 8 : i64
%0 = cc.address_of @foo.rodata_synth_0 : !cc.ptr<!cc.array<complex<f32> x 8>>
%4 = cc.create_state %3, %c8_i64 : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
%5 = cc.get_number_of_qubits %4 : (!cc.ptr<!cc.state>) -> i64
%4 = quake.create_state %3, %c8_i64 : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
%5 = quake.get_number_of_qubits %4 : (!cc.ptr<!cc.state>) -> i64
%6 = quake.alloca !quake.veq<?>[%5 : i64]
%7 = quake.init_state %6, %4 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>

Expand All @@ -841,7 +841,7 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
Before DeleteStates (delete-states):
``` func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
...
%4 = call @__nvqpp_cudaq_state_createFromData_fp32(%3, %c8_i64) : (!cc.ptr<i8>, i64) -> !cc.ptr<!cc.state>
%4 = quake.create_state %3, %c8_i64 : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
call @__nvqpp__mlirgen__sub_kernel(%4) : (!cc.ptr<!cc.state>) -> ()
return
}
Expand All @@ -851,9 +851,79 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
```
func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
...
%4 = call @__nvqpp_cudaq_state_createFromData_fp32(%3, %c8_i64) : (!cc.ptr<i8>, i64) -> !cc.ptr<!cc.state>
%4 = quake.create_state %3, %c8_i64 : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
call @__nvqpp__mlirgen__sub_kernel(%4) : (!cc.ptr<!cc.state>) -> ()
call @__nvqpp_cudaq_state_delete(%4) : (!cc.ptr<!cc.state>) -> ()
quake.delete_state %3 : !cc.ptr<!cc.state>
return
}
```
}];
}

def ReplaceStateWithKernel : Pass<"replace-state-with-kernel", "mlir::func::FuncOp"> {
let summary =
"Replace `quake.init_state` instructions with call to the kernel generating the state";
let description = [{
This optimization replaces `quake.init_state`, `quake.get_number_of_qubits`,
and `quake.materialize_state` operations invoked on state pointers during
argument synthesis for quantum devices.

Before this optimization, argument synthesis for state pointers for quantum
devices substituted a state created from the `quake.materialize_state`
operation for the state argument.

The `quake.materialize_state` operation accepts symbols for the synthesized
kernels `@num_qubits` and `@init` that argument synthesis generated from
the original kernel call that generated the state, e.g.,
the `cudaq::get_state` call that refers to the result of a specific quantum
kernel being invoked with a set of parameters

For example, for the user code:
```
state = cudaq::get_state(callee, args)
caller(state)
```

The argument synthesis generated the following new kernels from the `callee`
and synthesized them to substitute their arguments with `args`:
```
func.func @callee_num_qubits(arguments) -> i64
func.func @callee_init(qubits: !quake.veq<?>, arguments) -> !quake.veq<?>
```

The argument synthesis also substituted the state argument in the `caller`
with:
```
quake.materialize_state @callee_num_qubits @callee_init: !cc.ptr<!cc.state>
```

This optimization performs the replacements for the the following operations
that use a state produced by `quake.materialize_state @num_qubits @init`
operation:

- Replace `quake.get_number_of_qubits` operation by call to `@num_qubits`
- Replace `quake.init_state` operation by call to `@init`
- Clean up unused `quake.materialize_state` operation

For example:

Before ReplaceStateWithKernel (replace-state-with-kernel):
```
func.func @foo() {
%0 = quake.materialize_state @callee.num_qubits_0 @callee.init_0: !cc.ptr<!cc.state>
%1 = quake.get_number_of_qubits %0 : (!cc.ptr<!cc.state>) -> i64
%2 = quake.alloca !quake.veq<?>[%1 : i64]
%3 = quake.init_state %2, %0 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
return
}
```

After ReplaceStateWithKernel (replace-state-with-kernel):
```
func.func @foo() {
%1 = call @callee.num_qubits_0() : () -> i64
%2 = quake.alloca !quake.veq<?>[%1 : i64]
%3 = call @callee.init_0(%2) : (!quake.veq<?>) -> !quake.veq<?>
return
}
```
Expand Down
8 changes: 8 additions & 0 deletions lib/Optimizer/Dialect/CC/CCOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -980,6 +980,14 @@ struct FuseWithConstantArray

return success();
}
if (auto complexTy = dyn_cast<ComplexType>(extval.getType())) {
std::int32_t i = extval.getRawConstantIndices()[0];
auto cval = cast<ArrayAttr>(conarr.getConstantValues()[i]);
rewriter.replaceOpWithNewOp<complex::ConstantOp>(extval, complexTy,
cval);

return success();
}
}
return failure();
}
Expand Down
1 change: 1 addition & 0 deletions lib/Optimizer/Transforms/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ add_cudaq_library(OptTransforms
QuakeSynthesizer.cpp
RefToVeqAlloc.cpp
RegToMem.cpp
ReplaceStateWithKernel.cpp
SROA.cpp
StatePreparation.cpp
UnitarySynthesis.cpp
Expand Down
129 changes: 129 additions & 0 deletions lib/Optimizer/Transforms/ReplaceStateWithKernel.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
/*******************************************************************************
* Copyright (c) 2022 - 2025 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/Intrinsics.h"
#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
#include "cudaq/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"

namespace cudaq::opt {
#define GEN_PASS_DEF_REPLACESTATEWITHKERNEL
#include "cudaq/Optimizer/Transforms/Passes.h.inc"
} // namespace cudaq::opt

#define DEBUG_TYPE "replace-state-with-kernel"

using namespace mlir;

namespace {
// clang-format off
/// Replace `quake.get_number_of_qubits` by a call to a function
/// that computes the number of qubits for a state.
///
/// ```mlir
/// %0 = quake.materialize_state @callee.num_qubits_0 @callee.init_0 : !cc.ptr<!cc.state>
/// %1 = quake.get_number_of_qubits %0 : (!cc.ptr<!cc.state>) -> i64
/// ───────────────────────────────────────────
/// %1 = call @callee.num_qubits_0() : () -> i64
/// ```
// clang-format on
class ReplaceGetNumQubitsPattern
: public OpRewritePattern<quake::GetNumberOfQubitsOp> {
public:
using OpRewritePattern::OpRewritePattern;

LogicalResult matchAndRewrite(quake::GetNumberOfQubitsOp numQubits,
PatternRewriter &rewriter) const override {

auto stateOp = numQubits.getOperand();
auto materializeState = stateOp.getDefiningOp<quake::MaterializeStateOp>();
if (!materializeState)
return numQubits->emitError(
"ReplaceStateWithKernel: failed to replace `quake.get_num_qubits`");

auto numQubitsFunc = materializeState.getNumQubitsFunc();
rewriter.setInsertionPoint(numQubits);
rewriter.replaceOpWithNewOp<func::CallOp>(
numQubits, numQubits.getType(), numQubitsFunc, mlir::ValueRange{});
return success();
}
};

// clang-format off
/// Replace `quake.init_state` by a call to a (modified) kernel that produced
/// the state.
///
/// ```mlir
/// %0 = quake.materialize_state @callee.num_qubits_0 @callee.init_0 : !cc.ptr<!cc.state>
/// %3 = quake.init_state %2, %0 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
/// ───────────────────────────────────────────
/// %3 = call @callee.init_0(%2): (!quake.veq<?>) -> !quake.veq<?>
/// ```
// clang-format on
class ReplaceInitStatePattern
: public OpRewritePattern<quake::InitializeStateOp> {
public:
using OpRewritePattern::OpRewritePattern;

LogicalResult matchAndRewrite(quake::InitializeStateOp initState,
PatternRewriter &rewriter) const override {
auto allocaOp = initState.getOperand(0);
auto stateOp = initState.getOperand(1);

if (auto ptrTy = dyn_cast<cudaq::cc::PointerType>(stateOp.getType())) {
if (isa<cudaq::cc::StateType>(ptrTy.getElementType())) {
auto materializeState =
stateOp.getDefiningOp<quake::MaterializeStateOp>();
if (!materializeState)
return initState->emitError(
"ReplaceStateWithKernel: failed to replace `quake.init_state`");

auto initName = materializeState.getInitFunc();
rewriter.setInsertionPoint(initState);
rewriter.replaceOpWithNewOp<func::CallOp>(initState,
initState.getType(), initName,
mlir::ValueRange{allocaOp});

return success();
}
}
return failure();
}
};

class ReplaceStateWithKernelPass
: public cudaq::opt::impl::ReplaceStateWithKernelBase<
ReplaceStateWithKernelPass> {
public:
using ReplaceStateWithKernelBase::ReplaceStateWithKernelBase;

void runOnOperation() override {
auto *ctx = &getContext();
auto func = getOperation();
RewritePatternSet patterns(ctx);
patterns.insert<ReplaceGetNumQubitsPattern, ReplaceInitStatePattern>(ctx);

LLVM_DEBUG(llvm::dbgs()
<< "Before replace state with kernel: " << func << '\n');

if (failed(applyPatternsAndFoldGreedily(func.getOperation(),
std::move(patterns))))
signalPassFailure();

LLVM_DEBUG(llvm::dbgs()
<< "After replace state with kernel: " << func << '\n');
}
};
} // namespace
Loading
Loading