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

Decomposition updates for quake.control types #2186

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
94 changes: 75 additions & 19 deletions lib/Optimizer/Transforms/DecompositionPatterns.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
#include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Rewrite/FrozenRewritePatternSet.h"

Expand Down Expand Up @@ -40,10 +41,12 @@ inline Value createDivF(Location loc, Value numerator, double denominator,
return rewriter.create<arith::DivFOp>(loc, numerator, denominatorValue);
}

/// @brief Returns true if \p op contains any `ControlType` operands.
inline bool containsControlTypes(quake::OperatorInterface op) {
return llvm::any_of(op.getControls(), [](const Value &v) {
return v.getType().isa<quake::ControlType>();
/// @brief Returns true if \p op contains any `ControlType` operands that are
/// used outside of the block where \p op resides.
inline bool containsControlsUsedOutsideBlock(quake::OperatorInterface op) {
return llvm::any_of(op.getControls(), [op](Value v) {
return v.getType().isa<quake::ControlType>() &&
v.isUsedOutsideOfBlock(op->getBlock());
});
}

Expand Down Expand Up @@ -82,6 +85,37 @@ class QuakeOperatorCreator {
quake::WireType::get(rewriter.getContext()));
}

/// @brief Promote all `quake.control` types in \p controls to `quake.wire`
/// types.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This made me wonder if we don't just want to insist that the linear-ctrl-form pass isn't run first?

It seems like decomposition could be run in pruned control form though.

   %42 = quake.foo [%ctrl] %target : (!quake.control, !quake.target) -> !quake.target

might need to expand to

   %a42 = quake.bar %target : (!quake.target) -> !quake.target
   %b42 = quake.baz [%ctrl] %a42 : (!quake.control, !quake.target) -> !quake.target
   %c42 = quake.qux %b42 : (!quake.target) -> !quake.target
   %d42 = quake.quux [%ctrl] %c42 : (!quake.control, !quake.target) -> !quake.target
   %42 = quake.corge %d42 : (!quake.target) -> !quake.target

Some expansion patterns may require the %ctrl qubit to "magically" transform to a target qubit. In those cases, we'd have to add from_ctrl and to_ctrl "inside" the pattern.

   %a42 = quake.bar %target : (!quake.target) -> !quake.target
   %b42 = quake.baz [%ctrl] %a42 : (!quake.control, !quake.target) -> !quake.target
   %from = quake.from_ctrl %ctrl : (!quake.control) -> !quake.target
   %c42:2 = quake.qux %from, %b42 : (!quake.target, !quake.target) -> !quake.target
   %to = quake.to_ctrl %c42#0 : (!quake.target) -> !quake.control
   // Replace the next use of %ctrl (and all uses dominated by the above to_ctrl) with %to
   %d42 = quake.quux [%to] %c42#1 : (!quake.control, !quake.target) -> !quake.target
   %42 = quake.corge %d42 : (!quake.target) -> !quake.target

I guess the real question though is whether this effort is worth it or whether we should just run linear-ctrl-form.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add: this comment

   // Replace the next use of %ctrl (and all uses dominated by the above to_ctrl) with %to

would mean that a value of type !quake.control that is already passed via a terminator to a block argument would require no additional changes.

The sticky case is if the value is defined in another basic block and used in other basic blocks. In that case the from_ctrl and to_ctrl would add a phi node to the SSA graph. (I suspect that may need more work anyway, since we make a lot of assumptions about the code being a DAG, etc.)

Copy link
Collaborator Author

@bmhowe23 bmhowe23 Sep 4, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This made me wonder if we don't just want to insist that the linear-ctrl-form pass isn't run first?

It seems like decomposition could be run in pruned control form though.

This new promoteControls() is only run for specific decompositions that require it. Most decomposition passes do not require this. Requiring linear-ctrl-form be run before all decomposition passes, while functional, may be a bigger hammer than is currently required.

That being said - I'm ok with that ... I was just trying to make these decomposition passes support more flavors of the IR.

It seems like decomposition could be run in pruned control form though.

I agree. Note that the version in main already supports pruned control form for all except 4 decompositions, and these changes resolve those last 4, except for out-of-block use cases.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some expansion patterns may require the %ctrl qubit to "magically" transform to a target qubit. In those cases, we'd have to add from_ctrl and to_ctrl "inside" the pattern.

Yeah, this PR does that. Are you saying that you don't think this PR is doing that, or am I misinterpreting something?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My comment above is descriptive of and in context of the general problem and not within the confines of what this PR does (or does not) do.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This made me wonder if we don't just want to insist that the linear-ctrl-form pass isn't run first?

Just catching up on old reviews and I saw this one was open. Maybe I'm misinterpreting this, but this comment seems contrary to the spirit of #2145 (comment) because that comment seems to suggest that passes should handle all flavors of the IR while the one above (insist that the linear-ctrl-form pass is run first) makes it seem like it is ok to just assume that other passes will be run prior to this pass (and potentially after this pass to restore the control values). In fact, that other thread is the reason I made this PR, so I'm not sure what the next steps for this PR should be. Should we just change this pass back to not updating things when it sees quake.control types?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I see where this might be confusing. Let me see if I can't explain a bit more clearly.

  1. We don't want to reimplement the wheel in every pass. If there is a pass that does a particular transformation, use that pass instead of writing some new patterns that do almost the same thing. The reasons for this are: avoid duplication of effort, avoid introducing bugs, avoid long-term maintenance problems. For example, we really don't want a general folding pass, a pass that has custom folding of integer adds and multiplies as part of modifying certain calls, and a pass that custom folds calls to the math library as a side-effect of rescheduling blocks. The latter two passes have a lack of focus and instead of tackling the job they should, they are building a fiefdom of custom patterns that are unexpected and will likely, from experience, cause problems down the road.

  2. We want passes to tackle whatever input they are given without throwing errors in general. So, if we have a pass to rewrite high-level control flow to a primitive CFG, for example, the pass should run cleanly even if the IR is already represented in a primitive CFG. It just does nothing and the next pass can be applied. Or, if a new high-level construct was just added, the pass should just do what it can to convert the high-level structures it knows about and leave the new one there, which would indicate that someone needs to add the new op to the pass.

So to the point here, the pass could be written to be ignorant of the !quake.control type entirely. If it sees that type, it would just leave those ops alone and transform the ones it does know something about. Since there already is a general pass that eliminates !quake.control typed values, that pass might be used (since that is the focus and point of its existence) to perform that part of the work, while this pass has a narrower focus of doing decompositions in a reference/value semantics agnostic way. (i.e., only knowing about !quake.ref, !quake.wire.)

Hope that helps.

void promoteControls(Location loc, MutableArrayRef<Value> controls) {
origControls.assign(controls.begin(), controls.end());
for (auto &c : controls)
if (c.getType().isa<quake::ControlType>())
c = rewriter.create<quake::FromControlOp>(
loc, quake::WireType::get(rewriter.getContext()), c);
}

/// @brief Perform necessary conversion of `quake.wire` values to
/// `quake.control` values (complementing `promoteControls`). This also
/// replaces downstream uses of the original control with the new control in
/// case the decomposition modified the control.
void demoteWiresToControlsAndReplaceUses(Operation *op,
MutableArrayRef<Value> controls) {
auto ctrlTy = quake::ControlType::get(rewriter.getContext());
auto loc = op->getLoc();
DominanceInfo domInfo(op->getParentOfType<func::FuncOp>());
for (auto &&[oc, c] : llvm::zip_equal(origControls, controls))
if (oc.getType().isa<quake::ControlType>()) {
c = rewriter.create<quake::ToControlOp>(loc, ctrlTy, c);
// This is like oc.replaceAllUsesWith(c) except it checks for
// proper dominance before doing the replacement. This is important
// because the rewriter tends to work from the bottom up.
for (auto &use : llvm::make_early_inc_range(oc.getUses()))
if (domInfo.properlyDominates(c, use.getOwner()))
use.set(c);
}
}

/// Pluck out the values from \p newValues whose type is `WireType` and
/// replace all the \p op uses with those values.
void selectWiresAndReplaceUses(Operation *op, ValueRange newValues) {
Expand Down Expand Up @@ -224,6 +258,9 @@ class QuakeOperatorCreator {

private:
PatternRewriter &rewriter;
/// The original control values before some may have been promoted from
/// quake.control to quake.wire.
SmallVector<Value> origControls;
};

/// Check whether the operation has the correct number of controls.
Expand Down Expand Up @@ -661,11 +698,6 @@ struct CXToCZ : public OpRewritePattern<quake::XOp> {
PatternRewriter &rewriter) const override {
if (failed(checkNumControls(op, 1)))
return failure();
// This decomposition does not support `quake.control` types because the
// input controls are used as targets during this transformation.
if (containsControlTypes(op))
return failure();

// Op info
Location loc = op->getLoc();
Value target = op.getTarget();
Expand All @@ -675,7 +707,15 @@ struct CXToCZ : public OpRewritePattern<quake::XOp> {
if (negatedControls)
negControl = (*negatedControls)[0];

// TODO - Update this pattern to support threading modified controls
// throughout multiple blocks. qRewriter.demoteWiresToControlsAndReplaceUses
// does not currently handle that.
if (negControl && containsControlsUsedOutsideBlock(op))
return failure();

QuakeOperatorCreator qRewriter(rewriter);
if (negControl)
qRewriter.promoteControls(loc, controls);
qRewriter.create<quake::HOp>(loc, target);
if (negControl)
qRewriter.create<quake::XOp>(loc, controls);
Expand All @@ -684,6 +724,8 @@ struct CXToCZ : public OpRewritePattern<quake::XOp> {
qRewriter.create<quake::XOp>(loc, controls);
qRewriter.create<quake::HOp>(loc, target);

if (negControl)
qRewriter.demoteWiresToControlsAndReplaceUses(op, controls);
qRewriter.selectWiresAndReplaceUses(op, controls, target);
rewriter.eraseOp(op);
return success();
Expand Down Expand Up @@ -813,11 +855,6 @@ struct CCZToCX : public OpRewritePattern<quake::ZOp> {

LogicalResult matchAndRewrite(quake::ZOp op,
PatternRewriter &rewriter) const override {
// This decomposition does not support `quake.control` types because the
// input controls are used as targets during this transformation.
if (containsControlTypes(op))
return failure();

SmallVector<Value, 2> controls(2);
if (failed(checkAndExtractControls(op, controls, rewriter)))
return failure();
Expand All @@ -843,7 +880,14 @@ struct CCZToCX : public OpRewritePattern<quake::ZOp> {
}
}

// TODO - Update this pattern to support threading modified controls
// throughout multiple blocks. qRewriter.demoteWiresToControlsAndReplaceUses
// does not currently handle that.
if (containsControlsUsedOutsideBlock(op))
return failure();

QuakeOperatorCreator qRewriter(rewriter);
qRewriter.promoteControls(loc, controls);
qRewriter.create<quake::XOp>(loc, controls[1], target);
qRewriter.create<quake::TOp>(loc, /*isAdj=*/!negC0, target);
qRewriter.create<quake::XOp>(loc, controls[0], target);
Expand All @@ -860,6 +904,7 @@ struct CCZToCX : public OpRewritePattern<quake::ZOp> {

qRewriter.create<quake::TOp>(loc, /*isAdj=*/negC1, controls[0]);

qRewriter.demoteWiresToControlsAndReplaceUses(op, controls);
qRewriter.selectWiresAndReplaceUses(op, controls, target);
rewriter.eraseOp(op);
return success();
Expand All @@ -878,10 +923,6 @@ struct CZToCX : public OpRewritePattern<quake::ZOp> {

LogicalResult matchAndRewrite(quake::ZOp op,
PatternRewriter &rewriter) const override {
// This decomposition does not support `quake.control` types because the
// input controls are used as targets during this transformation.
if (containsControlTypes(op))
return failure();
if (failed(checkNumControls(op, 1)))
return failure();

Expand All @@ -894,7 +935,15 @@ struct CZToCX : public OpRewritePattern<quake::ZOp> {
if (negatedControls)
negControl = (*negatedControls)[0];

// TODO - Update this pattern to support threading modified controls
// throughout multiple blocks. qRewriter.demoteWiresToControlsAndReplaceUses
// does not currently handle that.
if (negControl && containsControlsUsedOutsideBlock(op))
return failure();

QuakeOperatorCreator qRewriter(rewriter);
if (negControl)
qRewriter.promoteControls(loc, controls);
qRewriter.create<quake::HOp>(loc, target);
if (negControl)
qRewriter.create<quake::XOp>(loc, controls);
Expand All @@ -903,6 +952,8 @@ struct CZToCX : public OpRewritePattern<quake::ZOp> {
qRewriter.create<quake::XOp>(loc, controls);
qRewriter.create<quake::HOp>(loc, target);

if (negControl)
qRewriter.demoteWiresToControlsAndReplaceUses(op, controls);
qRewriter.selectWiresAndReplaceUses(op, controls, target);
rewriter.eraseOp(op);
return success();
Expand Down Expand Up @@ -969,7 +1020,10 @@ struct CR1ToCX : public OpRewritePattern<quake::R1Op> {

LogicalResult matchAndRewrite(quake::R1Op op,
PatternRewriter &rewriter) const override {
if (containsControlTypes(op))
// TODO - Update this pattern to support threading modified controls
// throughout multiple blocks. qRewriter.demoteWiresToControlsAndReplaceUses
// does not currently handle that.
if (containsControlsUsedOutsideBlock(op))
return failure();

Value control;
Expand All @@ -994,6 +1048,7 @@ struct CR1ToCX : public OpRewritePattern<quake::R1Op> {
Value negHalfAngle = rewriter.create<arith::NegFOp>(loc, halfAngle);

QuakeOperatorCreator qRewriter(rewriter);
qRewriter.promoteControls(loc, control);
qRewriter.create<quake::R1Op>(loc, /*isAdj*/ negControl, halfAngle,
noControls, control);
qRewriter.create<quake::XOp>(loc, control, target);
Expand All @@ -1002,6 +1057,7 @@ struct CR1ToCX : public OpRewritePattern<quake::R1Op> {
qRewriter.create<quake::XOp>(loc, control, target);
qRewriter.create<quake::R1Op>(loc, halfAngle, noControls, target);

qRewriter.demoteWiresToControlsAndReplaceUses(op, control);
qRewriter.selectWiresAndReplaceUses(op, ValueRange{control, target});
rewriter.eraseOp(op);
return success();
Expand Down
1 change: 1 addition & 0 deletions test/Transforms/DecompositionPatterns/CCZToCX.qke
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
// RUN: cudaq-opt -pass-pipeline='builtin.module(decomposition{enable-patterns=CCZToCX})' %s | FileCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(decomposition{enable-patterns=CCZToCX})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CCZToCX})' %s | FileCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg,pruned-ctrl-form),decomposition{enable-patterns=CCZToCX})' %s | FileCheck %s

// Test the decomposition pattern with different control types. The FileCheck
// part of this test only cares about the sequence of operations. Correcteness
Expand Down
1 change: 1 addition & 0 deletions test/Transforms/DecompositionPatterns/CR1ToCX.qke
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
// RUN: cudaq-opt -pass-pipeline='builtin.module(decomposition{enable-patterns=CR1ToCX})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CR1ToCX})' %s | FileCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CR1ToCX})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg,pruned-ctrl-form),decomposition{enable-patterns=CR1ToCX})' %s | FileCheck %s

// Test the decomposition pattern with different control types. The FileCheck
// part of this test only cares about the sequence of operations. Correcteness
Expand Down
1 change: 1 addition & 0 deletions test/Transforms/DecompositionPatterns/CXToCZ.qke
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
// RUN: cudaq-opt -pass-pipeline='builtin.module(decomposition{enable-patterns=CXToCZ})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CXToCZ})' %s | FileCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CXToCZ})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg,pruned-ctrl-form),decomposition{enable-patterns=CXToCZ})' %s | FileCheck %s

// Test the decomposition pattern with different control types. The FileCheck
// part of this test only cares about the sequence of operations. Correcteness
Expand Down
1 change: 1 addition & 0 deletions test/Transforms/DecompositionPatterns/CZToCX.qke
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
// RUN: cudaq-opt -pass-pipeline='builtin.module(decomposition{enable-patterns=CZToCX})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CZToCX})' %s | FileCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg),decomposition{enable-patterns=CZToCX})' %s | CircuitCheck %s
// RUN: cudaq-opt -pass-pipeline='builtin.module(func.func(expand-control-veqs,memtoreg,pruned-ctrl-form),decomposition{enable-patterns=CZToCX})' %s | FileCheck %s

// Test the decomposition pattern with different control types. The FileCheck
// part of this test only cares about the sequence of operations. Correcteness
Expand Down
Loading