Skip to content

[core] Harden add-dealloc for IR in squirrelly mixed state. #5300

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.

Sign up for GitHub

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

Merged
schweitzpgi merged 1 commit into from
Aug 27, 2026
+124 −35
Merged

[core] Harden add-dealloc for IR in squirrelly mixed state. #5300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
57 changes: 47 additions & 10 deletions cudaq/lib/Optimizer/Transforms/AddDeallocs.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,13 @@ using namespace mlir;
using DeallocationMap = llvm::DenseMap<Operation *, bool>;
using RegionOpSet = llvm::DenseSet<Operation *>;

// Added to harden this pass when the IR is in a squirrelly mix of alloca and
// sink operations. Normally, the IR isn't in that state coming out of a bridge.
static bool isUnwrapImmediatelySunk(cudaq::quake::UnwrapOp unwrap) {
Value wire = unwrap.getResult();
return wire.hasOneUse() && isa<cudaq::quake::SinkOp>(*wire.user_begin());
}

namespace {
struct DeallocationAnalysisInfo {
DeallocationAnalysisInfo() = default;
Expand Down Expand Up @@ -61,6 +68,21 @@ struct DeallocationAnalysisInfo {
deallocMap[alloc] = true;
else
deallocMap.insert(std::make_pair(alloc, true));
} else if (auto unwrap = dyn_cast<cudaq::quake::UnwrapOp>(op);
unwrap && isUnwrapImmediatelySunk(unwrap)) {
auto val = unwrap.getRefValue();
Operation *alloc = val.getDefiningOp();
if (alloc && !isa<cudaq::quake::AllocaOp>(alloc)) {
auto initState = dyn_cast<cudaq::quake::InitializeStateOp>(alloc);
alloc =
initState ? initState.getTargets().getDefiningOp() : nullptr;
}
if (alloc) {
if (deallocMap.count(alloc))
deallocMap[alloc] = true;
else
deallocMap.insert(std::make_pair(alloc, true));
}
}
}
for (auto &[_, dealloced] : deallocMap)
Expand Down Expand Up @@ -89,6 +111,20 @@ class DeallocationAnalysis {
}

private:
// Record that the alloca defining \p val (following through an
// InitializeStateOp, if present) has already been deallocated -- whether
// by an explicit quake.dealloc or by the value-semantics equivalent.
bool markDeallocated(Value val) {
if (auto init = val.getDefiningOp<cudaq::quake::InitializeStateOp>())
val = init.getTargets();
auto alloc = val.getDefiningOp<cudaq::quake::AllocaOp>();
if (!alloc)
return false;
auto *op = alloc.getOperation();
allocMap[op] = true;
return true;
}

// Perform the analysis on \p func.
void performAnalysis(Operation *func) {
func->walk([this](Operation *o) {
Expand All @@ -105,20 +141,21 @@ class DeallocationAnalysis {
<< op->getParentOp() << '\n');
}
} else if (auto dealloc = dyn_cast<cudaq::quake::DeallocOp>(o)) {
auto val = dealloc.getReference();
if (auto init = val.getDefiningOp<cudaq::quake::InitializeStateOp>())
val = init.getTargets();
if (auto alloc = val.getDefiningOp<cudaq::quake::AllocaOp>()) {
auto *op = alloc.getOperation();
if (allocMap.count(op))
allocMap[op] = true;
else
allocMap.insert(std::make_pair(op, /*deallocated=*/true));
LLVM_DEBUG(llvm::dbgs() << "found dealloc of alloca: " << op << '\n');
if (markDeallocated(dealloc.getReference())) {
LLVM_DEBUG(llvm::dbgs()
<< "found dealloc of alloca: "
<< dealloc.getReference().getDefiningOp() << '\n');
} else {
dealloc->emitWarning("unable to determine associated allocation.");
hasErrors = true;
}
} else if (auto unwrap = dyn_cast<cudaq::quake::UnwrapOp>(o)) {
if (isUnwrapImmediatelySunk(unwrap) &&
markDeallocated(unwrap.getRefValue())) {
LLVM_DEBUG(llvm::dbgs()
<< "found unwrap+sink of alloca: "
<< unwrap.getRefValue().getDefiningOp() << '\n');
}
}
});
}
Expand Down
102 changes: 77 additions & 25 deletions cudaq/test/Transforms/add_dealloc-1.qke
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,38 +8,90 @@

// RUN: cudaq-opt --add-dealloc %s | FileCheck %s

