diff --git a/CHANGELOG.md b/CHANGELOG.md index 33db26d50f..6a4f014b36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,7 +15,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel - ✨ Add conversions between Jeff and QCO ([#1479], [#1548], [#1565]) ([**@denialhaag**]) - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568]) ([**@MatthiasReumann**]) - ✨ Add initial infrastructure for new QC and QCO MLIR dialects - ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1548], [#1550], [#1554], [#1570]) + ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1570]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**]) ### Changed @@ -343,6 +343,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool [#1549]: https://github.com/munich-quantum-toolkit/core/pull/1549 [#1548]: https://github.com/munich-quantum-toolkit/core/pull/1548 [#1547]: https://github.com/munich-quantum-toolkit/core/pull/1547 +[#1542]: https://github.com/munich-quantum-toolkit/core/pull/1542 [#1537]: https://github.com/munich-quantum-toolkit/core/pull/1537 [#1521]: https://github.com/munich-quantum-toolkit/core/pull/1521 [#1513]: https://github.com/munich-quantum-toolkit/core/pull/1513 diff --git a/docs/mlir/QTensor.md b/docs/mlir/QTensor.md new file mode 100644 index 0000000000..b12a912fbf --- /dev/null +++ b/docs/mlir/QTensor.md @@ -0,0 +1,7 @@ +--- +tocdepth: 3 +--- + +```{include} Dialects/QTensorDialect.md +:heading-offset: 1 +``` diff --git a/docs/mlir/index.md b/docs/mlir/index.md index 3b5e91b20c..f928fd86cd 100644 --- a/docs/mlir/index.md +++ b/docs/mlir/index.md @@ -8,7 +8,9 @@ We define multiple dialects, each with its dedicated purpose: - The {doc}`QCO dialect ` uses value semantics and is mainly designed for running optimizations. -Both dialects define various canonicalization and transformations that enable the compilation of quantum programs to native quantum hardware. +- The {doc}`QTensor dialect ` adds support for one-dimensional tensors of qubits with linear typing and is used in the QCO dialect to represent collections of qubits such as registers. + +These dialects define various canonicalization and transformations that enable the compilation of quantum programs to native quantum hardware. For intercompatibility, we provide {doc}`conversions ` between dialects. So far, this comprises conversions between QC and QCO as well as from QC to QIR. @@ -18,6 +20,7 @@ So far, this comprises conversions between QC and QCO as well as from QC to QIR. QC QCO +QTensor Conversions ``` diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt index 01edd5218a..dfc77dc6e2 100644 --- a/mlir/include/mlir/Dialect/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/CMakeLists.txt @@ -8,3 +8,4 @@ add_subdirectory(QC) add_subdirectory(QCO) +add_subdirectory(QTensor) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index c7955ad75e..a09c49ea7a 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -203,6 +203,182 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { [[nodiscard]] ClassicalRegister allocClassicalBitRegister(int64_t size, std::string name = "c") const; + //===--------------------------------------------------------------------===// + // QTensor operations + //===--------------------------------------------------------------------===// + + /** + * @brief Allocate a qubit tensor + * + * @details + * Allocates a one-dimensional tensor of !qco.qubit types with the given size + * if the size is a constant, otherwise the tensor has dynamic size. The + * qubits are initialized in the |0> state. The resulting tensor is added to + * the tracking. + * + * @param size Number of qubits (must be positive) + * @return The allocated tensor + * + * @par Example: + * ```c++ + * auto tensor = builder.qtensorAlloc(3); + * ``` + * ```mlir + * %tensor = qtensor.alloc(%c3) : tensor<3x!qco.qubit> + * ``` + */ + Value qtensorAlloc(const std::variant& size); + + /** + * @brief Allocate a qubit tensor from a list of qubit values + * + * @details + * Consumes the input qubits and creates a one-dimensional tensor of + * !qco.qubit types. The resulting tensor has a static size given by the + * number of input values. The consumed qubits are removed from the qubit + * tracking and the resulting tensor is added to the tracking. + * + * @param elements Inserted Qubits (must be valid/unconsumed) + * @return The allocated tensor + * + * @par Example: + * ```c++ + * auto tensor = builder.qtensorFromElements({q0, q1, q2}); + * ``` + * ```mlir + * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + * ``` + */ + Value qtensorFromElements(ValueRange elements); + + /** + * @brief Extract a qubit from a tensor + * + * @details + * Extracts a qubit from a one-dimensional tensor of qubits at the given index + * and returns the updated tensor and the extracted qubit. The extracted qubit + * is added to the qubit tracking and the tracking of the source tensor is + * updated. + * + * @param tensor Source tensor (must be valid/unconsumed) + * @param index The index from where the qubit is extracted + * @return Pair of (outTensor, extractedQubit) + * + * @par Example: + * ```c++ + * auto [outTensor, q0] = builder.qtensorExtract(tensor, 0); + * ``` + * ```mlir + * %outTensor, %q0 = qtensor.extract %tensor[%c0]: tensor<3x!qco.qubit> + * ``` + */ + std::pair + qtensorExtract(Value tensor, const std::variant& index); + + /** + * @brief Extract a qubit slice from a tensor + * + * @details + * Extracts a slice from a one-dimensional tensor of qubits at the given + * offset and size and returns the updated input tensor and the extracted + * tensor. The extracted tensor is added to the qubit tensor tracking and the + * tracking for the input tensor is updated. + * + * @param tensor Source tensor (must be valid/unconsumed) + * @param offset The offset from where the slice is extracted + * @param size The size of the extracted slice + * @return Pair of (outTensor, extractedSlice) + * + * @par Example: + * ```c++ + * auto [outTensor, extractedSlice] = builder.qtensorExtractSlice(tensor, 0, + * 2); + * ``` + * ```mlir + * %outTensor, %extractedSlice = qtensor.extract_slice %tensor[%c0][%c2] + * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + * ``` + */ + std::pair + qtensorExtractSlice(Value tensor, const std::variant& offset, + const std::variant& size); + + /** + * @brief Insert a qubit into a tensor + * + * @details + * Inserts a scalar qubit into the one-dimensional tensor of qubits at the + * given index. The inserted qubit is consumed and removed from the qubit + * tracking while the tracking for the source tensor is updated. + * + * @param scalar The scalar qubit that is inserted (must be valid/unconsumed) + * @param tensor The tensor where the qubit is inserted (must be + * valid/unconsumed) + * @param index The index into where the qubit is inserted + * @return The output tensor + * + * @par Example: + * ```c++ + * auto outTensor = builder.qtensorInsert(q0, tensor, 0); + * ``` + * ```mlir + * %outTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> + * ``` + */ + Value qtensorInsert(Value scalar, Value tensor, + const std::variant& index); + + /** + * @brief Insert a qubit slice into a tensor + * + * @details + * Inserts a one-dimensional tensor of qubits into another one-dimensional + * tensor of qubits at the given offset and size. The inserted tensor slice is + * consumed and removed from the tracking, while the tracking for the + * destination tensor is updated. + * + * @param sourceTensor The slice that is inserted (must be valid/unconsumed) + * @param destTensor The tensor where the slice is inserted (must be + * valid/unconsumed) + * @param offset The offset into where the slice is inserted + * @param size The size of the inserted slice + * @return The output tensor + * + * @par Example: + * ```c++ + * auto outTensor = builder.qtensorInsertSlice(slicedTensor, tensor, 0, 2); + * ``` + * ```mlir + * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2] + * : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + * ``` + */ + Value qtensorInsertSlice(Value sourceTensor, Value destTensor, + const std::variant& offset, + const std::variant& size); + + /** + * @brief Explicitly deallocate a tensor + * + * @details + * Validates and removes the tensor from tracking. Qubits or tensors of qubits + * that were extracted from the tensor but not inserted back again need to be + * deallocated separately. Optional; `finalize()` automatically deallocates + * all remaining tensors. + * + * @param tensor Tensor to deallocate (must be valid/unconsumed) + * @return Reference to this builder for method chaining + * + * @par Example: + * ```c++ + * builder.qtensorDealloc(tensor); + * ``` + * ```mlir + * qtensor.dealloc %tensor : tensor<3x!qco.qubit> + * ``` + */ + QCOProgramBuilder& qtensorDealloc(Value tensor); + //===--------------------------------------------------------------------===// // Measurement and Reset //===--------------------------------------------------------------------===// @@ -1055,7 +1231,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @brief Explicitly deallocate a qubit * * @details - * Validates and removes the qubit from tracking. Optional, finalize() + * Validates and removes the qubit from tracking. Optional; `finalize()` * automatically deallocates all remaining qubits. * * @param qubit Qubit to deallocate (must be valid/unconsumed) @@ -1124,10 +1300,10 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @brief Finalize the program and return the constructed module * * @details - * Automatically deallocates all remaining valid qubits, adds a return - * statement with exit code 0 (indicating successful execution), and - * transfers ownership of the module to the caller. - * The builder should not be used after calling this method. + * Automatically deallocates all remaining valid qubits and tensors of qubits, + * adds a return statement with exit code 0 (indicating successful execution), + * and transfers ownership of the module to the caller. The builder should not + * be used after calling this method. * * @return OwningOpRef containing the constructed quantum program module */ @@ -1176,5 +1352,27 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /// When an operation consumes a qubit and produces a new one, the old value /// is removed and the new output is added. llvm::DenseSet validQubits; + + /** + * @brief Validate that a tensor value is valid and unconsumed. This also + * checks if the tensor is one-dimensional and contains !qco.qubit as its + * values + * @param tensor Tensor value to validate + * @throws Aborts if tensor is not tracked (consumed or never created) + */ + void validateTensorValue(Value tensor) const; + + /** + * @brief Update tracking when an operation consumes and produces a tensor + * @param inputTensor Input tensor being consumed (must be valid) + * @param outputTensor New output tensor being produced + */ + void updateTensorTracking(Value inputTensor, Value outputTensor); + + /// Track valid (unconsumed) tensor SSA values for linear type enforcement. + /// Only values present in this set are valid for use in operations. + /// When an operation consumes a tensor and produces a new one, the old value + /// is removed and the new output is added. + llvm::DenseSet validTensors; }; } // namespace mlir::qco diff --git a/mlir/include/mlir/Dialect/QTensor/CMakeLists.txt b/mlir/include/mlir/Dialect/QTensor/CMakeLists.txt new file mode 100644 index 0000000000..b181a84fed --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +# All rights reserved. +# +# SPDX-License-Identifier: MIT +# +# Licensed under the MIT License + +add_subdirectory(IR) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt new file mode 100644 index 0000000000..0d66effe7c --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt @@ -0,0 +1,11 @@ +# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +# All rights reserved. +# +# SPDX-License-Identifier: MIT +# +# Licensed under the MIT License + +add_mlir_dialect(QTensorOps qtensor) + +add_mlir_doc(QTensorOps QTensorDialect Dialects/ -gen-dialect-doc -dialect=qtensor) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h new file mode 100644 index 0000000000..1d88688be8 --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#pragma once + +#include +#include + +#define DIALECT_NAME_QTensor "qtensor" + +//===----------------------------------------------------------------------===// +// QTensor Dialect +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/QTensor/IR/QTensorOpsDialect.h.inc" // IWYU pragma: export + +//===----------------------------------------------------------------------===// +// Types +//===----------------------------------------------------------------------===// + +#define GET_TYPEDEF_CLASSES +#include "mlir/Dialect/QTensor/IR/QTensorOpsTypes.h.inc" // IWYU pragma: export diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td new file mode 100644 index 0000000000..75122363fa --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td @@ -0,0 +1,30 @@ +// Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +// Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +// All rights reserved. +// +// SPDX-License-Identifier: MIT +// +// Licensed under the MIT License + +#ifndef MLIR_DIALECT_QTENSOR_IR_QTENSORDIALECT_TD +#define MLIR_DIALECT_QTENSOR_IR_QTENSORDIALECT_TD + +include "mlir/IR/DialectBase.td" + +def QTensorDialect : Dialect { + let name = "qtensor"; + + let summary = "The QTensor dialect for one-dimensional tensors of qubits with linear typing."; + + let description = [{ + The QTensor dialect is an adjusted variant of the standard tensor dialect of MLIR that supports linear typing for one-dimensional tensors of qubits. + In order to support linear typing, the extract operations of this dialect are modified so that they also return the updated tensor. + In addition, alloc/dealloc operations are added to the dialect to support the bulk allocation and deallocation of qubit tensors with linear types. + }]; + + let dependentDialects = ["::mlir::qco::QCODialect"]; + + let cppNamespace = "::mlir::qtensor"; +} + +#endif // MLIR_DIALECT_QTENSOR_IR_QTENSORDIALECT_TD diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h new file mode 100644 index 0000000000..677873c6e6 --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -0,0 +1,32 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#pragma once + +// Suppress warnings about ambiguous reversed operators in MLIR +// (see https://github.com/llvm/llvm-project/issues/45853) +#include +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wambiguous-reversed-operator" +#endif +#include +#ifdef __clang__ +#pragma clang diagnostic pop +#endif + +#include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" + +#include +#include + +#define GET_OP_CLASSES +#include "mlir/Dialect/QTensor/IR/QTensorOps.h.inc" // IWYU pragma: export diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td new file mode 100644 index 0000000000..9322eace45 --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -0,0 +1,261 @@ +// Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +// Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +// All rights reserved. +// +// SPDX-License-Identifier: MIT +// +// Licensed under the MIT License + +#ifndef MLIR_DIALECT_QTENSOR_IR_QTENSOROPS_TD +#define MLIR_DIALECT_QTENSOR_IR_QTENSOROPS_TD + +include "mlir/Dialect/QCO/IR/QCOTypes.td" +include "mlir/Dialect/QTensor/IR/QTensorDialect.td" + +include "mlir/IR/OpBase.td" +include "mlir/Interfaces/InferIntRangeInterface.td" +include "mlir/Interfaces/InferTypeOpInterface.td" +include "mlir/Interfaces/ShapedOpInterfaces.td" +include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/TilingInterface.td" + +//===----------------------------------------------------------------------===// +// Base Operation Classes +//===----------------------------------------------------------------------===// + +class QTensorOp traits = []> + : Op; + +//===----------------------------------------------------------------------===// +// Type Constraints +//===----------------------------------------------------------------------===// + +def Static1DQubitTensor : Type< + And<[ + 1DTensorOf<[QubitType]>.predicate, + HasStaticShapePred + ]>, + "Tensor must be statically shaped and elements must be of qubit type">; + +//===----------------------------------------------------------------------===// +// Operations +//===----------------------------------------------------------------------===// + +def AllocOp : QTensorOp<"alloc", [MemoryEffects<[MemAlloc]>]> { + let summary = "Tensor alloc operation"; + let description = [{ + Allocates a one-dimensional qubit tensor with the given size and returns the allocated tensor. + The qubits are initialized to the |0> state. If the size is a constant the resulting tensor has a static size, + otherwise the tensor has a dynamic size. + + Example: + ```mlir + %qtensor = qtensor.alloc(%c3) : tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins Index:$size); + let results = (outs 1DTensorOf<[QubitType]>:$result); + let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; + + let builders = [ + OpBuilder<(ins "Value":$size)> + ]; + + let hasVerifier = 1; +} + +def DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { + let summary = "Deallocate a tensor"; + let description = [{ + This operation deallocates the tensor and the values it holds, releasing its resources. + Qubits that were extracted but not inserted again need to be deallocated separately. + + Example: + ```mlir + qtensor.dealloc %tensor : tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins 1DTensorOf<[QubitType]>:$tensor); + let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; + + let hasCanonicalizer = 1; +} + +def FromElementsOp : QTensorOp<"from_elements", [ + Pure, + TypesMatchWith<"operand types match result element type", + "result", "elements", "SmallVector(" + "::llvm::cast($_self).getNumElements(), " + "::llvm::cast($_self).getElementType())"> + ]> { + let summary = "Create a tensor from a range of values"; + let description = [{ + The `qtensor.from_elements` operation is a wrapper operation of the `tensor.from_elements` operation. + This operation creates a one-dimensional qubit tensor using the operands as its input. + This also causes the input qubits to be consumed. + + Example: + ```mlir + %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins Variadic:$elements); + let results = (outs Static1DQubitTensor:$result); + let assemblyFormat = "$elements attr-dict `:` type($result)"; + + let builders = [ + // Special case builder for when `elements` has size >=1. + OpBuilder<(ins "ValueRange":$elements)> + ]; +} + +def ExtractOp : QTensorOp<"extract", [ + Pure, + TypesMatchWith<"result type matches element type of tensor", + "tensor", "result", + "::llvm::cast($_self).getElementType()">, + TypesMatchWith< + "returned tensor type matches input tensor", "tensor", "out_tensor", "$_self">]> { + let summary = "Extract element from tensor"; + let description = [{ + The `qtensor.extract` operation is the modified version of the standard `tensor.extract` + operation of the tensor dialect. It reads a ranked tensor and returns one qubit from the input tensor + at the given index. In addition, it also returns the updated input tensor as a result. + + Example: + ```mlir + %outTensor, %q0 = qtensor.extract %tensor[%c0] : tensor<4x!qco.qubit> + ``` + }]; + + let arguments = (ins 1DTensorOf<[QubitType]>:$tensor, Index:$index); + let results = (outs 1DTensorOf<[QubitType]>:$out_tensor, QubitType:$result); + let assemblyFormat = "$tensor `[` $index `]` attr-dict `:` type($tensor)"; + + let hasFolder = 1; + let hasVerifier = 1; +} + +def ExtractSliceOp : QTensorOp<"extract_slice", [ + Pure, + TypesMatchWith<"returned tensor type matches input tensor", + "tensor", "out_tensor", "$_self"> + ]> { + let summary = "Extract slice from tensor"; + let description = [{ + The `qtensor.extract_slice` operation is the modified version of the standard `tensor.extract_slice` + operation of the tensor dialect. It reads a one-dimensional qubit tensor and returns the extracted tensor of qubits specified by the + offset and size argument. In addition, it also returns the updated input tensor as result. + + The extract_slice operation supports the following arguments: + + - tensor: the "base" tensor from which to extract a slice. + - offset: the starting position in the base tensor from which the slice is + extracted. + - size: the length of the slice to extract from the base tensor. + + Example: + ```mlir + %outTensor, %extractedSlice = qtensor.extract_slice %tensor[%c0][%c2] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + ``` + }]; + + let arguments = (ins + 1DTensorOf<[QubitType]>:$tensor, + Index:$offset, + Index:$size + ); + let results = (outs 1DTensorOf<[QubitType]>:$out_tensor, 1DTensorOf<[QubitType]>:$result); + + let assemblyFormat = [{ + $tensor `[`$offset `]` `[`$size`]` + attr-dict `:` type($tensor) `to` type($result) + }]; + + let builders = [ + OpBuilder<(ins "Value":$tensor, "Value":$offset, + "Value":$size, + CArg<"ArrayRef", "{}">:$attrs)> + ]; + + let hasFolder = 1; + let hasVerifier = 1; +} + +def InsertOp : QTensorOp<"insert", [ + Pure, + TypesMatchWith<"result type matches type of dest", + "dest", "result", + "$_self">, + TypesMatchWith<"scalar type matches element type of dest", + "dest", "scalar", + "::llvm::cast($_self).getElementType()">]> { + let summary = "Insert element into tensor"; + let description = [{ + The `qtensor.insert` operation is a wrapper operation for the `tensor.insert` operation of the tensor dialect. + This operation inserts a qubit scalar into a ranked tensor as specified by the operation's index. The insertion + consumes the input qubit. + + Example: + ```mlir + %newTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins QubitType:$scalar, + 1DTensorOf<[QubitType]>:$dest, + Index:$index); + let results = (outs 1DTensorOf<[QubitType]>:$result); + let assemblyFormat = [{ + $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) + }]; + + let hasFolder = 1; + let hasVerifier = 1; +} + +def InsertSliceOp : QTensorOp<"insert_slice", [ + Pure, + TypesMatchWith<"expected result type to match dest type", + "dest", "result", "$_self"> + ]> { + let summary = "Insert slice into tensor"; + let description = [{ + The `qtensor.insert_slice` operation is a modified version of the `tensor.insert_slice` operation + of the tensor dialect. The operation inserts a tensor `source` into another tensor `dest` as specified by the + operation's offset and size arguments. This insertion consumes the `source` tensor of qubits. + + The insert_slice operation supports the following arguments: + + - source: the tensor that is inserted. The source is consumed after the insertion. + - dest: the tensor into which the source tensor is inserted. + - offset: the starting index in the `dest` tensor where the slice is inserted. + - size: the number of elements in the slice, which must match the size of the + source tensor type. + + Example: + ```mlir + %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2] : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins + 1DTensorOf<[QubitType]>:$source, + 1DTensorOf<[QubitType]>:$dest, + Index:$offset, + Index:$size + ); + let results = (outs 1DTensorOf<[QubitType]>:$result); + let assemblyFormat = [{ + $source `into` $dest `[`$offset `]` `[`$size`]` + attr-dict `:` type($source) `into` type($dest) + }]; + + let hasFolder = 1; + let hasVerifier = 1; +} + +#endif // MLIR_DIALECT_QTENSOR_IR_QTENSOROPS_TD diff --git a/mlir/lib/Conversion/QCOToQC/CMakeLists.txt b/mlir/lib/Conversion/QCOToQC/CMakeLists.txt index 9c403f8951..eba206ceda 100644 --- a/mlir/lib/Conversion/QCOToQC/CMakeLists.txt +++ b/mlir/lib/Conversion/QCOToQC/CMakeLists.txt @@ -16,6 +16,7 @@ add_mlir_conversion_library( LINK_LIBS MLIRQCDialect MLIRQCODialect + MLIRQTensorDialect MLIRArithDialect MLIRFuncDialect MLIRTransforms diff --git a/mlir/lib/Conversion/QCToQCO/CMakeLists.txt b/mlir/lib/Conversion/QCToQCO/CMakeLists.txt index 5a1c1b9797..0474683bfd 100644 --- a/mlir/lib/Conversion/QCToQCO/CMakeLists.txt +++ b/mlir/lib/Conversion/QCToQCO/CMakeLists.txt @@ -16,6 +16,7 @@ add_mlir_conversion_library( LINK_LIBS MLIRQCDialect MLIRQCODialect + MLIRQTensorDialect MLIRArithDialect MLIRFuncDialect MLIRTransforms diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt index c0d91fa751..d8d7ce7b9f 100644 --- a/mlir/lib/Dialect/CMakeLists.txt +++ b/mlir/lib/Dialect/CMakeLists.txt @@ -9,3 +9,4 @@ add_subdirectory(QCO) add_subdirectory(QIR) add_subdirectory(QC) +add_subdirectory(QTensor) diff --git a/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt b/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt index 4ca81dcd93..e55eab1497 100644 --- a/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt +++ b/mlir/lib/Dialect/QCO/Builder/CMakeLists.txt @@ -15,6 +15,7 @@ add_mlir_library( MLIRFuncDialect MLIRSCFDialect MLIRQCODialect + MLIRQTensorDialect DISABLE_INSTALL) mqt_mlir_target_use_project_options(MLIRQCOProgramBuilder) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index a4e9d21bb8..05b55ff6dc 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -12,11 +12,14 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include "mlir/Dialect/Utils/Utils.h" #include #include #include +#include #include #include #include @@ -43,7 +46,7 @@ QCOProgramBuilder::QCOProgramBuilder(MLIRContext* context) : ImplicitLocOpBuilder( FileLineColLoc::get(context, "", 1, 1), context), ctx(context), module(ModuleOp::create(*this)) { - ctx->loadDialect(); + ctx->loadDialect(); } void QCOProgramBuilder::initialize() { @@ -161,6 +164,151 @@ void QCOProgramBuilder::updateQubitTracking(Value inputQubit, validQubits.insert(outputQubit); } +void QCOProgramBuilder::validateTensorValue(Value tensor) const { + if (!validTensors.contains(tensor)) { + llvm::errs() << "Attempting to use an invalid tensor SSA value. " + << "The value may have been consumed by a previous operation " + << "or was never created through this builder.\n"; + llvm::reportFatalUsageError( + "Invalid tensor value used (either consumed or not tracked)"); + } + + auto tensorType = llvm::dyn_cast(tensor.getType()); + if (!tensorType || tensorType.getRank() != 1) { + llvm::reportFatalUsageError("Tensor must be of 1-D RankedTensorType!"); + } + if (!llvm::isa(tensorType.getElementType())) { + llvm::reportFatalUsageError("Elements must be of QubitType!"); + } +} + +void QCOProgramBuilder::updateTensorTracking(Value inputTensor, + Value outputTensor) { + // Validate the input tensor + validateTensorValue(inputTensor); + + // Remove the input (consumed) value from tracking + validTensors.erase(inputTensor); + + // Add the output (new) value to tracking + validTensors.insert(outputTensor); +} + +//===----------------------------------------------------------------------===// +// QTensor Operations +//===----------------------------------------------------------------------===// + +Value QCOProgramBuilder::qtensorAlloc( + const std::variant& size) { + checkFinalized(); + auto sizeValue = utils::variantToValue(*this, getLoc(), size); + + auto allocOp = qtensor::AllocOp::create(*this, sizeValue); + auto result = allocOp.getResult(); + validTensors.insert(result); + return result; +} + +Value QCOProgramBuilder::qtensorFromElements(ValueRange elements) { + checkFinalized(); + + if (elements.empty()) { + llvm::reportFatalUsageError("Elements must contain at least one qubit"); + } + + for (auto element : elements) { + if (!llvm::isa(element.getType())) { + llvm::reportFatalUsageError("Elements must be QubitType!"); + } + validateQubitValue(element); + validQubits.erase(element); + } + + auto fromElementsOp = qtensor::FromElementsOp::create(*this, elements); + auto result = fromElementsOp.getResult(); + validTensors.insert(result); + return result; +} + +std::pair +QCOProgramBuilder::qtensorExtract(Value tensor, + const std::variant& index) { + checkFinalized(); + + auto indexValue = utils::variantToValue(*this, getLoc(), index); + auto extractOp = qtensor::ExtractOp::create(*this, tensor, indexValue); + auto qubit = extractOp.getResult(); + auto outTensor = extractOp.getOutTensor(); + + validQubits.insert(qubit); + updateTensorTracking(tensor, outTensor); + + return {outTensor, qubit}; +} + +std::pair QCOProgramBuilder::qtensorExtractSlice( + Value tensor, const std::variant& offset, + const std::variant& size) { + checkFinalized(); + + auto offsetValue = utils::variantToValue(*this, getLoc(), offset); + auto sizesValue = utils::variantToValue(*this, getLoc(), size); + auto extractSliceOp = + qtensor::ExtractSliceOp::create(*this, tensor, offsetValue, sizesValue); + auto slicedTensor = extractSliceOp.getResult(); + auto outTensor = extractSliceOp.getOutTensor(); + + validTensors.insert(slicedTensor); + updateTensorTracking(tensor, outTensor); + + return {outTensor, slicedTensor}; +} + +Value QCOProgramBuilder::qtensorInsert( + Value scalar, Value tensor, const std::variant& index) { + checkFinalized(); + + auto indexValue = utils::variantToValue(*this, getLoc(), index); + auto insertOp = qtensor::InsertOp::create(*this, scalar, tensor, indexValue); + + auto outTensor = insertOp.getResult(); + + validateQubitValue(scalar); + validQubits.erase(scalar); + updateTensorTracking(tensor, outTensor); + return outTensor; +} + +Value QCOProgramBuilder::qtensorInsertSlice( + Value source, Value dest, const std::variant& offset, + const std::variant& size) { + checkFinalized(); + + auto offsetValue = utils::variantToValue(*this, getLoc(), offset); + auto sizeValue = utils::variantToValue(*this, getLoc(), size); + auto insertSliceOp = qtensor::InsertSliceOp::create(*this, source, dest, + offsetValue, sizeValue); + + auto outTensor = insertSliceOp.getResult(); + + validateTensorValue(source); + validTensors.erase(source); + updateTensorTracking(dest, outTensor); + + return outTensor; +} + +QCOProgramBuilder& QCOProgramBuilder::qtensorDealloc(Value tensor) { + checkFinalized(); + + validateTensorValue(tensor); + validTensors.erase(tensor); + + qtensor::DeallocOp::create(*this, tensor); + + return *this; +} + //===----------------------------------------------------------------------===// // Measurement and Reset //===----------------------------------------------------------------------===// @@ -776,22 +924,36 @@ OwningOpRef QCOProgramBuilder::finalize() { "Insertion point is not in entry block of main function"); } - // Automatically deallocate all still-allocated qubits - // Sort qubits for deterministic output - llvm::SmallVector sortedQubits(validQubits.begin(), validQubits.end()); - llvm::sort(sortedQubits, [](Value a, Value b) { + auto blockOrderComparator = [](Value a, Value b) { auto* opA = a.getDefiningOp(); auto* opB = b.getDefiningOp(); if (!opA || !opB || opA->getBlock() != opB->getBlock()) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); } return opA->isBeforeInBlock(opB); - }); + }; + + // Automatically deallocate all still-allocated qubits + // Sort qubits for deterministic output + llvm::SmallVector sortedQubits(validQubits.begin(), validQubits.end()); + llvm::sort(sortedQubits, blockOrderComparator); + for (auto qubit : sortedQubits) { DeallocOp::create(*this, qubit); } + // Automatically deallocate all still-allocated tensors + // Sort tensors for deterministic output + llvm::SmallVector sortedTensors(validTensors.begin(), + validTensors.end()); + llvm::sort(sortedTensors, blockOrderComparator); + + for (auto tensor : sortedTensors) { + qtensor::DeallocOp::create(*this, tensor); + } + validQubits.clear(); + validTensors.clear(); // Create constant 0 for successful exit code auto exitCode = intConstant(0); diff --git a/mlir/lib/Dialect/QTensor/CMakeLists.txt b/mlir/lib/Dialect/QTensor/CMakeLists.txt new file mode 100644 index 0000000000..b181a84fed --- /dev/null +++ b/mlir/lib/Dialect/QTensor/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +# All rights reserved. +# +# SPDX-License-Identifier: MIT +# +# Licensed under the MIT License + +add_subdirectory(IR) diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt new file mode 100644 index 0000000000..98aa36b0c5 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -0,0 +1,49 @@ +# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM +# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH +# All rights reserved. +# +# SPDX-License-Identifier: MIT +# +# Licensed under the MIT License + +file(GLOB_RECURSE OPERATIONS "${CMAKE_CURRENT_SOURCE_DIR}/Operations/*.cpp") + +add_mlir_dialect_library( + MLIRQTensorDialect + QTensorOps.cpp + ${OPERATIONS} + ADDITIONAL_HEADER_DIRS + ${MQT_MLIR_SOURCE_INCLUDE_DIR}/mlir/Dialect/QTensor + DEPENDS + MLIRQTensorOpsIncGen + LINK_LIBS + PRIVATE + MLIRIR + MLIRArithDialect + MLIRInferTypeOpInterface + MLIRSideEffectInterfaces + PUBLIC + MLIRQCODialect + DISABLE_INSTALL) + +mqt_mlir_target_use_project_options(MLIRQTensorDialect) + +# collect header files +file(GLOB_RECURSE IR_HEADERS_SOURCE "${MQT_MLIR_SOURCE_INCLUDE_DIR}/mlir/Dialect/QTensor/IR/*.h") +file(GLOB_RECURSE IR_HEADERS_BUILD "${MQT_MLIR_BUILD_INCLUDE_DIR}/mlir/Dialect/QTensor/IR/*.inc") + +# add public headers using file sets +target_sources( + MLIRQTensorDialect + PUBLIC FILE_SET + HEADERS + BASE_DIRS + ${MQT_MLIR_SOURCE_INCLUDE_DIR} + FILES + ${IR_HEADERS_SOURCE} + FILE_SET + HEADERS + BASE_DIRS + ${MQT_MLIR_BUILD_INCLUDE_DIR} + FILES + ${IR_HEADERS_BUILD}) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp new file mode 100644 index 0000000000..898b8b6412 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -0,0 +1,61 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace mlir; +using namespace mlir::qtensor; + +void AllocOp::build(OpBuilder& builder, OperationState& result, Value size) { + auto sizeValue = getConstantIntValue(size); + if (sizeValue) { + assert(*sizeValue > 0 && "qtensor.alloc size must be positive"); + } + + auto resultType = + RankedTensorType::get({sizeValue ? *sizeValue : ShapedType::kDynamic}, + qco::QubitType::get(builder.getContext())); + build(builder, result, resultType, size); +} + +LogicalResult AllocOp::verify() { + auto resultType = cast(getResult().getType()); + auto sizeValue = getConstantIntValue(getSize()); + auto resultSize = resultType.getShape()[0]; + + if (sizeValue && *sizeValue <= 0) { + return emitOpError("Constant size operand must be positive"); + } + if (sizeValue.has_value() == resultType.isDynamicDim(0)) { + return emitOpError("Size operand and result type must both be static or " + "both be dynamic, but got ") + << (sizeValue ? "static size with dynamic result" + : "dynamic size with static result"); + } + if (sizeValue && resultSize != *sizeValue) { + return emitOpError("Constant size operand (") + << *sizeValue << ") does not match static result size (" + << resultSize << ")"; + } + + return success(); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp new file mode 100644 index 0000000000..90f076ede1 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp @@ -0,0 +1,50 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +namespace { + +/** + * @brief Remove matching allocation and deallocation pairs without operations + * between them. + */ +struct RemoveAllocDeallocPair final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(DeallocOp op, + PatternRewriter& rewriter) const override { + // Check whether the tensor is directly defined by a qtensor::AllocOp. + auto tensor = op.getTensor(); + auto allocOp = tensor.getDefiningOp(); + if (!allocOp) { + return failure(); + } + + // Remove the AllocOp and the DeallocOp + rewriter.eraseOp(op); + rewriter.eraseOp(allocOp); + return success(); + } +}; + +} // namespace + +void DeallocOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp new file mode 100644 index 0000000000..27e8de6995 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +LogicalResult ExtractOp::verify() { + auto tensorDim = getTensor().getType().getDimSize(0); + auto index = getConstantIntValue(getIndex()); + + if (index) { + if (*index < 0) { + return emitOpError("Index must be non-negative"); + } + if (!ShapedType::isDynamic(tensorDim) && *index >= tensorDim) { + return emitOpError("Index exceeds tensor dimension"); + } + } + return success(); +} + +/** + * @brief If an ExtractOp consumes an InsertOp with the same index, + * return the scalar and the destTensor from the InsertOp directly. + */ +static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { + auto insertOp = extractOp.getTensor().getDefiningOp(); + if (!insertOp) { + return nullptr; + } + + Value insertIndex = insertOp.getIndex(); + Value extractIndex = extractOp.getIndex(); + + if (getAsOpFoldResult(insertIndex) != getAsOpFoldResult(extractIndex)) { + return nullptr; + } + + return insertOp; +} + +LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, + SmallVectorImpl& results) { + if (auto insertOp = foldExtractAfterInsert(*this)) { + results.emplace_back(insertOp.getDest()); + results.emplace_back(insertOp.getScalar()); + return success(); + } + + return failure(); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp new file mode 100644 index 0000000000..3c7648b04d --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -0,0 +1,103 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value tensor, + Value offset, Value size, + ArrayRef attrs) { + auto tensorType = cast(tensor.getType()); + auto sizeValue = getConstantIntValue(size); + auto resultType = RankedTensorType::get( + {sizeValue ? *sizeValue : ShapedType::kDynamic}, + tensorType.getElementType(), tensorType.getEncoding()); + + result.addAttributes(attrs); + build(b, result, {tensor.getType(), resultType}, tensor, offset, size); +} + +LogicalResult ExtractSliceOp::verify() { + auto tensorDim = getTensor().getType().getDimSize(0); + auto resultDim = getResult().getType().getDimSize(0); + auto constOffset = getConstantIntValue(getOffset()); + auto constSize = getConstantIntValue(getSize()); + + if (constOffset && *constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } + + if (constSize && *constSize <= 0) { + return emitOpError("Size must be positive"); + } + + if (constOffset && constSize && !ShapedType::isDynamic(tensorDim)) { + if (*constOffset + *constSize > tensorDim) { + return emitOpError("Offset + Size exceeds source dimension"); + } + } + + if (constSize && !ShapedType::isDynamic(resultDim)) { + if (resultDim != *constSize) { + return emitOpError("Result tensor dimension must match size operand"); + } + } + + return success(); +} + +/** + * @brief If an ExtractSliceOp consumes an InsertSliceOp with the same offset + * and size, return the sourceTensor and the destTensor from the InsertSliceOp + * directly. + */ +static InsertSliceOp +foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { + auto insertSliceOp = + extractSliceOp.getTensor().getDefiningOp(); + if (!insertSliceOp) { + return nullptr; + } + + auto insertOffset = insertSliceOp.getOffset(); + auto extractOffset = extractSliceOp.getOffset(); + auto insertSize = insertSliceOp.getSize(); + auto extractSize = extractSliceOp.getSize(); + + if (getAsOpFoldResult(insertOffset) != getAsOpFoldResult(extractOffset) || + getAsOpFoldResult(insertSize) != getAsOpFoldResult(extractSize)) { + return nullptr; + } + + return insertSliceOp; +} + +LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, + SmallVectorImpl& results) { + if (auto insertOp = foldExtractAfterInsertSlice(*this)) { + results.emplace_back(insertOp.getDest()); + results.emplace_back(insertOp.getSource()); + return success(); + } + + return failure(); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp new file mode 100644 index 0000000000..9d3b6a70f9 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include + +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +void FromElementsOp::build(OpBuilder& builder, OperationState& result, + ValueRange elements) { + assert(!elements.empty() && "Expected at least one element"); + auto resultType = RankedTensorType::get( + {static_cast(elements.size())}, elements.front().getType()); + build(builder, result, resultType, elements); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp new file mode 100644 index 0000000000..982d4a6335 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +LogicalResult InsertOp::verify() { + auto dstDim = getDest().getType().getDimSize(0); + auto index = getConstantIntValue(getIndex()); + + if (index) { + if (*index < 0) { + return emitOpError("Index must be non-negative"); + } + if (!ShapedType::isDynamic(dstDim) && *index >= dstDim) { + return emitOpError("Index exceeds tensor dimension"); + } + } + + return success(); +} + +/** + * @brief If an InsertOp consumes an ExtractOp with the same index, + * return the tensor from the extractOp directly. + */ +static Value foldInsertAfterExtract(InsertOp insertOp) { + auto extractOp = insertOp.getScalar().getDefiningOp(); + + if (!extractOp) { + return nullptr; + } + if (insertOp.getDest() != extractOp.getOutTensor()) { + return nullptr; + } + + auto insertIndex = insertOp.getIndex(); + auto extractIndex = extractOp.getIndex(); + + if (getAsOpFoldResult(insertIndex) != getAsOpFoldResult(extractIndex)) { + return nullptr; + } + + return extractOp.getTensor(); +} + +OpFoldResult InsertOp::fold(FoldAdaptor /*adaptor*/) { + if (auto result = foldInsertAfterExtract(*this)) { + return result; + } + + return {}; +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp new file mode 100644 index 0000000000..1a7b6526ab --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -0,0 +1,86 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +LogicalResult InsertSliceOp::verify() { + auto srcDim = getSource().getType().getDimSize(0); + auto dstDim = getDest().getType().getDimSize(0); + auto constOffset = getConstantIntValue(getOffset()); + auto constSize = getConstantIntValue(getSize()); + + if (constOffset && *constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } + + if (constSize && *constSize <= 0) { + return emitOpError("Size must be positive"); + } + + if (constSize && !ShapedType::isDynamic(srcDim)) { + if (*constSize != srcDim) { + return emitOpError("Size must match source dimension"); + } + } + + if (constOffset && constSize && !ShapedType::isDynamic(dstDim)) { + if (*constSize > dstDim || *constOffset > dstDim - *constSize) { + return emitOpError("Offset + Size exceeds destination dimension"); + } + } + + return success(); +} + +/** + * @brief If an InsertSliceOp consumes an ExtractSliceOp with the same offset + * and size, return the sourceTensor from the extractSliceOp directly. + */ +static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { + auto extractSliceOp = + insertSliceOp.getSource().getDefiningOp(); + if (!extractSliceOp) { + return nullptr; + } + + if (extractSliceOp.getOutTensor() != insertSliceOp.getDest()) { + return nullptr; + } + + auto insertOffset = insertSliceOp.getOffset(); + auto extractOffset = extractSliceOp.getOffset(); + auto insertSize = insertSliceOp.getSize(); + auto extractSize = extractSliceOp.getSize(); + + if (getAsOpFoldResult(insertOffset) != getAsOpFoldResult(extractOffset) || + getAsOpFoldResult(insertSize) != getAsOpFoldResult(extractSize)) { + return nullptr; + } + + return extractSliceOp.getTensor(); +} + +OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { + if (auto result = foldInsertAfterExtractSlice(*this)) { + return result; + } + + return {}; +} diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp new file mode 100644 index 0000000000..c26ed12d62 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -0,0 +1,57 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated + +// The following headers are needed for some template instantiations. +// IWYU pragma: begin_keep +#include +#include +// IWYU pragma: end_keep + +using namespace mlir; +using namespace mlir::qtensor; + +//===----------------------------------------------------------------------===// +// Dialect +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/QTensor/IR/QTensorOpsDialect.cpp.inc" + +void QTensorDialect::initialize() { + // NOLINTNEXTLINE(clang-analyzer-core.StackAddressEscape) + addTypes< +#define GET_TYPEDEF_LIST +#include "mlir/Dialect/QTensor/IR/QTensorOpsTypes.cpp.inc" + + >(); + + addOperations< +#define GET_OP_LIST +#include "mlir/Dialect/QTensor/IR/QTensorOps.cpp.inc" + + >(); +} + +//===----------------------------------------------------------------------===// +// Types +//===----------------------------------------------------------------------===// + +#define GET_TYPEDEF_CLASSES +#include "mlir/Dialect/QTensor/IR/QTensorOpsTypes.cpp.inc" + +//===----------------------------------------------------------------------===// +// Operations +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "mlir/Dialect/QTensor/IR/QTensorOps.cpp.inc" diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index e7b4853f34..ad4981d1b7 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -12,6 +12,7 @@ #include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h" #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Support/IRVerification.h" #include "mlir/Support/Passes.h" #include "qco_programs.h" @@ -57,7 +58,8 @@ class QCOTest : public testing::TestWithParam { void SetUp() override { // Register all necessary dialects DialectRegistry registry; - registry.insert(); + registry.insert(); context = std::make_unique(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); @@ -75,6 +77,7 @@ TEST_P(QCOTest, ProgramEquivalence) { ASSERT_TRUE(program); printer.record(program.get(), "Original QCO IR" + name); EXPECT_TRUE(verify(*program).succeeded()); + runCanonicalizationPasses(program.get()); printer.record(program.get(), "Canonicalized QCO IR" + name); EXPECT_TRUE(verify(*program).succeeded()); @@ -1060,3 +1063,56 @@ INSTANTIATE_TEST_SUITE_P( QCOTestCase{"AllocDeallocPair", MQT_NAMED_BUILDER(allocDeallocPair), MQT_NAMED_BUILDER(emptyQCO)})); /// @} + +/// \name QTensor/QTensor.cpp +/// @{ +INSTANTIATE_TEST_SUITE_P( + QTensorTest, QCOTest, + testing::Values( + QCOTestCase{"QTensorAlloc", MQT_NAMED_BUILDER(qtensorAlloc), + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{"QTensorAllocDealloc", MQT_NAMED_BUILDER(qtensorDealloc), + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{"QTensorFromElements", + MQT_NAMED_BUILDER(qtensorFromElements), + MQT_NAMED_BUILDER(qtensorFromElements)}, + QCOTestCase{"QTensorExtract", MQT_NAMED_BUILDER(qtensorExtract), + MQT_NAMED_BUILDER(qtensorExtract)}, + QCOTestCase{"QTensorInsert", MQT_NAMED_BUILDER(qtensorInsert), + MQT_NAMED_BUILDER(qtensorInsert)}, + QCOTestCase{"QTensorExtractSlice", + MQT_NAMED_BUILDER(qtensorExtractSlice), + MQT_NAMED_BUILDER(qtensorExtractSlice)}, + QCOTestCase{"QTensorInsertSlice", MQT_NAMED_BUILDER(qtensorInsertSlice), + MQT_NAMED_BUILDER(qtensorInsertSlice)}, + QCOTestCase{"QTensorExtractInsertSameIndex", + MQT_NAMED_BUILDER(qtensorExtractInsertSameIndex), + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{"QTensorExtractInsertIndexMismatch", + MQT_NAMED_BUILDER(qtensorExtractInsertIndexMismatch), + MQT_NAMED_BUILDER(qtensorExtractInsertIndexMismatch)}, + QCOTestCase{"QTensorInsertExtractSameIndex", + MQT_NAMED_BUILDER(qtensorInsertExtractSameIndex), + MQT_NAMED_BUILDER(qtensorInsert)}, + QCOTestCase{"QTensorInsertExtractIndexMismatch", + MQT_NAMED_BUILDER(qtensorInsertExtractIndexMismatch), + MQT_NAMED_BUILDER(qtensorInsertExtractIndexMismatch)}, + QCOTestCase{"QTensorExtractSliceInsertSliceSameOffset", + MQT_NAMED_BUILDER(qtensorExtractSliceInsertSliceSameOffset), + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{ + "QTensorExtractSliceInsertSliceOffsetMismatch", + MQT_NAMED_BUILDER(qtensorExtractSliceInsertSliceOffsetMismatch), + MQT_NAMED_BUILDER(qtensorExtractSliceInsertSliceOffsetMismatch)}, + QCOTestCase{"QTensorInsertSliceExtractSliceSameOffset", + MQT_NAMED_BUILDER(qtensorInsertSliceExtractSliceSameOffset), + MQT_NAMED_BUILDER(qtensorInsertSlice)}, + QCOTestCase{ + "QTensorInsertSliceExtractSliceOffsetMismatch", + MQT_NAMED_BUILDER(qtensorInsertSliceExtractSliceOffsetMismatch), + MQT_NAMED_BUILDER(qtensorInsertSliceExtractSliceOffsetMismatch)}, + QCOTestCase{ + "QTensorExtractSliceExtractInsertInsertSlice", + MQT_NAMED_BUILDER(qtensorExtractSliceExtractInsertInsertSlice), + MQT_NAMED_BUILDER(qtensorAlloc)})); +/// @} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 5c6a42bf1e..1ef16df6d0 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2133,4 +2133,127 @@ void nestedFalseIf(QCOProgramBuilder& b) { return llvm::to_vector(innerResult); }); } + +void qtensorAlloc(QCOProgramBuilder& b) { b.qtensorAlloc(3); } + +void qtensorDealloc(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + b.qtensorDealloc(qtensor); +} + +void qtensorFromElements(QCOProgramBuilder& b) { + auto q0 = b.allocQubit(); + auto q1 = b.allocQubit(); + auto q2 = b.allocQubit(); + b.qtensorFromElements({q0, q1, q2}); +} + +void qtensorExtract(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + b.qtensorExtract(qtensor, 0); +} + +void qtensorInsert(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); + auto q1 = b.h(q0); + b.qtensorInsert(q1, extractOutTensor, 0); +} + +void qtensorExtractSlice(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + b.qtensorExtractSlice(qtensor, 0, 2); +} + +void qtensorInsertSlice(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractSliceOutTensor, slicedTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + auto [extractOutTensor, q0] = b.qtensorExtract(slicedTensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.qtensorInsert(q1, extractOutTensor, 0); + b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); +} + +void qtensorExtractInsertIndexMismatch(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); + b.qtensorInsert(q0, extractOutTensor, 1); +} + +void qtensorExtractInsertSameIndex(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); + b.qtensorInsert(q0, extractOutTensor, 0); +} + +void qtensorExtractSliceInsertSliceOffsetMismatch(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractSliceOutTensor, slicedTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 1, 2); +} + +void qtensorExtractSliceInsertSliceSameOffset(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractSliceOutTensor, slicedTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 0, 2); +} + +void qtensorInsertExtractIndexMismatch(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.qtensorInsert(q1, extractOutTensor, 0); + auto [extractOutTensor1, q2] = b.qtensorExtract(insertOutTensor, 1); + b.qtensorInsert(q2, extractOutTensor1, 0); +} + +void qtensorInsertExtractSameIndex(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.qtensorInsert(q1, extractOutTensor, 0); + auto [extractOutTensor1, q2] = b.qtensorExtract(insertOutTensor, 0); + b.qtensorInsert(q2, extractOutTensor1, 0); +} + +void qtensorInsertSliceExtractSliceOffsetMismatch(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractSliceOutTensor, slicedTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + auto [extractOutTensor, q0] = b.qtensorExtract(slicedTensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.qtensorInsert(q1, extractOutTensor, 0); + auto insertSliceOutTensor = + b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); + auto [extractSliceOutTensor1, slicedTensor1] = + b.qtensorExtractSlice(insertSliceOutTensor, 1, 2); + b.qtensorInsertSlice(slicedTensor1, extractSliceOutTensor1, 0, 2); +} + +void qtensorInsertSliceExtractSliceSameOffset(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractSliceOutTensor, slicedTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + auto [extractOutTensor, q0] = b.qtensorExtract(slicedTensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.qtensorInsert(q1, extractOutTensor, 0); + auto insertSliceOutTensor = + b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); + auto [extractSliceOutTensor1, slicedTensor1] = + b.qtensorExtractSlice(insertSliceOutTensor, 0, 2); + b.qtensorInsertSlice(slicedTensor1, extractSliceOutTensor1, 0, 2); +} + +void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto [extractSliceOutTensor, slicedTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + auto [extractOutTensor, q0] = b.qtensorExtract(slicedTensor, 0); + auto insertOutTensor = b.qtensorInsert(q0, extractOutTensor, 0); + b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); +} + } // namespace mlir::qco diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index c0374325b2..ba26e42969 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -974,4 +974,63 @@ void nestedTrueIf(QCOProgramBuilder& b); /// Creates a circuit with a nested if operation in the else branch that uses /// the same condition. void nestedFalseIf(QCOProgramBuilder& b); + +// --- QTensor Operations -------------------------------------------------- // + +/// Allocates a tensor of size `3`. +void qtensorAlloc(QCOProgramBuilder& b); + +/// Allocates and explicitly deallocates a tensor. +void qtensorDealloc(QCOProgramBuilder& b); + +/// Constructs a tensor with from_elements. +void qtensorFromElements(QCOProgramBuilder& b); + +/// Extracts a qubit from a tensor. +void qtensorExtract(QCOProgramBuilder& b); + +/// Inserts a qubit into a tensor. +void qtensorInsert(QCOProgramBuilder& b); + +/// Extracts a slice from a tensor. +void qtensorExtractSlice(QCOProgramBuilder& b); + +/// Inserts a slice into a tensor. +void qtensorInsertSlice(QCOProgramBuilder& b); + +/// Extracts a qubit from a tensor and inserts it immediately at a different +/// index. +void qtensorExtractInsertIndexMismatch(QCOProgramBuilder& b); + +/// Extracts a qubit from a tensor and inserts it immediately at the same index. +void qtensorExtractInsertSameIndex(QCOProgramBuilder& b); + +/// Inserts a qubit into a tensor and extracts it immediately at a different +/// index. +void qtensorInsertExtractIndexMismatch(QCOProgramBuilder& b); + +/// Inserts a qubit into a tensor and extracts it immediately at the same index. +void qtensorInsertExtractSameIndex(QCOProgramBuilder& b); + +/// Extracts a slice of qubits from a tensor and inserts it immediately at a +/// different offset. +void qtensorExtractSliceInsertSliceOffsetMismatch(QCOProgramBuilder& b); + +/// Extracts a slice of qubits from a tensor and inserts it immediately at the +/// same offset. +void qtensorExtractSliceInsertSliceSameOffset(QCOProgramBuilder& b); + +/// Inserts a slice of qubits into a tensor and extracts it immediately at a +/// different offset. +void qtensorInsertSliceExtractSliceOffsetMismatch(QCOProgramBuilder& b); + +/// Inserts a slice of qubits into a tensor and extracts it immediately at the +/// same offset. +void qtensorInsertSliceExtractSliceSameOffset(QCOProgramBuilder& b); + +/// Extracts a slice of qubits, extracts a qubit from the slice, inserts the +/// qubit back into the slice, and inserts the slice back into the tensor +/// immediately at the same index and offset. +void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b); + } // namespace mlir::qco