diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d19fb6511..8a3b9e9988 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ releases may include breaking changes. [#1807], [#1808], [#1815], [#1824], [#1869], [#1872], [#1914], [#1925], [#1927], [#1935], [#1936], [#1938], [#1975], [#1976], [#2006], [#2014], [#2015], [#2017], [#2026], [#2028], [#2054], [#2058], [#2125], [#2136], - [#2149], [#2150], [#2158], [#2194], [#2210], [#2211], [#2220]) + [#2149], [#2150], [#2158], [#2194], [#2196], [#2210], [#2211], [#2220]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**], [**@J4MMlE**]) @@ -890,6 +890,7 @@ for previous changelogs._ [#2216]: https://github.com/munich-quantum-toolkit/core/pull/2216 [#2203]: https://github.com/munich-quantum-toolkit/core/pull/2203 [#2214]: https://github.com/munich-quantum-toolkit/core/pull/2214 +[#2196]: https://github.com/munich-quantum-toolkit/core/pull/2196 [#2194]: https://github.com/munich-quantum-toolkit/core/pull/2194 [#2193]: https://github.com/munich-quantum-toolkit/core/pull/2193 [#2184]: https://github.com/munich-quantum-toolkit/core/pull/2184 diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 87f1d1a17c..f32717a6f3 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1882,6 +1882,33 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { const std::variant& index, ValueRange yieldedValues); + //===--------------------------------------------------------------------===// + // Additional functions + //===--------------------------------------------------------------------===// + + /// Starts a private function and returns its entry-block arguments. + /// + /// Functions must be defined before operations in `main` and cannot nest. + SmallVector startFunction(StringRef name, TypeRange argTypes, + TypeRange resultTypes); + + /// Ends the active function and restores the surrounding builder scope. + /// + /// Every linear value in the function must be returned or consumed. + void endFunction(ValueRange returnValues); + + /// Calls a completed function and returns its results. + /// + /// Linear operands and results are paired by following the callee body. + /// Mapping failure is a usage error. + SmallVector call(StringRef callee, ValueRange operands); + + /// Returns the `!qco.qubit` type. + Type getQubitType(); + + /// Returns `tensor`. + Type getQubitTensorType(int64_t size); + //===--------------------------------------------------------------------===// // Finalization //===--------------------------------------------------------------------===// @@ -2070,6 +2097,9 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /// Ensure static and dynamic qubit allocation modes are not mixed. void ensureAllocationMode(AllocationMode requestedMode); + + // Insertion point to restore after finishing an additional function. + OpBuilder::InsertPoint savedInsertionPoint; }; } // namespace qco } // namespace mlir diff --git a/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt b/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt index d8f282e109..8c4c60f853 100644 --- a/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt +++ b/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt @@ -17,7 +17,9 @@ add_mlir_library( MLIRMQTDialect MLIRSCFDialect MLIRQCODialect + MLIRQCOUtils MLIRQTensorDialect + MLIRQTensorUtils PRIVATE MLIRMQTUtils) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index c0bab129df..2c20a4c92e 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -18,8 +18,10 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/QCOUtils.h" +#include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include "mlir/Dialect/QTensor/Utils/TensorIterator.h" #include #include @@ -37,6 +39,7 @@ #include #include #include +#include #include #include #include @@ -1409,6 +1412,215 @@ QCOProgramBuilder::scfCondition(Value reg, return scfCondition(condition, yieldedValues); } +//===----------------------------------------------------------------------===// +// Additional Functions +//===----------------------------------------------------------------------===// + +Type QCOProgramBuilder::getQubitType() { return QubitType::get(ctx); } + +Type QCOProgramBuilder::getQubitTensorType(int64_t size) { + return RankedTensorType::get({size}, getQubitType()); +} + +static bool isQubitTensor(Type type) { + auto tensorType = dyn_cast(type); + return tensorType && isa(tensorType.getElementType()); +} + +SmallVector QCOProgramBuilder::startFunction(StringRef name, + TypeRange argTypes, + TypeRange resultTypes) { + checkFinalized(); + + if (SymbolTable::lookupSymbolIn(module, name) != nullptr) { + llvm::reportFatalUsageError("Function with the same name already exists"); + } + + if (savedInsertionPoint.isSet()) { + llvm::reportFatalUsageError( + "Cannot start a function while another one is being built"); + } + + // Defining callees first prevents their bodies from capturing values from + // main and removes the need to preserve partially built main state. + if (!getInsertionBlock()->empty()) { + llvm::reportFatalUsageError( + "Functions must be defined before operations in main"); + } + savedInsertionPoint = saveInsertionPoint(); + + setInsertionPointToEnd(cast(module).getBody()); + func::FuncOp funcOp = + func::FuncOp::create(*this, name, getFunctionType(argTypes, resultTypes)); + // The interprocedural passes only consider functions that are not externally + // visible, so additional functions are private by default. + funcOp.setPrivate(); + + Block& entryBlock = funcOp.getBody().emplaceBlock(); + SmallVector locs(argTypes.size(), getLoc()); + entryBlock.addArguments(argTypes, locs); + setInsertionPointToStart(&entryBlock); + + SmallVector args; + for (BlockArgument arg : entryBlock.getArguments()) { + if (isa(arg.getType())) { + validQubits.insert(arg); + } else if (isQubitTensor(arg.getType())) { + // A tensor argument acts like a register the callee owns for the + // duration of the call, so give it its own register id. + validTensors.insert(Tensor{arg, tensorCounter++}); + } + args.emplace_back(arg); + } + + return args; +} + +void QCOProgramBuilder::endFunction(ValueRange returnValues) { + checkFinalized(); + + if (!savedInsertionPoint.isSet()) { + llvm::reportFatalUsageError( + "endFunction() called without a matching startFunction()"); + } + + func::FuncOp funcOp = cast(getInsertionBlock()->getParentOp()); + if (!llvm::equal(returnValues.getTypes(), funcOp.getResultTypes())) { + llvm::reportFatalUsageError( + "Return values do not match the declared function result types"); + } + + for (Value value : returnValues) { + if (isa(value.getType())) { + validateQubitValue(value); + validQubits.erase(value); + } else if (isQubitTensor(value.getType())) { + validateTensorValue(value); + validTensors.erase(value); + } + } + + // Only values created inside the function are tracked at this point, so + // anything left over has escaped. + if (!validQubits.empty()) { + llvm::reportFatalUsageError( + "Function body has qubit values that are neither returned nor " + "consumed"); + } + if (!validTensors.empty()) { + llvm::reportFatalUsageError( + "Function body has tensor values that are neither returned nor " + "deallocated"); + } + + func::ReturnOp::create(*this, returnValues); + + OpBuilder::InsertPoint insertionPoint = savedInsertionPoint; + savedInsertionPoint = {}; + restoreInsertionPoint(insertionPoint); +} + +SmallVector QCOProgramBuilder::call(StringRef callee, + ValueRange operands) { + checkFinalized(); + + func::FuncOp funcOp = dyn_cast_or_null( + SymbolTable::lookupSymbolIn(module, getStringAttr(callee))); + if (!funcOp) { + llvm::reportFatalUsageError("Callee not found in module"); + } + + if (!llvm::equal(operands.getTypes(), funcOp.getArgumentTypes())) { + llvm::reportFatalUsageError( + "Call operands do not match the declared function argument types"); + } + + // Re-insert qubits that were extracted from a tensor operand, so the callee + // receives a complete register. Every other construct that hands a tensor to + // a nested region does the same before building it. Unlike those, a call may + // also carry classical operands, so `prepareInitArgs` cannot be used here; + // qubits passed alongside their tensor are excluded from the re-insertion. + DenseSet qubitOperandSet; + for (Value operand : operands) { + if (isa(operand.getType())) { + qubitOperandSet.insert(operand); + } + } + SmallVector preparedOperands; + preparedOperands.reserve(operands.size()); + SmallVector qubitOperands; + SmallVector tensorOperands; + for (Value operand : operands) { + Value prepared = isQubitTensor(operand.getType()) + ? prepareInitArg(operand, &qubitOperandSet) + : operand; + preparedOperands.emplace_back(prepared); + if (isa(prepared.getType())) { + validateQubitValue(prepared); + qubitOperands.emplace_back(prepared); + } else if (isQubitTensor(prepared.getType())) { + validateTensorValue(prepared); + tensorOperands.emplace_back(prepared); + } + } + + func::CallOp callOp = func::CallOp::create(*this, funcOp, preparedOperands); + + // Thread each qubit operand into the result that continues its wire. The + // correspondence is derived from the callee body instead of assumed to be + // positional, so a callee that hands its qubits back in a different order + // than it takes them is tracked the way it actually behaves. + CallQubitMapping qubitMapping; + DenseSet continuedResults; + for (Value operand : qubitOperands) { + auto resultOr = qubitMapping.getResultForOperand(callOp, operand); + if (failed(resultOr)) { + llvm::reportFatalUsageError( + "Cannot derive linear-value correspondence for callee"); + } + Value result = *resultOr; + if (!result) { + // The callee keeps this qubit. + validQubits.erase(operand); + continue; + } + updateQubitTracking(operand, result); + continuedResults.insert(result); + } + // Qubit tensors are threaded the same way, using the tensor counterpart of + // the mapping above. + qtensor::CallTensorMapping tensorMapping; + for (Value operand : tensorOperands) { + auto resultOr = tensorMapping.getResultForOperand(callOp, operand); + if (failed(resultOr)) { + llvm::reportFatalUsageError( + "Cannot derive linear-value correspondence for callee"); + } + Value result = *resultOr; + if (!result) { + // The callee keeps this tensor. + validTensors.erase(operand); + continue; + } + updateTensorTracking(operand, result); + continuedResults.insert(result); + } + + // Results without a corresponding operand were created by the callee. + for (Value result : callOp.getResults()) { + if (continuedResults.contains(result)) { + continue; + } + if (isa(result.getType())) { + validQubits.insert(result); + } else if (isQubitTensor(result.getType())) { + validTensors.insert(Tensor{result, tensorCounter++}); + } + } + + return SmallVector(callOp.getResults()); +} + //===----------------------------------------------------------------------===// // Finalization //===----------------------------------------------------------------------===// diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 2ef1cf24e3..329165979b 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -323,6 +323,155 @@ TEST_F(QCOTest, BuilderSupportsIndependentClassicalRegisterInitialization) { "undefined"); } +TEST_F(QCOTest, BuilderSupportsAdditionalFunctions) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + + SmallVector args = + builder.startFunction("thread", {qubitType}, {qubitType}); + builder.endFunction({builder.h(args[0])}); + + Value qubit = builder.allocQubit(); + Value tensor = builder.qtensorAlloc(2); + SmallVector results = builder.call("thread", {qubit}); + builder.sink(results[0]); + builder.qtensorDealloc(tensor); + EXPECT_TRUE(builder.finalize()); +} + +TEST_F(QCOTest, BuilderTracksKeptAndCreatedLinearValues) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + Type tensorType = builder.getQubitTensorType(2); + + SmallVector args = builder.startFunction( + "replace", {qubitType, tensorType}, {qubitType, tensorType}); + builder.sink(args[0]); + builder.qtensorDealloc(args[1]); + builder.endFunction({builder.allocQubit(), builder.qtensorAlloc(2)}); + + Value qubit = builder.allocQubit(); + Value tensor = builder.qtensorAlloc(2); + SmallVector results = builder.call("replace", {qubit, tensor}); + builder.sink(results[0]); + builder.qtensorDealloc(results[1]); + EXPECT_TRUE(builder.finalize()); +} + +TEST_F(QCOTest, BuilderRejectsInvalidFunctionStateAndSymbols) { + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + builder.startFunction("f", {qubitType}, {qubitType}); + builder.startFunction("g", {qubitType}, {qubitType}); + }, + "Cannot start a function while another one is being built"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + builder.endFunction({}); + }, + "endFunction\\(\\) called without a matching startFunction\\(\\)"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + builder.call("does_not_exist", {}); + }, + "Callee not found in module"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + SmallVector args = + builder.startFunction("f", {qubitType}, {qubitType}); + builder.call("f", {args[0]}); + }, + "Cannot derive linear-value correspondence for callee"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + SmallVector args = + builder.startFunction("f", {qubitType}, {qubitType}); + builder.endFunction({args[0]}); + builder.startFunction("f", {qubitType}, {qubitType}); + }, + "Function with the same name already exists"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + builder.allocQubit(); + Type qubitType = builder.getQubitType(); + builder.startFunction("f", {qubitType}, {qubitType}); + }, + "Functions must be defined before operations in main"); +} + +TEST_F(QCOTest, BuilderRejectsInvalidFunctionValues) { + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + SmallVector args = + builder.startFunction("f", {qubitType}, {qubitType}); + auto measured = builder.measure(args[0]); + builder.sink(measured.first); + builder.endFunction({measured.second}); + }, + "Return values do not match the declared function result types"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + SmallVector args = + builder.startFunction("f", {qubitType}, {qubitType}); + builder.endFunction({args[0]}); + builder.call("f", {builder.floatConstant(0.5)}); + }, + "Call operands do not match the declared function argument types"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + SmallVector args = + builder.startFunction("f", {qubitType}, {qubitType}); + builder.allocQubit(); + builder.endFunction({args[0]}); + }, + "neither returned nor consumed"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Type qubitType = builder.getQubitType(); + SmallVector args = + builder.startFunction("f", {qubitType}, {qubitType}); + builder.qtensorAlloc(2); + builder.endFunction({args[0]}); + }, + "neither returned nor deallocated"); +} + TEST_F(QCOTest, DirectSingleQubitPowBuilder) { QCOProgramBuilder builder(context.get()); builder.initialize();