module {
func.func @__nvqpp__mlirgen__FromState() {
%c4 = arith.constant 4 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = cc.address_of @__nvqpp__rodata_init_0 : !cc.ptr<!cc.array<f64 x 4>>
%1 = quake.alloca !quake.veq<4>
%2 = quake.init_state %1, %0 : (!quake.veq<4>, !cc.ptr<!cc.array<f64 x 4>>) -> !quake.veq<4>
%3 = cc.loop while ((%arg0 = %c0) -> (index)) {
%4 = arith.cmpi slt, %arg0, %c4 : index
cc.condition %4(%arg0 : index)
} do {
^bb0(%arg0: index):
%4 = quake.extract_ref %2[%arg0] : (!quake.veq<4>, index) -> !quake.ref
quake.h %4 : (!quake.ref) -> ()
cc.continue %arg0 : index
} step {
^bb0(%arg0: index):
%4 = arith.addi %arg0, %c1 : index
cc.continue %4 : index
} {invariant}
return
}
cc.global constant private @__nvqpp__rodata_init_0 (dense<[1.000000e+00, 0.000000e+00, 5.000000e-01, 5.000000e-01]> : tensor<4xf64>) : !cc.array<f64 x 4>
func.func @__nvqpp__mlirgen__FromState() {
%c4 = arith.constant 4 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = cc.address_of @__nvqpp__rodata_init_0 : !cc.ptr<!cc.array<f64 x 4>>
%1 = quake.alloca !quake.veq<4>
%2 = quake.init_state %1, %0 : (!quake.veq<4>, !cc.ptr<!cc.array<f64 x 4>>) -> !quake.veq<4>
%3 = cc.loop while ((%arg0 = %c0) -> (index)) {
%4 = arith.cmpi slt, %arg0, %c4 : index
cc.condition %4(%arg0 : index)
} do {
^bb0(%arg0: index):
%4 = quake.extract_ref %2[%arg0] : (!quake.veq<4>, index) -> !quake.ref
quake.h %4 : (!quake.ref) -> ()
cc.continue %arg0 : index
} step {
^bb0(%arg0: index):
%4 = arith.addi %arg0, %c1 : index
cc.continue %4 : index
} {invariant}
return
}

cc.global constant private @__nvqpp__rodata_init_0 (dense<[1.000000e+00, 0.000000e+00, 5.000000e-01, 5.000000e-01]> : tensor<4xf64>) : !cc.array<f64 x 4>

// CHECK-LABEL: func.func @__nvqpp__mlirgen__FromState() {
// CHECK: %[[VAL_3:.*]] = cc.address_of @__nvqpp__rodata_init_0 : !cc.ptr<!cc.array<f64 x 4>>
// CHECK: %[[VAL_4:.*]] = quake.alloca !quake.veq<4>
// CHECK: %[[VAL_5:.*]] = quake.init_state %[[VAL_4]], %[[VAL_3]] : (!quake.veq<4>, !cc.ptr<!cc.array<f64 x 4>>) -> !quake.veq<4>
// CHECK: quake.dealloc %[[VAL_5]] : !quake.veq<4>
// CHECK: return
// CHECK: }
// CHECK: cc.global constant private @__nvqpp__rodata_init_0 (dense<[1.000000e+00, 0.000000e+00, 5.000000e-01, 5.000000e-01]> : tensor<4xf64>) : !cc.array<f64 x 4>
// CHECK: cc.global constant private @__nvqpp__rodata_init_0


// Issue 5296
func.func @already_sunk() {
%0 = quake.alloca !quake.ref
%1 = quake.alloca !quake.ref
%2 = quake.concat %0, %1 : (!quake.ref, !quake.ref) -> !quake.veq<2>
%r = quake.extract_ref %2[0] : (!quake.veq<2>) -> !quake.ref
quake.h %r : (!quake.ref) -> ()
%3 = quake.unwrap %0 : (!quake.ref) -> !quake.wire
quake.sink %3 : !quake.wire
%4 = quake.unwrap %1 : (!quake.ref) -> !quake.wire
quake.sink %4 : !quake.wire
return
}

// CHECK-LABEL: func.func @already_sunk() {
// CHECK: %[[VAL_0:.*]] = quake.alloca !quake.ref
// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.ref
// CHECK: %[[VAL_2:.*]] = quake.concat %[[VAL_0]], %[[VAL_1]] : (!quake.ref, !quake.ref) -> !quake.veq<2>
// CHECK: %[[VAL_6:.*]] = quake.extract_ref %[[VAL_2]][0] : (!quake.veq<2>) -> !quake.ref
// CHECK: quake.h %[[VAL_6]] : (!quake.ref) -> ()
// CHECK: %[[VAL_7:.*]] = quake.unwrap %[[VAL_0]] : (!quake.ref) -> !quake.wire
// CHECK: quake.sink %[[VAL_7]] : !quake.wire
// CHECK: %[[VAL_8:.*]] = quake.unwrap %[[VAL_1]] : (!quake.ref) -> !quake.wire
// CHECK: quake.sink %[[VAL_8]] : !quake.wire
// CHECK-NOT: quake.dealloc
// CHECK: return
// CHECK: }

// Issue 5296
func.func @one_sunk_one_not() {
%0 = quake.alloca !quake.ref
%1 = quake.alloca !quake.ref
quake.h %0 : (!quake.ref) -> ()
%2 = quake.unwrap %0 : (!quake.ref) -> !quake.wire
quake.sink %2 : !quake.wire
quake.h %1 : (!quake.ref) -> ()
return
}

// CHECK-LABEL: func.func @one_sunk_one_not() {
// CHECK: %[[VAL_9:.*]] = quake.alloca !quake.ref
// CHECK: %[[VAL_10:.*]] = quake.alloca !quake.ref
// CHECK: quake.h %[[VAL_9]] : (!quake.ref) -> ()
// CHECK: %[[VAL_11:.*]] = quake.unwrap %[[VAL_9]] : (!quake.ref) -> !quake.wire
// CHECK: quake.sink %[[VAL_11]] : !quake.wire
// CHECK: quake.h %[[VAL_10]] : (!quake.ref) -> ()
// CHECK: quake.dealloc %[[VAL_10]] : !quake.ref
// CHECK: return
// CHECK: }



Loading