diff --git a/mlir/lib/Dialect/QIR/Utils/QIRUtils.cpp b/mlir/lib/Dialect/QIR/Utils/QIRUtils.cpp index 5929e072c5..c315130ee7 100644 --- a/mlir/lib/Dialect/QIR/Utils/QIRUtils.cpp +++ b/mlir/lib/Dialect/QIR/Utils/QIRUtils.cpp @@ -282,15 +282,26 @@ LLVM::LLVMFuncOp getOrCreateFunctionDeclaration(OpBuilder& builder, builder.setInsertionPointToEnd(moduleOp.getBody()); fnDecl = LLVM::LLVMFuncOp::create(builder, op->getLoc(), fnName, fnType); + } - // Add irreversible attribute to irreversible quantum operations - if (fnName == QIR_MEASURE || fnName == QIR_RESET) { - fnDecl->setAttr("passthrough", - builder.getStrArrayAttr({::qir::IRREVERSIBLE_ATTR})); - } + auto function = cast(fnDecl); + if (fnName != QIR_MEASURE && fnName != QIR_RESET) { + return function; } - return cast(fnDecl); + const auto irreversible = builder.getStringAttr(::qir::IRREVERSIBLE_ATTR); + const auto passthrough = function->getAttrOfType("passthrough"); + if (passthrough && llvm::is_contained(passthrough, irreversible)) { + return function; + } + + SmallVector entries; + if (passthrough) { + entries.append(passthrough.begin(), passthrough.end()); + } + entries.push_back(irreversible); + function->setAttr("passthrough", builder.getArrayAttr(entries)); + return function; } LLVM::AddressOfOp createResultLabel(OpBuilder& builder, Operation* op, diff --git a/mlir/unittests/Dialect/QIR/IR/test_qir_ir.cpp b/mlir/unittests/Dialect/QIR/IR/test_qir_ir.cpp index 384f0889d3..a3fe86125a 100644 --- a/mlir/unittests/Dialect/QIR/IR/test_qir_ir.cpp +++ b/mlir/unittests/Dialect/QIR/IR/test_qir_ir.cpp @@ -180,6 +180,42 @@ TEST_F(QIRTest, BuilderReturnsCompleteClassicalRegister) { EXPECT_FALSE(returnedRegister.array); } +TEST_F(QIRTest, ReusedIrreversibleDeclarationsPreservePassthroughIdempotently) { + OpBuilder builder(context.get()); + const auto location = builder.getUnknownLoc(); + auto moduleOp = ModuleOp::create(location); + builder.setInsertionPointToStart(moduleOp.getBody()); + const auto ptrType = LLVM::LLVMPointerType::get(context.get()); + const auto voidType = LLVM::LLVMVoidType::get(context.get()); + const auto nounwind = builder.getStringAttr("nounwind"); + const auto targetCPU = builder.getStrArrayAttr({"target-cpu", "generic"}); + + for (const StringRef name : {StringRef(QIR_MEASURE), StringRef(QIR_RESET)}) { + const SmallVector parameters(name == QIR_MEASURE ? 2 : 1, ptrType); + const auto functionType = LLVM::LLVMFunctionType::get(voidType, parameters); + auto declaration = + LLVM::LLVMFuncOp::create(builder, location, name, functionType); + declaration->setAttr("passthrough", + builder.getArrayAttr({nounwind, targetCPU})); + + EXPECT_EQ( + getOrCreateFunctionDeclaration(builder, moduleOp, name, functionType), + declaration); + EXPECT_EQ( + getOrCreateFunctionDeclaration(builder, moduleOp, name, functionType), + declaration); + + const auto passthrough = + declaration->getAttrOfType("passthrough"); + ASSERT_TRUE(passthrough); + ASSERT_EQ(passthrough.size(), 3U); + EXPECT_EQ(passthrough[0], nounwind); + EXPECT_EQ(passthrough[1], targetCPU); + EXPECT_EQ(llvm::count(passthrough, builder.getStringAttr("irreversible")), + 1); + } +} + TEST_F(QIRTest, AdaptiveBuilderSelectsControlledSpecializationsByArity) { auto module = QIRProgramBuilder::build( context.get(),