diff --git a/mlir/lib/Dialect/QC/IR/QubitManagement/DeallocOp.cpp b/mlir/lib/Dialect/QC/IR/QubitManagement/DeallocOp.cpp index 26b6f7b6f6..b8559de464 100644 --- a/mlir/lib/Dialect/QC/IR/QubitManagement/DeallocOp.cpp +++ b/mlir/lib/Dialect/QC/IR/QubitManagement/DeallocOp.cpp @@ -53,7 +53,9 @@ struct HoistStaticQubit final : OpRewritePattern { LogicalResult matchAndRewrite(StaticOp op, PatternRewriter& rewriter) const override { auto funcOp = op->getParentOfType(); - if (!funcOp || op->getBlock() == &funcOp.getBody().front()) { + if (!funcOp || + op->getParentWithTrait() != funcOp || + op->getBlock() == &funcOp.getBody().front()) { return failure(); } rewriter.moveOpBefore(op, &funcOp.getBody().front(), diff --git a/mlir/unittests/Dialect/QC/IR/test_qc_ir.cpp b/mlir/unittests/Dialect/QC/IR/test_qc_ir.cpp index daaceed3fa..fa17cf9476 100644 --- a/mlir/unittests/Dialect/QC/IR/test_qc_ir.cpp +++ b/mlir/unittests/Dialect/QC/IR/test_qc_ir.cpp @@ -181,6 +181,30 @@ TEST_F(QCTest, CleanupHoistsAndCoalescesStaticQubits) { EXPECT_EQ(staticOps, 2U); } +TEST_F(QCTest, CleanupDoesNotHoistStaticQubitsAcrossIsolationBoundaries) { + auto module = parseSourceString(R"mlir( + module { + func.func @main() { + "builtin.module"() ({ + %q = qc.static 0 : !qc.qubit + qc.x %q : !qc.qubit + }) : () -> () + return + } + } + )mlir", + context.get()); + ASSERT_TRUE(module); + ASSERT_TRUE(succeeded(verify(*module))); + ASSERT_TRUE(succeeded(runQCCleanupPipeline(*module))); + ASSERT_TRUE(succeeded(verify(*module))); + + StaticOp staticOp; + module->walk([&](StaticOp op) { staticOp = op; }); + ASSERT_TRUE(staticOp); + EXPECT_TRUE(isa(staticOp->getParentOp())); +} + TEST_F(QCTest, BuilderRejectsMixedStaticAndDynamicQubitAllocationModes) { EXPECT_DEATH( {