From 7e259404dc17b66ad3cf30b159eb0e6e8695dbef Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 26 Feb 2026 16:21:47 +0100 Subject: [PATCH 001/108] add base qtensor implementation --- mlir/include/mlir/Dialect/CMakeLists.txt | 1 + mlir/include/mlir/Dialect/QC/IR/QCOps.td | 2 +- .../mlir/Dialect/QTensor/CMakeLists.txt | 9 + .../mlir/Dialect/QTensor/IR/CMakeLists.txt | 10 + .../mlir/Dialect/QTensor/IR/QTensorBase.td | 39 + .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 51 ++ .../mlir/Dialect/QTensor/IR/QTensorOps.h | 32 + .../mlir/Dialect/QTensor/IR/QTensorOps.td | 234 ++++++ mlir/lib/Conversion/QCOToQC/CMakeLists.txt | 2 + mlir/lib/Conversion/QCToQCO/CMakeLists.txt | 1 + mlir/lib/Dialect/CMakeLists.txt | 1 + mlir/lib/Dialect/QCO/Builder/CMakeLists.txt | 1 + mlir/lib/Dialect/QTensor/CMakeLists.txt | 9 + mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 46 ++ mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 671 ++++++++++++++++++ 15 files changed, 1108 insertions(+), 1 deletion(-) create mode 100644 mlir/include/mlir/Dialect/QTensor/CMakeLists.txt create mode 100644 mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt create mode 100644 mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td create mode 100644 mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h create mode 100644 mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h create mode 100644 mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td create mode 100644 mlir/lib/Dialect/QTensor/CMakeLists.txt create mode 100644 mlir/lib/Dialect/QTensor/IR/CMakeLists.txt create mode 100644 mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp 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/QC/IR/QCOps.td b/mlir/include/mlir/Dialect/QC/IR/QCOps.td index 5956e39a56..035067bf3a 100644 --- a/mlir/include/mlir/Dialect/QC/IR/QCOps.td +++ b/mlir/include/mlir/Dialect/QC/IR/QCOps.td @@ -60,7 +60,7 @@ class QCType traits = []> let mnemonic = typeMnemonic; } -def QubitType : QCType<"Qubit", "qubit"> { +def QubitType : QCType<"Qubit", "qubit", [MemRefElementTypeInterface]> { let summary = "QC qubit reference type"; let description = [{ The `!qc.qubit` type represents a reference to a quantum bit in the 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..81bd22fbc9 --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt @@ -0,0 +1,10 @@ +# 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) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td new file mode 100644 index 0000000000..04df4da1ad --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td @@ -0,0 +1,39 @@ +// 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 + +//===- TensorBase.td - Base definitions for tensor dialect -*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef TENSOR_BASE +#define TENSOR_BASE + +include "mlir/IR/OpBase.td" + +def QTensorDialect : Dialect { + let name = "qtensor"; + let cppNamespace = "::mlir::qtensor"; + + let description = [{ + TODO + }]; + + let hasCanonicalizer = 1; + let hasConstantMaterializer = 1; + let dependentDialects = [ + "affine::AffineDialect", + "arith::ArithDialect", + "complex::ComplexDialect", + ]; +} + +#endif // TENSOR_BASE 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..122c683c2f --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -0,0 +1,51 @@ +/* + * 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 "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/ViewLikeInterface.h" + +#define DIALECT_NAME_QTensor "qtensor" + +//===----------------------------------------------------------------------===// +// QTensor Dialect Helpers +//===----------------------------------------------------------------------===// + +namespace mlir { + +/// Return the list of Range (i.e. offset, size, stride). Each Range +/// entry contains either the dynamic value or a ConstantIndexOp constructed +/// with `b` at location `loc`. +SmallVector getOrCreateRanges(OffsetSizeAndStrideOpInterface op, + OpBuilder& b, Location loc); + +} // namespace mlir + +//===----------------------------------------------------------------------===// +// 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 + +//===----------------------------------------------------------------------===// +// QTensor Dialect Helpers +//===----------------------------------------------------------------------===// + +namespace mlir::qtensor {} // namespace mlir::qtensor 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..75fc520754 --- /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) +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wambiguous-reversed-operator" +#endif +#include +#include +#ifdef __clang__ +#pragma clang diagnostic pop +#endif + +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" + +#include +#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..bd2d077a68 --- /dev/null +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -0,0 +1,234 @@ +// 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 + +//===- QTensorOps.td - QTensor op definitions ----------------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef QTENSOROPS +#define QTENSOROPS + +include "mlir/Dialect/QTensor/IR/QTensorBase.td" +include "mlir/Interfaces/CastInterfaces.td" +include "mlir/Interfaces/ControlFlowInterfaces.td" +include "mlir/Interfaces/DestinationStyleOpInterface.td" +include "mlir/Interfaces/InferIntRangeInterface.td" +include "mlir/Interfaces/InferTypeOpInterface.td" +include "mlir/Interfaces/ParallelCombiningOpInterface.td" +include "mlir/Interfaces/ShapedOpInterfaces.td" +include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/TilingInterface.td" +include "mlir/Interfaces/ViewLikeInterface.td" +include "mlir/IR/OpAsmInterface.td" + +class QTensorOp traits = []> + : Op; + +// Base class for ops with static/dynamic offset, sizes and strides +// attributes/arguments. +class QTensorOpWithOffsetSizesAndStrides traits = []> + : QTensorOp { + code extraBaseClassDeclaration = [{ + /// Return the type of the base tensor operand. + ::mlir::RankedTensorType getSourceType() { + return ::llvm::cast(getSource().getType()); + } + + /// Return the type of the result tensor. + ::mlir::RankedTensorType getResultType() { + return ::llvm::cast(getResult().getType()); + } + + /// Return the dynamic sizes for this subview operation if specified. + ::mlir::Operation::operand_range getDynamicSizes() { return getSizes(); } + + }]; +} +def QTensor_ExtractOp : QTensorOp<"extract", [ + DeclareOpInterfaceMethods, + 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 = "element extraction operation"; + let description = [{ + The `tensor.extract` op reads a ranked tensor and returns one element as + specified by the given indices. The result of the op is a value with the + same type as the elements of the tensor. The arity of indices must match + the rank of the accessed value. All indices should all be of `index` type. + + Example: + + ```mlir + %4 = tensor.extract %t[%1, %2] : tensor<4x4xi32> + %5 = tensor.extract %rt[%1, %2] : tensor + ``` + }]; + + let arguments = (ins AnyRankedTensor:$tensor, Variadic:$indices); + let results = (outs AnyType:$result, AnyRankedTensor:$out_tensor); + let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)"; + + let hasFolder = 1; + let hasVerifier = 1; +} + + +//===----------------------------------------------------------------------===// +// ExtractSliceOp +//===----------------------------------------------------------------------===// + +def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", [ + DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, + AttrSizedOperandSegments, + Pure, + OffsetSizeAndStrideOpInterface, + TypesMatchWith<"returned tensor type matches input tensor", + "source", "out_source", "$_self"> + ]> { + let summary = "extract slice operation"; + let description = [{ + + }]; + + let arguments = (ins + AnyRankedTensor:$source, + Variadic:$offsets, + Variadic:$sizes, + Variadic:$strides, + DenseI64ArrayAttr:$static_offsets, + DenseI64ArrayAttr:$static_sizes, + DenseI64ArrayAttr:$static_strides + ); + let results = (outs AnyRankedTensor:$result, AnyRankedTensor:$out_source); + + let assemblyFormat = [{ + $source `` + custom($offsets, $static_offsets) + custom($sizes, $static_sizes) + custom($strides, $static_strides) + attr-dict `:` type($source) `to` type($result) + }]; + + let builders = [ + // Build an ExtractSliceOp with mixed static and dynamic entries and + // inferred result type. + OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, + "ArrayRef":$sizes, "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with mixed static and dynamic entries and custom + // result type. If the type passed is nullptr, it is inferred. + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with dynamic entries and custom result type. If + // the type passed is nullptr, it is inferred. + OpBuilder<(ins "Value":$source, "ValueRange":$offsets, + "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with dynamic entries and inferred result type. + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with mixed static and dynamic entries packed in + // a Range vector. + OpBuilder<(ins "Value":$source, "ArrayRef":$ranges, + CArg<"ArrayRef", "{}">:$attrs)>, + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + /// The result of an extract_slice is always a tensor. + // TODO: deprecate + RankedTensorType getType() { + return getResultType(); + } + + /// Compute the rank-reduction mask that can be applied to map the source + /// tensor type to the result tensor type by dropping unit dims. + std::optional> + computeRankReductionMask() { + return ::mlir::computeRankReductionMask(getSourceType().getShape(), + getType().getShape()); + }; + + /// An extract_slice result type can be inferred, when it is not + /// rank-reduced, from the source type and the static representation of + /// offsets, sizes and strides. Special sentinels encode the dynamic case. + static RankedTensorType inferResultType( + RankedTensorType sourceTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + static RankedTensorType inferResultType( + RankedTensorType sourceTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + + /// If the rank is reduced (i.e. the desiredResultRank is smaller than the + /// number of sizes), drop as many size 1 as needed to produce an inferred type + /// with the desired rank. + /// + /// Note that there may be multiple ways to compute this rank-reduced type: + /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. + /// + /// To disambiguate, this function always drops the first 1 sizes occurrences. + static RankedTensorType inferCanonicalRankReducedResultType( + unsigned resultRank, + RankedTensorType sourceRankedTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + static RankedTensorType inferCanonicalRankReducedResultType( + unsigned resultRank, + RankedTensorType sourceRankedTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + + /// Return the expected rank of each of the`static_offsets`, `static_sizes` + /// and `static_strides` attributes. + std::array getArrayAttrMaxRanks() { + unsigned rank = getSourceType().getRank(); + return {rank, rank, rank}; + } + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } + + /// Return the dimensions of the source that are dropped in the + /// result when the result is rank-reduced. + llvm::SmallBitVector getDroppedDims(); + + /// Given a `value`, asserted to be of RankedTensorType, build an + /// ExtractSliceOp that results in a rank-reducing extract to the desired + /// tensor shape and return the new value created. + /// If the shape of `value` is already the `desiredShape`, just return + /// `value`. + /// If the shape of `value` cannot be rank-reduced to `desiredShape`, fail. + static FailureOr rankReduceIfNeeded( + OpBuilder &b, Location loc, Value value, ArrayRef desiredShape); + }]; + + let hasCanonicalizer = 1; + let hasFolder = 1; + let hasVerifier = 1; +} + +#endif // TENSOR_OPS diff --git a/mlir/lib/Conversion/QCOToQC/CMakeLists.txt b/mlir/lib/Conversion/QCOToQC/CMakeLists.txt index c50ca147fe..8b85d395bf 100644 --- a/mlir/lib/Conversion/QCOToQC/CMakeLists.txt +++ b/mlir/lib/Conversion/QCOToQC/CMakeLists.txt @@ -16,9 +16,11 @@ add_mlir_library( LINK_LIBS MLIRQCDialect MLIRQCODialect + MLIRQTensorDialect MLIRArithDialect MLIRFuncDialect MLIRTransforms + MLIRViewLikeInterface MLIRFuncTransforms DISABLE_INSTALL) diff --git a/mlir/lib/Conversion/QCToQCO/CMakeLists.txt b/mlir/lib/Conversion/QCToQCO/CMakeLists.txt index 80a07eda6d..8cb17e8a5e 100644 --- a/mlir/lib/Conversion/QCToQCO/CMakeLists.txt +++ b/mlir/lib/Conversion/QCToQCO/CMakeLists.txt @@ -16,6 +16,7 @@ add_mlir_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/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..97a3088c2d --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -0,0 +1,46 @@ +# 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_library( + MLIRQTensorDialect + QTensorOps.cpp + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/mlir/include/mlir/Dialect/QTensor + DEPENDS + MLIRQTensorOpsIncGen + LINK_LIBS + PUBLIC + MLIRIR + MLIRArithDialect + MLIRViewLikeInterface + MLIRInferTypeOpInterface + MLIRSideEffectInterfaces + Eigen3::Eigen + 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/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp new file mode 100644 index 0000000000..2e143562bb --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -0,0 +1,671 @@ +/* + * 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 + */ + +//===----------------------------------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include "mlir/Dialect/Affine/IR/AffineOps.h" // for affine::AffineDialect +#include "mlir/Dialect/Arith/IR/Arith.h" // for arith::ArithDialect +#include "mlir/Dialect/Complex/IR/Complex.h" // for complex::ComplexDialect +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributeInterfaces.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/TypeUtilities.h" +#include "mlir/Interfaces/DestinationStyleOpInterface.h" +#include "mlir/Interfaces/InferIntRangeInterface.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Interfaces/Utils/InferIntRangeCommon.h" +#include "mlir/Interfaces/ViewLikeInterface.h" +#include "mlir/Support/LLVM.h" + +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallBitVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/MathExtras.h" + +#include +#include +#include +#include + +// 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; + +/// Materialize a single constant operation from a given attribute value with +/// the desired resultant type. +Operation* QTensorDialect::materializeConstant(OpBuilder& builder, + Attribute value, Type type, + Location loc) { + if (auto op = arith::ConstantOp::materialize(builder, value, type, loc)) { + return op; + } + if (complex::ConstantOp::isBuildableWith(value, type)) { + return builder.create(loc, type, + llvm::cast(value)); + } + return nullptr; +} + +/// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or +/// rank-extending tensor.insert_slice op. +static llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, + ArrayRef mixedSizes) { + llvm::SmallBitVector droppedDims(mixedSizes.size()); + int64_t shapePos = reducedShape.size() - 1; + + for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { + size_t idx = mixedSizes.size() - size.index() - 1; + // Rank-reduced dims must have a static unit dimension. + bool isStaticUnitSize = + isa(size.value()) && + llvm::cast(cast(size.value())).getInt() == 1; + + if (shapePos < 0) { + // There are no more dims in the reduced shape. All remaining sizes must + // be rank-reduced dims. + assert(isStaticUnitSize && "expected unit dim"); + droppedDims.set(idx); + continue; + } + + // Dim is preserved if the size is not a static 1. + if (!isStaticUnitSize) { + --shapePos; + continue; + } + + // Dim is preserved if the reduced shape dim is also 1. + if (reducedShape[shapePos] == 1) { + --shapePos; + continue; + } + + // Otherwise: Dim is dropped. + droppedDims.set(idx); + } + + assert(shapePos < 0 && "dimension mismatch"); + return droppedDims; +} + +//===----------------------------------------------------------------------===// +// ExtractOp +//===----------------------------------------------------------------------===// + +void ExtractOp::getAsmResultNames( + function_ref setNameFn) { + setNameFn(getResult(), "q_extracted"); +} + +LogicalResult ExtractOp::verify() { + // Verify the # indices match if we have a ranked type. + auto tensorType = llvm::cast(getTensor().getType()); + if (tensorType.getRank() != static_cast(getIndices().size())) { + return emitOpError("incorrect number of indices for extract_element"); + } + return success(); +} + +/// If we have an ExtractOp consuming an InsertOp with the same +/// indices, we can return the InsertOp's scalar directly. +// TODO: This only checks the immediate producer; extend to go up the +// insert/extract chain if the slices are disjoint. +static Value foldExtractAfterInsert(ExtractOp extractOp) { + auto insertOp = extractOp.getTensor().getDefiningOp(); + + auto isSame = [](Value a, Value b) { + return getAsOpFoldResult(a) == getAsOpFoldResult(b); + }; + if (insertOp && insertOp.getScalar().getType() == extractOp.getType(0) && + llvm::equal(insertOp.getIndices(), extractOp.getIndices(), isSame)) { + return insertOp.getScalar(); + } + return {}; +} + +LogicalResult ExtractOp::fold(FoldAdaptor adaptor, + SmallVectorImpl& results) { + // Collect the constant indices into the tensor. + SmallVector indices; + for (Attribute indice : adaptor.getIndices()) { + if (!indice || !llvm::isa(indice)) { + return failure(); + } + indices.push_back(llvm::cast(indice).getInt()); + } + + // Fold extract(from_elements(...)). + if (auto fromElementsOp = + getTensor().getDefiningOp()) { + auto tensorType = llvm::cast(fromElementsOp.getType()); + auto rank = tensorType.getRank(); + assert(static_cast(indices.size()) == tensorType.getRank() && + "rank mismatch"); + int flatIndex = 0; + int stride = 1; + for (int i = rank - 1; i >= 0; --i) { + flatIndex += indices[i] * stride; + stride *= tensorType.getDimSize(i); + } + // Prevent out of bounds accesses. This can happen in invalid code that + // will never execute. + if (static_cast(fromElementsOp.getElements().size()) <= flatIndex || + flatIndex < 0) { + return failure(); + } + results.push_back(fromElementsOp.getElements()[flatIndex]); + results.push_back(getTensor()); + return success(); + } + + if (Value result = foldExtractAfterInsert(*this)) { + results.push_back(result); + results.push_back(getTensor()); + return success(); + } + + return failure(); +} + +//===----------------------------------------------------------------------===// +// ExtractSliceOp +//===----------------------------------------------------------------------===// + +void ExtractSliceOp::getAsmResultNames( + function_ref setNameFn) { + setNameFn(getResult(), "q_extracted_slice"); +} + +/// An extract_slice result type can be inferred, when it is not +/// rank-reduced, from the source type and the static representation of +/// offsets, sizes and strides. Special sentinels encode the dynamic case. +RankedTensorType ExtractSliceOp::inferResultType( + RankedTensorType sourceTensorType, ArrayRef staticOffsets, + ArrayRef staticSizes, ArrayRef staticStrides) { + // An extract_slice op may specify only a leading subset of offset/sizes/ + // strides in which case we complete with offset=0, sizes from memref type + // and strides=1. + assert(static_cast(staticSizes.size()) == + sourceTensorType.getRank() && + "unexpected staticSizes not equal to rank of source"); + return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), + sourceTensorType.getEncoding()); +} + +RankedTensorType ExtractSliceOp::inferResultType( + RankedTensorType sourceTensorType, ArrayRef offsets, + ArrayRef sizes, ArrayRef strides) { + SmallVector staticSizes; + std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes); + assert(static_cast(staticSizes.size()) == + sourceTensorType.getRank() && + "unexpected staticSizes not equal to rank of source"); + return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), + sourceTensorType.getEncoding()); +} + +/// If the rank is reduced (i.e. the desiredResultRank is smaller than the +/// number of sizes), drop as many size 1 as needed to produce an inferred +/// type with the desired rank. +/// +/// Note that there may be multiple ways to compute this rank-reduced type: +/// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. +/// +/// To disambiguate, this function always drops the first 1 sizes occurrences. +RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( + unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides) { + // Type inferred in the absence of rank-reducing behavior. + auto inferredType = llvm::cast( + inferResultType(sourceRankedTensorType, offsets, sizes, strides)); + int rankDiff = inferredType.getRank() - desiredResultRank; + if (rankDiff > 0) { + auto shape = inferredType.getShape(); + llvm::SmallBitVector dimsToProject = + getPositionsOfShapeOne(rankDiff, shape); + SmallVector projectedShape; + // Best effort rank-reducing: drop 1s in order. + for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) { + if (!dimsToProject.test(pos)) { + projectedShape.push_back(shape[pos]); + } + } + inferredType = + RankedTensorType::get(projectedShape, inferredType.getElementType()); + } + return inferredType; +} + +RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( + unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides) { + SmallVector staticOffsets; + SmallVector staticSizes; + SmallVector staticStrides; + SmallVector dynamicOffsets; + SmallVector dynamicSizes; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + return ExtractSliceOp::inferCanonicalRankReducedResultType( + desiredResultRank, sourceRankedTensorType, staticOffsets, staticSizes, + staticStrides); +} + +/// Build an ExtractSliceOp with mixed static and dynamic entries and custom +/// result type. If the type passed is nullptr, it is inferred. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + SmallVector staticOffsets; + SmallVector staticSizes; + SmallVector staticStrides; + SmallVector dynamicOffsets; + SmallVector dynamicSizes; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + auto sourceRankedTensorType = llvm::cast(source.getType()); + // Structuring implementation this way avoids duplication between builders. + if (!resultType) { + resultType = llvm::cast(ExtractSliceOp::inferResultType( + sourceRankedTensorType, staticOffsets, staticSizes, staticStrides)); + } + result.addAttributes(attrs); + build(b, result, resultType, source, dynamicOffsets, dynamicSizes, + dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), + b.getDenseI64ArrayAttr(staticSizes), + b.getDenseI64ArrayAttr(staticStrides)); +} + +/// Build an ExtractSliceOp with mixed static and dynamic entries and inferred +/// result type. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); +} + +/// Build an ExtractSliceOp with mixed static and dynamic entries packed into +/// a Range vector. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, + ArrayRef ranges, + ArrayRef attrs) { + auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); + build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); +} + +/// Build an ExtractSliceOp with dynamic entries and custom result type. If +/// the type passed is nullptr, it is inferred. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + SmallVector offsetValues = llvm::to_vector<4>( + llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); + SmallVector sizeValues = llvm::to_vector<4>( + llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = llvm::to_vector<4>( + llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + build(b, result, resultType, source, offsetValues, sizeValues, strideValues); +} + +/// Build an ExtractSliceOp with dynamic entries and inferred result type. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); +} + +static LogicalResult produceSliceErrorMsg(SliceVerificationResult result, + Operation* op, + RankedTensorType expectedType) { + switch (result) { + case SliceVerificationResult::Success: + return success(); + case SliceVerificationResult::RankTooLarge: + return op->emitError("expected rank to be smaller or equal to ") + << "the other rank. "; + case SliceVerificationResult::SizeMismatch: + return op->emitError("expected type to be ") + << expectedType << " or a rank-reduced version. (size mismatch) "; + case SliceVerificationResult::ElemTypeMismatch: + return op->emitError("expected element type to be ") + << expectedType.getElementType(); + default: + llvm_unreachable("unexpected extract_slice op verification result"); + } +} + +/// Verifier for ExtractSliceOp. +LogicalResult ExtractSliceOp::verify() { + RankedTensorType sourceType = getSourceType(); + + // Verify result type against inferred type. + RankedTensorType expectedType = ExtractSliceOp::inferResultType( + sourceType, getMixedOffsets(), getMixedSizes(), getMixedStrides()); + SliceVerificationResult result = isRankReducedType(expectedType, getType()); + if (result != SliceVerificationResult::Success) { + return produceSliceErrorMsg(result, *this, expectedType); + } + + // Verify that offsets, sizes, strides do not run out-of-bounds with respect + // to the source tensor. + SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( + sourceType.getShape(), getStaticOffsets(), getStaticSizes(), + getStaticStrides(), /*generateErrorMessage=*/true); + if (!boundsResult.isValid) { + return getOperation()->emitError(boundsResult.errorMessage); + } + + return success(); +} + +llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { + return ::getDroppedDims(getType().getShape(), getMixedSizes()); +} + +FailureOr +ExtractSliceOp::rankReduceIfNeeded(OpBuilder& b, Location loc, Value value, + ArrayRef desiredShape) { + auto sourceTensorType = llvm::dyn_cast(value.getType()); + assert(sourceTensorType && "not a ranked tensor type"); + auto sourceShape = sourceTensorType.getShape(); + if (sourceShape.equals(desiredShape)) { + return value; + } + auto maybeRankReductionMask = + mlir::computeRankReductionMask(sourceShape, desiredShape); + if (!maybeRankReductionMask) { + return failure(); + } + return tensor::createCanonicalRankReducingExtractSliceOp( + b, loc, value, + RankedTensorType::Builder(sourceTensorType).setShape(desiredShape)); +} + +LogicalResult ExtractSliceOp::reifyResultShapes( + OpBuilder& builder, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { + reifiedReturnShapes.resize(1); + reifiedReturnShapes[0].reserve(getType().getRank()); + SmallVector mixedSizes = getMixedSizes(); + llvm::SmallBitVector droppedDims = getDroppedDims(); + for (const auto& size : enumerate(mixedSizes)) { + if (droppedDims.test(size.index())) { + continue; + } + reifiedReturnShapes[0].push_back(size.value()); + } + return success(); +} + +namespace { +/// Pattern to rewrite an extract_slice op with tensor::Cast arguments. +/// This essentially pushes memref_cast past its consuming slice when +/// `canFoldIntoConsumerOp` is true. +/// +/// Example: +/// ``` +/// %0 = tensor.cast %V : tensor<16x16xf32> to tensor +/// %1 = tensor.extract_slice %0[0, 0][3, 4][1, 1] : tensor to +/// tensor<3x4xf32> +/// ``` +/// is rewritten into: +/// ``` +/// %0 = tensor.extract_slice %V[0, 0][3, 4][1, 1] : tensor<16x16xf32> to +/// tensor<3x4xf32> %1 = tensor.cast %0: tensor<3x4xf32> to tensor<3x4xf32> +/// ``` +class ExtractSliceOpCastFolder final : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(ExtractSliceOp sliceOp, + PatternRewriter& rewriter) const override { + // Any constant operand, just return to let the constant folder kick in. + if (llvm::any_of(sliceOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) { + return failure(); + } + + auto castOp = sliceOp.getSource().getDefiningOp(); + if (!castOp) { + return failure(); + } + + if (!canFoldIntoConsumerOp(castOp)) { + return failure(); + } + + // Pattern does not apply if the produced op would not verify. + SliceBoundsVerificationResult sliceResult = verifyInBoundsSlice( + cast(castOp.getSource().getType()).getShape(), + sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), + sliceOp.getStaticStrides()); + if (!sliceResult.isValid) { + return failure(); + } + + // Create folded extract. + Location loc = sliceOp.getLoc(); + auto newResult = rewriter.create( + loc, sliceOp.getType(), castOp.getSource(), sliceOp.getOffsets(), + sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(), + sliceOp.getStaticSizes(), sliceOp.getStaticStrides()); + rewriter.replaceOp(sliceOp, newResult->getResult(0)); + rewriter.replaceOp(castOp, newResult->getResult(1)); + return success(); + } +}; + +/// Slice elements from `values` into `outValues`. `counts` represents the +/// numbers of elements to stride in the original values for each dimension. +/// The output values can be used to construct a DenseElementsAttr. +template +static void sliceElements(IterTy values, ArrayRef counts, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides, + llvm::SmallVectorImpl* outValues) { + assert(offsets.size() == sizes.size()); + assert(offsets.size() == strides.size()); + if (offsets.empty()) { + return; + } + + int64_t offset = offsets.front(); + int64_t size = sizes.front(); + int64_t stride = strides.front(); + if (offsets.size() == 1) { + for (int64_t i = 0; i < size; ++i, offset += stride) { + outValues->push_back(*(values + offset)); + } + + return; + } + + for (int64_t i = 0; i < size; ++i, offset += stride) { + auto begin = values + offset * counts.front(); + sliceElements(begin, counts.drop_front(), + offsets.drop_front(), sizes.drop_front(), + strides.drop_front(), outValues); + } +} + +} // namespace + +/// Return the canonical type of the result of an extract_slice op. +struct SliceReturnTypeCanonicalizer { + RankedTensorType operator()(ExtractSliceOp op, + ArrayRef mixedOffsets, + ArrayRef mixedSizes, + ArrayRef mixedStrides) { + return ExtractSliceOp::inferCanonicalRankReducedResultType( + op.getType().getRank(), op.getSourceType(), mixedOffsets, mixedSizes, + mixedStrides); + } +}; + +/// A canonicalizer wrapper to replace ExtractSliceOps. +struct SliceCanonicalizer { + void operator()(PatternRewriter& rewriter, ExtractSliceOp op, + ExtractSliceOp newOp) { + Value replacement = newOp.getResult(); + if (replacement.getType() != op.getType()) { + replacement = rewriter.create(op.getLoc(), op.getType(), + replacement); + } + rewriter.replaceOp(op, replacement); + } +}; + +void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add< + OpWithOffsetSizesAndStridesConstantArgumentFolder< + ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>, + ExtractSliceOpCastFolder>(context); +} + +// +static LogicalResult +foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, + ShapedType shapedType) { + OpBuilder b(op.getContext()); + for (OpFoldResult opFold : op.getMixedOffsets()) { + if (getConstantIntValue(opFold) != static_cast(0)) { + return failure(); + } + } + // Rank-reducing noops only need to inspect the leading dimensions: + // llvm::zip is appropriate. + auto shape = shapedType.getShape(); + for (auto it : llvm::zip(op.getMixedSizes(), shape)) { + if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) { + return failure(); + } + } + for (OpFoldResult opFold : op.getMixedStrides()) { + if (getConstantIntValue(opFold) != static_cast(1)) { + return failure(); + } + } + return success(); +} + +/// If we have an ExtractSliceOp consuming an InsertSliceOp with the same +/// slice, we can return the InsertSliceOp's source directly. +// TODO: This only checks the immediate producer; extend to go up the +// insert/extract chain if the slices are disjoint. +static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { + auto insertOp = extractOp.getSource().getDefiningOp(); + + auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; + if (insertOp && insertOp.getSource().getType() == extractOp.getType() && + insertOp.isSameAs(extractOp, isSame)) { + return insertOp.getSource(); + } + + return {}; +} + +LogicalResult ExtractSliceOp::fold(FoldAdaptor adaptor, + SmallVectorImpl& results) { + + if (getSourceType() == getType() && + succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { + results.push_back(this->getSource()); + return success(); + } + if (Value slice = foldExtractAfterInsertSlice(*this)) { + results.push_back(slice); + return success(); + } + + return failure(); +} + +Value mlir::tensor::createCanonicalRankReducingExtractSliceOp( + OpBuilder& b, Location loc, Value tensor, RankedTensorType targetType) { + auto rankedTensorType = llvm::cast(tensor.getType()); + unsigned rank = rankedTensorType.getRank(); + SmallVector offsets(rank, b.getIndexAttr(0)); + SmallVector sizes = getMixedSizes(b, loc, tensor); + SmallVector strides(rank, b.getIndexAttr(1)); + return b.createOrFold(loc, targetType, tensor, + offsets, sizes, strides); +} + +void QTensorDialect::getCanonicalizationPatterns( + RewritePatternSet& results) const { + // results.add(getContext()); +} +//===----------------------------------------------------------------------===// +// 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" From 8f65b5ab026eb60fced4d72930131d7e14df5b8b Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 26 Feb 2026 17:14:18 +0100 Subject: [PATCH 002/108] add additional operations --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 125 +++++++++++++++ mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 145 ++++++++++++++++++ 2 files changed, 270 insertions(+) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index bd2d077a68..a7e91590d0 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -231,4 +231,129 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", let hasVerifier = 1; } +def QTensor_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 = "tensor from elements operation."; + let description = [{ + Create a N-D tensor from a range of same-type arguments. The number of + provided `elements` should equal to the number of the elements in the + result type. The `elements` correspond to a flattened tensor. + + Example: + + ```mlir + tensor.from_elements %a, %b, %c, %d, %e, %f : tensor<2x3xindex> + ``` + + will result in a tensor + + [[%a, %b, %c] + [%d, %e, %f]] + }]; + + let arguments = (ins Variadic:$elements); + let results = (outs AnyStaticShapeTensor:$result); + + let assemblyFormat = "$elements attr-dict `:` type($result)"; + + let builders = [ + // Special case builder for when `elements` has size >=1. + OpBuilder<(ins "ValueRange":$elements)> + ]; + + let hasCanonicalizer = 1; +} + +def QTensor_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 = "element insertion operation"; + let description = [{ + }]; + + let arguments = (ins AnyType:$scalar, + AnyRankedTensor:$dest, + Variadic:$indices); + let results = (outs AnyRankedTensor:$result); + let assemblyFormat = [{ + $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) + }]; + + let hasCanonicalizer = 1; +} + +def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ + AttrSizedOperandSegments, + Pure, + TypesMatchWith<"expected result type to match dest type", + "dest", "result", "$_self"> + ]> { + let summary = "insert_slice operation"; + let description = [{ + }]; + + let arguments = (ins + AnyRankedTensor:$source, + AnyRankedTensor:$dest, + Variadic:$offsets, + Variadic:$sizes, + Variadic:$strides, + DenseI64ArrayAttr:$static_offsets, + DenseI64ArrayAttr:$static_sizes, + DenseI64ArrayAttr:$static_strides + ); + let results = (outs AnyRankedTensor:$result); + + let assemblyFormat = [{ + $source `into` $dest `` + custom($offsets, $static_offsets) + custom($sizes, $static_sizes) + custom($strides, $static_strides) + attr-dict `:` type($source) `into` type($dest) + }]; + + let builders = [ + // Build a InsertSliceOp with mixed static and dynamic entries and inferred + // result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a InsertSliceOp with dynamic entries and inferred result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an InsertSliceOp with mixed static and dynamic entries packed in + // a Range vector and inferred result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$ranges, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an InsertSliceOp with mixed static and dynamic sizes, offsets set + // to 0, strides set to 1 and inferred result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$sizes, + CArg<"ArrayRef", "{}">:$attrs)>, + ]; + + let hasCanonicalizer = 1; +} + +def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { + let summary = "Deallocate a tensor"; + let description = [{ + }]; + + let arguments = (ins AnyRankedTensor:$tensor); + let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; +} #endif // TENSOR_OPS diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 2e143562bb..99dbf32024 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -635,6 +635,151 @@ void QTensorDialect::getCanonicalizationPatterns( RewritePatternSet& results) const { // results.add(getContext()); } + +//===----------------------------------------------------------------------===// +// FromElementsOp +//===----------------------------------------------------------------------===// + +void FromElementsOp::build(OpBuilder& builder, OperationState& result, + ValueRange elements) { + assert(!elements.empty() && "expected at least one element"); + Type resultType = RankedTensorType::get( + {static_cast(elements.size())}, elements.front().getType()); + build(builder, result, resultType, elements); +} +namespace { + +struct ConvertFromElementsOpToTensorOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(qtensor::FromElementsOp fromElementsOp, + PatternRewriter& rewriter) const final { + + rewriter.replaceOpWithNewOp( + fromElementsOp, fromElementsOp.getElements()); + + return success(); + } +}; + +} // namespace + +void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} + +//===----------------------------------------------------------------------===// +// InsertOp +//===----------------------------------------------------------------------===// + +namespace { + +struct ConvertInsertOpToTensorOp : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(qtensor::InsertOp insertOp, + PatternRewriter& rewriter) const final { + rewriter.replaceOpWithNewOp( + insertOp, insertOp.getScalar(), insertOp.getDest(), + insertOp.getIndices()); + return success(); + } +}; + +} // namespace + +void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} + +//===----------------------------------------------------------------------===// +// InsertSliceOp +//===----------------------------------------------------------------------===// + +// Build a InsertSliceOp with mixed static and dynamic entries. +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + SmallVector staticOffsets; + SmallVector staticSizes; + SmallVector staticStrides; + SmallVector dynamicOffsets; + SmallVector dynamicSizes; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + result.addAttributes(attrs); + build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, + dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), + b.getDenseI64ArrayAttr(staticSizes), + b.getDenseI64ArrayAttr(staticStrides)); +} + +/// Build an InsertSliceOp with mixed static and dynamic entries packed into a +/// Range vector. +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ArrayRef ranges, + ArrayRef attrs) { + auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); + build(b, result, source, dest, offsets, sizes, strides, attrs); +} + +// Build a InsertSliceOp with dynamic entries. +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + SmallVector offsetValues = llvm::to_vector<4>( + llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); + SmallVector sizeValues = llvm::to_vector<4>( + llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = llvm::to_vector<4>( + llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + build(b, result, source, dest, offsetValues, sizeValues, strideValues); +} + +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ArrayRef sizes, + ArrayRef attrs) { + Attribute zeroIdxAttr = b.getIndexAttr(0); + Attribute oneIdxAttr = b.getIndexAttr(1); + SmallVector writeStrides(sizes.size(), oneIdxAttr); + SmallVector writeOffsets(sizes.size(), zeroIdxAttr); + build(b, result, source, dest, writeOffsets, sizes, writeStrides, attrs); +} + +namespace { + +struct ConvertInsertSliceOpToTensorOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(qtensor::InsertSliceOp insertSliceOp, + PatternRewriter& rewriter) const final { + rewriter.replaceOpWithNewOp( + insertSliceOp, + insertSliceOp.getResult().getType(), // explicit result type + insertSliceOp.getSource(), insertSliceOp.getDest(), + insertSliceOp.getOffsets(), insertSliceOp.getSizes(), + insertSliceOp.getStrides(), + insertSliceOp.getStaticOffsets(), // static integer attrs + insertSliceOp.getStaticSizes(), insertSliceOp.getStaticStrides()); + return success(); + } +}; + +} // namespace + +void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} + //===----------------------------------------------------------------------===// // Dialect //===----------------------------------------------------------------------===// From affb356618050df67a18d61ce97ec3ef5e8a5ce3 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sun, 1 Mar 2026 14:19:19 +0100 Subject: [PATCH 003/108] adjust build functions --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 9 +++- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 43 +++++++++++++++---- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index a7e91590d0..a9d8c43a6b 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -132,6 +132,10 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with mixed static and dynamic entries and custom // result type. If the type passed is nullptr, it is inferred. + OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, "ArrayRef":$offsets, "ArrayRef":$sizes, "ArrayRef":$strides, @@ -142,7 +146,10 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", "ValueRange":$sizes, "ValueRange":$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with dynamic entries and inferred result type. - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with mixed static and dynamic entries packed in diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 99dbf32024..32f57758cf 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -285,7 +285,8 @@ RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( /// Build an ExtractSliceOp with mixed static and dynamic entries and custom /// result type. If the type passed is nullptr, it is inferred. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, + RankedTensorType resultType, + RankedTensorType outSourceType, Value source, ArrayRef offsets, ArrayRef sizes, ArrayRef strides, @@ -306,11 +307,20 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, sourceRankedTensorType, staticOffsets, staticSizes, staticStrides)); } result.addAttributes(attrs); - build(b, result, resultType, source, dynamicOffsets, dynamicSizes, - dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), + build(b, result, {resultType, outSourceType}, source, dynamicOffsets, + dynamicSizes, dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), b.getDenseI64ArrayAttr(staticSizes), b.getDenseI64ArrayAttr(staticStrides)); } +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + build(b, result, resultType, cast(source.getType()), source, + offsets, sizes, strides, attrs); +} /// Build an ExtractSliceOp with mixed static and dynamic entries and inferred /// result type. @@ -319,7 +329,8 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, ArrayRef sizes, ArrayRef strides, ArrayRef attrs) { - build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); + build(b, result, RankedTensorType(), cast(source.getType()), + source, offsets, sizes, strides, attrs); } /// Build an ExtractSliceOp with mixed static and dynamic entries packed into @@ -328,13 +339,15 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, ArrayRef ranges, ArrayRef attrs) { auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); - build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); + build(b, result, RankedTensorType(), cast(source.getType()), + source, offsets, sizes, strides, attrs); } /// Build an ExtractSliceOp with dynamic entries and custom result type. If /// the type passed is nullptr, it is inferred. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, + RankedTensorType resultType, + RankedTensorType outSourceType, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef attrs) { SmallVector offsetValues = llvm::to_vector<4>( @@ -343,14 +356,23 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); SmallVector strideValues = llvm::to_vector<4>( llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, resultType, source, offsetValues, sizeValues, strideValues); + build(b, result, resultType, outSourceType, source, offsetValues, sizeValues, + strideValues); +} +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + build(b, result, resultType, resultType, source, offsets, sizes, strides, + attrs); } /// Build an ExtractSliceOp with dynamic entries and inferred result type. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef attrs) { - build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); + build(b, result, RankedTensorType(), cast(source.getType()), + source, offsets, sizes, strides, attrs); } static LogicalResult produceSliceErrorMsg(SliceVerificationResult result, @@ -546,11 +568,12 @@ struct SliceCanonicalizer { void operator()(PatternRewriter& rewriter, ExtractSliceOp op, ExtractSliceOp newOp) { Value replacement = newOp.getResult(); + Value outSource = newOp.getOutSource(); if (replacement.getType() != op.getType()) { replacement = rewriter.create(op.getLoc(), op.getType(), replacement); } - rewriter.replaceOp(op, replacement); + rewriter.replaceOp(op, {replacement, outSource}); } }; @@ -610,10 +633,12 @@ LogicalResult ExtractSliceOp::fold(FoldAdaptor adaptor, if (getSourceType() == getType() && succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { results.push_back(this->getSource()); + results.push_back(getSource()); return success(); } if (Value slice = foldExtractAfterInsertSlice(*this)) { results.push_back(slice); + results.push_back(getSource()); return success(); } From c7dc6a5b9f09a5db5d989417dc66e92895507af7 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sun, 1 Mar 2026 14:19:44 +0100 Subject: [PATCH 004/108] add builders in QCO --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 22 ++++ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 110 +++++++++++++++++- 2 files changed, 131 insertions(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index ec96738fab..b4d5b971d6 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -202,6 +202,28 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { [[nodiscard]] ClassicalRegister allocClassicalBitRegister(int64_t size, std::string name = "c") const; + Value allocateTensor(int64_t size); + + Value fromElements(ValueRange elements); + + std::pair extract(Value tensor, + const std::variant& index); + + std::pair + extractSlice(Value tensor, const std::variant& offset, + const std::variant& sizes, + const std::variant& strides); + + Value insert(Value scalar, Value tensor, + const std::variant& index); + + Value insertSlice(Value sourceTensor, Value destTensor, + const std::variant& offset, + const std::variant& sizes, + const std::variant& strides); + + QCOProgramBuilder& deallocTensor(Value tensor); + //===--------------------------------------------------------------------===// // Measurement and Reset //===--------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index bd45a8d174..a4a7e32686 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -12,6 +12,8 @@ #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 @@ -134,6 +136,108 @@ QCOProgramBuilder::allocClassicalBitRegister(const int64_t size, return {.name = std::move(name), .size = size}; } +Value QCOProgramBuilder::allocateTensor(int64_t size) { + checkFinalized(); + + if (size <= 0) { + llvm::reportFatalUsageError("Size must be positive"); + } + + llvm::SmallVector qubits; + qubits.reserve(static_cast(size)); + for (int64_t i = 0; i < size; ++i) { + auto allocOp = AllocOp::create(*this); + qubits.emplace_back(allocOp); + } + auto fromElementsOp = qtensor::FromElementsOp::create(*this, qubits); + validQubits.insert(fromElementsOp); + return fromElementsOp.getResult(); +} +Value QCOProgramBuilder::fromElements(ValueRange elements) { + auto fromElementsOp = qtensor::FromElementsOp::create(*this, elements); + validQubits.insert(fromElementsOp); + return fromElementsOp.getResult(); +} + +std::pair +QCOProgramBuilder::extract(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); + updateQubitTracking(tensor, outTensor); + + return {qubit, outTensor}; +} + +std::pair +QCOProgramBuilder::extractSlice(Value tensor, + const std::variant& offset, + const std::variant& sizes, + const std::variant& strides) { + checkFinalized(); + + auto offsetValue = utils::variantToValue(*this, getLoc(), offset); + auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); + auto stridesValue = utils::variantToValue(*this, getLoc(), strides); + auto extractSliceOp = qtensor::ExtractSliceOp::create( + *this, tensor, offsetValue, sizesValue, stridesValue); + auto slicedTensor = extractSliceOp.getResult(); + auto outTensor = extractSliceOp.getOutSource(); + + validQubits.insert(slicedTensor); + updateQubitTracking(tensor, outTensor); + + return {slicedTensor, outTensor}; +} + +Value QCOProgramBuilder::insert(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(); + + validQubits.erase(scalar); + updateQubitTracking(tensor, outTensor); + return outTensor; +} + +Value QCOProgramBuilder::insertSlice( + Value source, Value dest, const std::variant& offset, + const std::variant& sizes, + const std::variant& strides) { + checkFinalized(); + + auto offsetValue = utils::variantToValue(*this, getLoc(), offset); + auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); + auto stridesValue = utils::variantToValue(*this, getLoc(), strides); + auto insertSliceOp = qtensor::InsertSliceOp::create( + *this, source, dest, offsetValue, sizesValue, stridesValue); + + auto outTensor = insertSliceOp.getResult(); + + validQubits.erase(source); + updateQubitTracking(dest, outTensor); + + return outTensor; +} + +QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { + checkFinalized(); + + qtensor::DeallocOp::create(*this, tensor); + + validQubits.erase(tensor); + return *this; +} //===----------------------------------------------------------------------===// // Linear Type Tracking Helpers //===----------------------------------------------------------------------===// @@ -787,7 +891,11 @@ OwningOpRef QCOProgramBuilder::finalize() { return opA->isBeforeInBlock(opB); }); for (auto qubit : sortedQubits) { - DeallocOp::create(*this, qubit); + if (llvm::isa(qubit.getType())) { + DeallocOp::create(*this, qubit); + } else { + qtensor::DeallocOp::create(*this, qubit); + } } validQubits.clear(); From efa7aff438d33e75a3d89b60ef5e74a66ecba473 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sun, 1 Mar 2026 14:35:22 +0100 Subject: [PATCH 005/108] add canonicalization for extract --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 1 + mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index a9d8c43a6b..c558fbed24 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -84,6 +84,7 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ let hasFolder = 1; let hasVerifier = 1; + let hasCanonicalizer = 1; } diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 32f57758cf..3f7a9c514e 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -132,6 +132,23 @@ LogicalResult ExtractOp::verify() { } return success(); } +struct ExtractFromTensorCast : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(ExtractOp extract, + PatternRewriter& rewriter) const final { + auto tensorCast = extract.getTensor().getDefiningOp(); + if (!tensorCast) { + return failure(); + } + if (!llvm::isa(tensorCast.getSource().getType())) { + return failure(); + } + rewriter.replaceOpWithNewOp(extract, tensorCast.getSource(), + extract.getIndices()); + return success(); + } +}; /// If we have an ExtractOp consuming an InsertOp with the same /// indices, we can return the InsertOp's scalar directly. @@ -194,6 +211,10 @@ LogicalResult ExtractOp::fold(FoldAdaptor adaptor, return failure(); } +void ExtractOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} //===----------------------------------------------------------------------===// // ExtractSliceOp //===----------------------------------------------------------------------===// From 3c3901ec10198291f071ce87161d197a5d09d798 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 16:12:27 +0100 Subject: [PATCH 006/108] add additional canonicalization --- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 106 +++++++++------------ 1 file changed, 44 insertions(+), 62 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 3f7a9c514e..c8be577349 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -8,45 +8,29 @@ * Licensed under the MIT License */ -//===----------------------------------------------------------------------===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// - #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" // for affine::AffineDialect -#include "mlir/Dialect/Arith/IR/Arith.h" // for arith::ArithDialect -#include "mlir/Dialect/Complex/IR/Complex.h" // for complex::ComplexDialect #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated -#include "mlir/Dialect/Utils/StaticValueUtils.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinAttributeInterfaces.h" -#include "mlir/IR/BuiltinTypeInterfaces.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/TypeUtilities.h" -#include "mlir/Interfaces/DestinationStyleOpInterface.h" -#include "mlir/Interfaces/InferIntRangeInterface.h" -#include "mlir/Interfaces/LoopLikeInterface.h" -#include "mlir/Interfaces/Utils/InferIntRangeCommon.h" -#include "mlir/Interfaces/ViewLikeInterface.h" -#include "mlir/Support/LLVM.h" - -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallBitVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/MathExtras.h" +#include +#include #include +#include +#include +#include // for affine::AffineDialect +#include // for arith::ArithDialect #include +#include // for complex::ComplexDialect #include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include // The following headers are needed for some template instantiations. @@ -178,30 +162,6 @@ LogicalResult ExtractOp::fold(FoldAdaptor adaptor, indices.push_back(llvm::cast(indice).getInt()); } - // Fold extract(from_elements(...)). - if (auto fromElementsOp = - getTensor().getDefiningOp()) { - auto tensorType = llvm::cast(fromElementsOp.getType()); - auto rank = tensorType.getRank(); - assert(static_cast(indices.size()) == tensorType.getRank() && - "rank mismatch"); - int flatIndex = 0; - int stride = 1; - for (int i = rank - 1; i >= 0; --i) { - flatIndex += indices[i] * stride; - stride *= tensorType.getDimSize(i); - } - // Prevent out of bounds accesses. This can happen in invalid code that - // will never execute. - if (static_cast(fromElementsOp.getElements().size()) <= flatIndex || - flatIndex < 0) { - return failure(); - } - results.push_back(fromElementsOp.getElements()[flatIndex]); - results.push_back(getTensor()); - return success(); - } - if (Value result = foldExtractAfterInsert(*this)) { results.push_back(result); results.push_back(getTensor()); @@ -464,7 +424,7 @@ ExtractSliceOp::rankReduceIfNeeded(OpBuilder& b, Location loc, Value value, } LogicalResult ExtractSliceOp::reifyResultShapes( - OpBuilder& builder, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { + OpBuilder& /*builder*/, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { reifiedReturnShapes.resize(1); reifiedReturnShapes[0].reserve(getType().getRank()); SmallVector mixedSizes = getMixedSizes(); @@ -541,10 +501,10 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern { /// numbers of elements to stride in the original values for each dimension. /// The output values can be used to construct a DenseElementsAttr. template -static void sliceElements(IterTy values, ArrayRef counts, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides, - llvm::SmallVectorImpl* outValues) { +void sliceElements(IterTy values, ArrayRef counts, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides, + llvm::SmallVectorImpl* outValues) { assert(offsets.size() == sizes.size()); assert(offsets.size() == strides.size()); if (offsets.empty()) { @@ -734,11 +694,33 @@ struct ConvertInsertOpToTensorOp : public OpRewritePattern { } }; +struct InsertFromExtractOp : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOp insertOp, + PatternRewriter& rewriter) const final { + auto extractOp = insertOp.getScalar().getDefiningOp(); + if (!extractOp) { + return failure(); + } + if (insertOp.getDest() != extractOp.getOutTensor()) { + return failure(); + } + if (insertOp.getIndices() != extractOp.getIndices()) { + return failure(); + } + + rewriter.replaceOp(insertOp, extractOp.getTensor()); + rewriter.eraseOp(extractOp); + return success(); + } +}; + } // namespace void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add(context); } //===----------------------------------------------------------------------===// From fd009b600c5f9263708c03b849c7bac522975df0 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 16:14:09 +0100 Subject: [PATCH 007/108] add tests for qtensor --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 26 +++++++++++- mlir/unittests/programs/qco_programs.cpp | 40 +++++++++++++++++++ mlir/unittests/programs/qco_programs.h | 21 ++++++++++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 75c10cc9a4..5011a8edc1 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" @@ -22,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -46,7 +48,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(); @@ -1020,3 +1023,24 @@ 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{"AllocTensor", MQT_NAMED_BUILDER(allocTensor), + MQT_NAMED_BUILDER(allocTensor)}, + QCOTestCase{"AllocDeallocTensor", MQT_NAMED_BUILDER(allocDeallocTensor), + MQT_NAMED_BUILDER(allocTensor)}, + QCOTestCase{"ExtractTensor", MQT_NAMED_BUILDER(extractTensor), + MQT_NAMED_BUILDER(extractTensor)}, + QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(insertTensor), + MQT_NAMED_BUILDER(insertTensor)}, + QCOTestCase{"ExtractSliceTensor", MQT_NAMED_BUILDER(extractSliceTensor), + MQT_NAMED_BUILDER(extractSliceTensor)}, + QCOTestCase{"InsertSliceTensor", MQT_NAMED_BUILDER(insertSliceTensor), + MQT_NAMED_BUILDER(insertSliceTensor)}, + QCOTestCase{"ExtractInsert", MQT_NAMED_BUILDER(extractInsertTensor), + MQT_NAMED_BUILDER(allocTensor)})); +/// @} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 3312c4c173..57b58c7d53 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2049,4 +2049,44 @@ void nestedFalseIf(QCOProgramBuilder& b) { return llvm::to_vector(innerResult); }); } + +void allocTensor(QCOProgramBuilder& b) { b.allocateTensor(3); } + +void allocDeallocTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + b.deallocTensor(qtensor); +} + +void extractTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + // does not work for now + // b.extract(qtensor, 0); +} + +void insertTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + auto [q0, extractOutTensor] = b.extract(qtensor, 0); + auto q1 = b.h(q0); + b.insert(q1, extractOutTensor, 0); +} + +void extractSliceTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + b.extractSlice(qtensor, 0, 2, 1); +} + +void insertSliceTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); + auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.insert(q1, extractOutTensor, 0); + b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2, 1); +} + +void extractInsertTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + auto [q0, extractOutTensor] = b.extract(qtensor, 0); + b.insert(q0, extractOutTensor, 0); +} } // namespace mlir::qco diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 4f3130e07b..51312d73b6 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -932,4 +932,25 @@ 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); + +/// Allocates a tensor of size `3`. +void allocTensor(QCOProgramBuilder& b); + +/// Allocates and explicitly deallocates a tensor. +void allocDeallocTensor(QCOProgramBuilder& b); + +/// Extracts a qubit from a tensor. +void extractTensor(QCOProgramBuilder& b); + +/// Inserts a qubit in a tensor. +void insertTensor(QCOProgramBuilder& b); + +/// Extracts a slice from a tensor. +void extractSliceTensor(QCOProgramBuilder& b); + +/// Inserts a slice from a tensor. +void insertSliceTensor(QCOProgramBuilder& b); + +/// Extract a qubit from a tensor and insert it immediately. +void extractInsertTensor(QCOProgramBuilder& b); } // namespace mlir::qco From 8cefca0cb92ec541bfcea2c5589d77158cc1fa6d Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 17:20:28 +0100 Subject: [PATCH 008/108] copy tensor insertslice --- .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 23 +- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 46 +- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 399 ++++++++++++++++-- 3 files changed, 418 insertions(+), 50 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h index 122c683c2f..10f630e3c5 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -10,27 +10,14 @@ #pragma once -#include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include "mlir/IR/Dialect.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/Interfaces/ViewLikeInterface.h" +#include +#include +#include +#include +#include #define DIALECT_NAME_QTensor "qtensor" -//===----------------------------------------------------------------------===// -// QTensor Dialect Helpers -//===----------------------------------------------------------------------===// - -namespace mlir { - -/// Return the list of Range (i.e. offset, size, stride). Each Range -/// entry contains either the dynamic value or a ConstantIndexOp constructed -/// with `b` at location `loc`. -SmallVector getOrCreateRanges(OffsetSizeAndStrideOpInterface op, - OpBuilder& b, Location loc); - -} // namespace mlir - //===----------------------------------------------------------------------===// // QTensor Dialect //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index c558fbed24..bf99b86b9b 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -300,14 +300,19 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let hasCanonicalizer = 1; } -def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ +def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ + DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, AttrSizedOperandSegments, + DestinationStyleOpInterface, Pure, + OffsetSizeAndStrideOpInterface, TypesMatchWith<"expected result type to match dest type", "dest", "result", "$_self"> ]> { let summary = "insert_slice operation"; let description = [{ + }]; let arguments = (ins @@ -345,15 +350,42 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ // a Range vector and inferred result type. OpBuilder<(ins "Value":$source, "Value":$dest, "ArrayRef":$ranges, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an InsertSliceOp with mixed static and dynamic sizes, offsets set - // to 0, strides set to 1 and inferred result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ArrayRef":$sizes, - CArg<"ArrayRef", "{}">:$attrs)>, + CArg<"ArrayRef", "{}">:$attrs)> ]; + let extraClassDeclaration = extraBaseClassDeclaration # [{ + /// The result of a insert_slice is always a tensor. + // TODO: Deprecate this method. + RankedTensorType getType() { + return getResultType(); + } + + /// The `dest` type is the same as the result type. + RankedTensorType getDestType() { + return getResultType(); + } + + /// Return the expected rank of each of the`static_offsets`, `static_sizes` + /// and `static_strides` attributes. + std::array getArrayAttrMaxRanks() { + unsigned rank = getResultType().getRank(); + return {rank, rank, rank}; + } + + /// Return the dimensions of the dest that are omitted to insert a source + /// when the result is rank-extended. + llvm::SmallBitVector getDroppedDims(); + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; } + + MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } + }]; + let hasCanonicalizer = 1; + let hasFolder = 1; + let hasVerifier = 1; } def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index c8be577349..1f1242fd0c 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -330,7 +331,8 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, RankedTensorType resultType, RankedTensorType outSourceType, Value source, ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { + ValueRange strides, + ArrayRef /*attrs*/) { SmallVector offsetValues = llvm::to_vector<4>( llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); SmallVector sizeValues = llvm::to_vector<4>( @@ -608,7 +610,7 @@ static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { return {}; } -LogicalResult ExtractSliceOp::fold(FoldAdaptor adaptor, +LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { if (getSourceType() == getType() && @@ -726,6 +728,10 @@ void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, //===----------------------------------------------------------------------===// // InsertSliceOp //===----------------------------------------------------------------------===// +void InsertSliceOp::getAsmResultNames( + function_ref setNameFn) { + setNameFn(getResult(), "inserted_slice"); +} // Build a InsertSliceOp with mixed static and dynamic entries. void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, @@ -761,7 +767,8 @@ void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, // Build a InsertSliceOp with dynamic entries. void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value dest, ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { + ValueRange strides, + ArrayRef /*attrs*/) { SmallVector offsetValues = llvm::to_vector<4>( llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); SmallVector sizeValues = llvm::to_vector<4>( @@ -771,41 +778,383 @@ void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, build(b, result, source, dest, offsetValues, sizeValues, strideValues); } -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ArrayRef sizes, - ArrayRef attrs) { - Attribute zeroIdxAttr = b.getIndexAttr(0); - Attribute oneIdxAttr = b.getIndexAttr(1); - SmallVector writeStrides(sizes.size(), oneIdxAttr); - SmallVector writeOffsets(sizes.size(), zeroIdxAttr); - build(b, result, source, dest, writeOffsets, sizes, writeStrides, attrs); +/// Rank-reducing type verification for both InsertSliceOp and +/// ParallelInsertSliceOp. +static SliceVerificationResult verifyInsertSliceOp( + RankedTensorType srcType, RankedTensorType dstType, + ArrayRef staticOffsets, ArrayRef staticSizes, + ArrayRef staticStrides, RankedTensorType* expectedType = nullptr) { + // insert_slice is the inverse of extract_slice, use the same type + // inference. + RankedTensorType expected = ExtractSliceOp::inferResultType( + dstType, staticOffsets, staticSizes, staticStrides); + if (expectedType != nullptr) { + *expectedType = expected; + } + return isRankReducedType(expected, srcType); +} + +/// Verifier for InsertSliceOp. +LogicalResult InsertSliceOp::verify() { + // Verify result type against inferred type. + RankedTensorType expectedType; + SliceVerificationResult result = + verifyInsertSliceOp(getSourceType(), getType(), getStaticOffsets(), + getStaticSizes(), getStaticStrides(), &expectedType); + if (result != SliceVerificationResult::Success) { + return produceSliceErrorMsg(result, *this, expectedType); + } + + // Verify that offsets, sizes, strides do not run out-of-bounds with respect + // to the destination tensor. + SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( + getDestType().getShape(), getStaticOffsets(), getStaticSizes(), + getStaticStrides(), /*generateErrorMessage=*/true); + if (!boundsResult.isValid) { + return getOperation()->emitError(boundsResult.errorMessage); + } + + return success(); +} + +/// If we have two consecutive InsertSliceOp writing to the same slice, we +/// can mutate the second InsertSliceOp's destination to the first one's. +/// +/// Example: +/// +/// ```mlir +/// %0 = tensor.insert_slice %slice0 into %input[0, 0] [64, 64] [1, 1] +/// %1 = tensor.insert_slice %slice1 into %0[0, 0] [64, 64] [1, 1] +/// ``` +/// +/// folds into: +/// +/// ```mlir +/// %1 = tensor.insert_slice %slice1 into %input[0, 0] [64, 64] [1, 1] +/// ``` +/// +/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { + auto prevInsertOp = insertOp.getDest().getDefiningOp(); + + auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; + if (!prevInsertOp || + prevInsertOp.getSource().getType() != insertOp.getSource().getType() || + !prevInsertOp.isSameAs(insertOp, isSame)) { + return failure(); + } + + insertOp.getDestMutable().assign(prevInsertOp.getDest()); + return success(); +} + +/// Folds round-trip extract/insert slice op pairs. +/// Example: +/// ```mlir +/// %0 = tensor.extract_slice %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] +/// %1 = tensor.insert_slice %0 into %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] +/// ``` +/// can be folded into %val. +static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { + auto extractOp = insertOp.getSource().getDefiningOp(); + auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; + if (!extractOp || extractOp.getOutSource() != insertOp.getDest() || + !extractOp.isSameAs(insertOp, isSame)) { + return nullptr; + } + return extractOp.getSource(); +} + +OpFoldResult InsertSliceOp::fold(FoldAdaptor) { + if (getSourceType().hasStaticShape() && getType().hasStaticShape() && + getSourceType() == getType() && + succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { + return this->getSource(); + } + if (succeeded(foldInsertAfterInsertSlice(*this))) { + return getResult(); + } + if (auto result = foldInsertAfterExtractSlice(*this)) { + return result; + } + if (llvm::any_of(getMixedSizes(), isZeroInteger)) { + return getDest(); + } + return {}; +} + +LogicalResult InsertSliceOp::reifyResultShapes( + OpBuilder& builder, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { + reifiedReturnShapes.resize(1, SmallVector(getType().getRank())); + reifiedReturnShapes[0] = tensor::getMixedSizes(builder, getLoc(), getDest()); + return success(); } namespace { +/// Pattern to rewrite a insert_slice op with constant arguments. +/// +/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +template +class InsertSliceOpConstantArgumentFolder final + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; -struct ConvertInsertSliceOpToTensorOp - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + PatternRewriter& rewriter) const override { + SmallVector mixedOffsets(insertSliceOp.getMixedOffsets()); + SmallVector mixedSizes(insertSliceOp.getMixedSizes()); + SmallVector mixedStrides(insertSliceOp.getMixedStrides()); + + // No constant operands were folded, just return; + if (failed(foldDynamicOffsetSizeList(mixedOffsets)) && + failed(foldDynamicOffsetSizeList(mixedSizes)) && + failed(foldDynamicStrideList(mixedStrides))) { + return failure(); + } - LogicalResult matchAndRewrite(qtensor::InsertSliceOp insertSliceOp, - PatternRewriter& rewriter) const final { - rewriter.replaceOpWithNewOp( - insertSliceOp, - insertSliceOp.getResult().getType(), // explicit result type - insertSliceOp.getSource(), insertSliceOp.getDest(), - insertSliceOp.getOffsets(), insertSliceOp.getSizes(), - insertSliceOp.getStrides(), - insertSliceOp.getStaticOffsets(), // static integer attrs - insertSliceOp.getStaticSizes(), insertSliceOp.getStaticStrides()); + // Pattern does not apply if the produced op would not verify. + SliceBoundsVerificationResult sliceResult = + verifyInBoundsSlice(insertSliceOp.getDest().getType().getShape(), + mixedOffsets, mixedSizes, mixedStrides); + if (!sliceResult.isValid) { + return failure(); + } + + // Create the new op in canonical form. + auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType( + insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(), + mixedOffsets, mixedSizes, mixedStrides); + Value toInsert = insertSliceOp.getSource(); + if (sourceType != insertSliceOp.getSourceType()) { + OpBuilder::InsertionGuard g(rewriter); + // The only difference between InsertSliceOp and ParallelInsertSliceOp + // is that the insertion point is just before the ParallelCombiningOp in + // the parallel case. + if (std::is_same_v) { + rewriter.setInsertionPoint(insertSliceOp->getParentOp()); + } + toInsert = rewriter.create(insertSliceOp.getLoc(), + sourceType, toInsert); + } + rewriter.replaceOpWithNewOp( + insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets, + mixedSizes, mixedStrides); + return success(); + } +}; + +/// Fold tensor_casts with insert_slice operations. If the source or +/// destination tensor is a tensor_cast that removes static type information, +/// the cast is folded into the insert_slice operation. E.g.: +/// +/// ```mlir +/// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor +/// %2 = tensor.insert_slice %1 into ... : tensor into ... +/// ``` +/// +/// folds into: +/// +/// ```mlir +/// %2 = tensor.insert_slice %0 into ... : tensor<8x16xf32> into ... +/// ``` +/// +/// Note: When folding a cast on the destination tensor, the result of the +/// insert_slice operation is casted to ensure that the type of the result did +/// not change. +/// +/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +template +struct InsertSliceOpCastFolder final : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + PatternRewriter& rewriter) const override { + if (llvm::any_of(insertSliceOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) { + return failure(); + } + + auto getSourceOfCastOp = [](Value v) -> std::optional { + auto castOp = v.getDefiningOp(); + if (!castOp || !canFoldIntoConsumerOp(castOp)) { + return std::nullopt; + } + return castOp.getSource(); + }; + std::optional sourceCastSource = + getSourceOfCastOp(insertSliceOp.getSource()); + std::optional destCastSource = + getSourceOfCastOp(insertSliceOp.getDest()); + if (!sourceCastSource && !destCastSource) { + return failure(); + } + + auto src = + (sourceCastSource ? *sourceCastSource : insertSliceOp.getSource()); + auto dst = (destCastSource ? *destCastSource : insertSliceOp.getDest()); + auto srcType = llvm::dyn_cast(src.getType()); + auto dstType = llvm::dyn_cast(dst.getType()); + if (!srcType || !dstType) { + return failure(); + } + + // The tensor.cast source could have additional static information not seen + // in the insert slice op static sizes, so we ignore dynamic dims when + // computing the rank reduction mask. + SmallVector staticSizes(insertSliceOp.getStaticSizes()); + auto rankReductionMask = computeRankReductionMask( + staticSizes, srcType.getShape(), /*matchDynamic=*/true); + if (!rankReductionMask.has_value()) { + return failure(); + } + // Replace dimensions in the insert slice op with corresponding static dims + // from the cast source type. If the insert slice sizes have static dims + // that are not static in the tensor.cast source (i.e., when the cast op + // casts a dynamic dim to static), the dim should not be replaced, and the + // pattern will fail later in `verifyInsertSliceOp`. + SmallVector mixedSizes(insertSliceOp.getMixedSizes()); + int64_t rankReducedIdx = 0; + for (auto [idx, size] : enumerate(staticSizes)) { + if (!rankReductionMask.value().contains(idx) && + !srcType.isDynamicDim(rankReducedIdx)) { + mixedSizes[idx] = getAsIndexOpFoldResult( + rewriter.getContext(), srcType.getDimSize(rankReducedIdx)); + size = srcType.getDimSize(rankReducedIdx++); + } + } + + // Pattern does not apply if the produced op would not verify. + if (verifyInsertSliceOp(srcType, dstType, insertSliceOp.getStaticOffsets(), + staticSizes, insertSliceOp.getStaticStrides()) != + SliceVerificationResult::Success) { + return failure(); + } + SliceBoundsVerificationResult sliceResult = + verifyInBoundsSlice(dstType.getShape(), insertSliceOp.getMixedOffsets(), + mixedSizes, insertSliceOp.getMixedStrides()); + if (!sliceResult.isValid) { + return failure(); + } + + Operation* replacement = rewriter.create( + insertSliceOp.getLoc(), src, dst, insertSliceOp.getMixedOffsets(), + mixedSizes, insertSliceOp.getMixedStrides()); + + // In the parallel case there is no result and so nothing to cast. + bool isParallelInsert = + std::is_same::value; + if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { + replacement = rewriter.create(insertSliceOp.getLoc(), + insertSliceOp.getDestType(), + replacement->getResult(0)); + } + rewriter.replaceOp(insertSliceOp, replacement->getResults()); return success(); } }; +/// If additional static type information can be deduced from a insert_slice's +/// size operands, insert an explicit cast of the op's source operand. This +/// enables other canonicalization patterns that are matching for tensor_cast +/// ops such as `ForOpTensorCastFolder` in SCF. +/// +/// Example: +/// +/// ```mlir +/// %r = tensor.insert_slice %0 into %1[...] [64, 64] [1, 1] +/// : tensor into ... +/// ``` +/// +/// folds into: +/// +/// ```mlir +/// %tmp = tensor.cast %0 : tensor to tensor<64x64xf32> +/// %r = tensor.insert_slice %tmp into %1[...] [64, 64] [1, 1] +/// : tensor<64x64xf32> into ... +/// ``` +/// +/// This patterns works with both InsertSliceOp and ParallelInsertSliceOp. +template +struct InsertSliceOpSourceCastInserter final + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + PatternRewriter& rewriter) const override { + RankedTensorType srcType = insertSliceOp.getSourceType(); + if (srcType.getRank() != insertSliceOp.getDestType().getRank()) { + return failure(); + } + SmallVector newSrcShape(srcType.getShape()); + for (int64_t i = 0; i < srcType.getRank(); ++i) { + if (std::optional constInt = + getConstantIntValue(insertSliceOp.getMixedSizes()[i])) { + // Bail on invalid IR. + if (*constInt < 0) { + return failure(); + } + newSrcShape[i] = *constInt; + } + } + if (!hasValidSizesOffsets(newSrcShape)) { + return failure(); + } + + RankedTensorType newSrcType = RankedTensorType::get( + newSrcShape, srcType.getElementType(), srcType.getEncoding()); + if (srcType == newSrcType || + !mlir::tensor::preservesStaticInformation(srcType, newSrcType) || + !tensor::CastOp::areCastCompatible(srcType, newSrcType)) { + return failure(); + } + + // newSrcType is: + // 1) Different from srcType. + // 2) "More static" than srcType. + // 3) Cast-compatible with srcType. + // Insert the cast. + OpBuilder::InsertionGuard g(rewriter); + // The only difference between InsertSliceOp and ParallelInsertSliceOp is + // that the insertion point is just before the ParallelCombiningOp in the + // parallel case. + if (std::is_same_v) { + rewriter.setInsertionPoint(insertSliceOp->getParentOp()); + } + Value cast = rewriter.create( + insertSliceOp.getLoc(), newSrcType, insertSliceOp.getSource()); + rewriter.replaceOpWithNewOp( + insertSliceOp, cast, insertSliceOp.getDest(), + insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(), + insertSliceOp.getMixedStrides()); + return success(); + } +}; } // namespace +llvm::SmallBitVector InsertSliceOp::getDroppedDims() { + return ::getDroppedDims(getSourceType().getShape(), getMixedSizes()); +} + void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add, + InsertSliceOpCastFolder, + InsertSliceOpSourceCastInserter>(context); +} + +Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder& b, + Location loc, + Value tensor, + Value dest) { + auto rankedTensorType = llvm::cast(dest.getType()); + unsigned rank = rankedTensorType.getRank(); + SmallVector offsets(rank, b.getIndexAttr(0)); + SmallVector sizes = getMixedSizes(b, loc, dest); + SmallVector strides(rank, b.getIndexAttr(1)); + return b.createOrFold(loc, tensor, dest, offsets, + sizes, strides); } //===----------------------------------------------------------------------===// From 78bf30bb678904ffe2951e35e9ae0977061b8955 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 17:21:59 +0100 Subject: [PATCH 009/108] add additional tests --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 9 ++++++++- mlir/unittests/programs/qco_programs.cpp | 15 +++++++++++++++ mlir/unittests/programs/qco_programs.h | 9 ++++++++- 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 5011a8edc1..b561f91500 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1042,5 +1042,12 @@ INSTANTIATE_TEST_SUITE_P( QCOTestCase{"InsertSliceTensor", MQT_NAMED_BUILDER(insertSliceTensor), MQT_NAMED_BUILDER(insertSliceTensor)}, QCOTestCase{"ExtractInsert", MQT_NAMED_BUILDER(extractInsertTensor), - MQT_NAMED_BUILDER(allocTensor)})); + MQT_NAMED_BUILDER(allocTensor)}, + QCOTestCase{"ExtractSliceInsertSlice", + MQT_NAMED_BUILDER(extractSliceInsertSliceTensor), + MQT_NAMED_BUILDER(allocTensor)}, + QCOTestCase{ + "ExtractSliceExtractInsertInsertSliceTensor", + MQT_NAMED_BUILDER(extractSliceExtractInsertInsertSliceTensor), + MQT_NAMED_BUILDER(allocTensor)})); /// @} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 57b58c7d53..4e07bf7aa6 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2089,4 +2089,19 @@ void extractInsertTensor(QCOProgramBuilder& b) { auto [q0, extractOutTensor] = b.extract(qtensor, 0); b.insert(q0, extractOutTensor, 0); } + +void extractSliceInsertSliceTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); + b.insertSlice(slicedTensor, extractSliceOutTensor, 0, 2, 1); +} + +void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { + auto qtensor = b.allocateTensor(3); + auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); + auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); + auto insertOutTensor = b.insert(q0, extractOutTensor, 0); + b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2, 1); +} + } // namespace mlir::qco diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 51312d73b6..2cba4e458f 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -951,6 +951,13 @@ void extractSliceTensor(QCOProgramBuilder& b); /// Inserts a slice from a tensor. void insertSliceTensor(QCOProgramBuilder& b); -/// Extract a qubit from a tensor and insert it immediately. +/// Extracts a qubit from a tensor and insert it immediately. void extractInsertTensor(QCOProgramBuilder& b); + +/// Extracts a slice of qubits from a tensor and insert it immediately. +void extractSliceInsertSliceTensor(QCOProgramBuilder& b); + +/// Extracts a slice of qubits, a qubit from the slice, insert the qubit back to +/// the slice and the slice back to the tensor immediately. +void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b); } // namespace mlir::qco From 596636c221ee2bdc8454558936fb696ab3c2d6ed Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 17:59:22 +0100 Subject: [PATCH 010/108] add docstrings in QCOProgramBuilder --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index b4d5b971d6..3d1f3ef0ef 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -202,26 +202,136 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { [[nodiscard]] ClassicalRegister allocClassicalBitRegister(int64_t size, std::string name = "c") const; + //===--------------------------------------------------------------------===// + // QTensor operations + //===--------------------------------------------------------------------===// + + /** + * @brief Allocate a qubit tensor + * @param size Number of qubits (must be positive) + * @return The allocated tensor + * + * @par Example: + * ```c++ + * auto tensor = builder.allocateTensor(3); + * ``` + * ```mlir + * %q0 = qco.alloc : !qco.qubit + * %q1 = qco.alloc : !qco.qubit + * %q2 = qco.alloc : !qco.qubit + * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + * ``` + */ Value allocateTensor(int64_t size); + /** + * @brief Allocate a qubit tensor from a list of qubit values + * @param elements Inserted Qubits + * @return The allocated tensor + * + * @par Example: + * ```c++ + * auto tensor = builder.fromElements({q0, q1, q2}); + * ``` + * ```mlir + * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + * ``` + */ Value fromElements(ValueRange elements); + /** + * @brief Extract a qubit from a tensor + * @param tensor Source tensor + * @param index The index from where the qubit is extracted + * @return Pair of (extractedQubit, outTensor) + * + * @par Example: + * ```c++ + * auto [q0, outTensor] = builder.extract(tensor, 0); + * ``` + * ```mlir + * %q0, %outTensor = qtensor.extract %tensor[%c0]: tensor<3x!qco.qubit> + * ``` + */ std::pair extract(Value tensor, const std::variant& index); + /** + * @brief Extract a qubit slice from a tensor + * @param tensor Source tensor + * @param offset The offset from where the slice is extracted + * @param size The size of the extracted slice + * @param strides The strides from where the values are extracted + * @return Pair of (extractedSlice, outTensor) + * + * @par Example: + * ```c++ + * auto [extractedSlice, outTensor] = builder.extract_slice(tensor, 0, 2, 1); + * ``` + * ```mlir + * %extractedSlice, %outTensor = qtensor.extract_slice %tensor[%c0][%c2][%c1] + * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + * ``` + */ std::pair extractSlice(Value tensor, const std::variant& offset, const std::variant& sizes, const std::variant& strides); + /** + * @brief insert a qubit into a tensor + * @param scalar The scalar qubit that is inserted + * @param tensor The tensor where the qubit is inserted + * @param index The index into where the qubit is inserted + * @return The output tensor + * + * @par Example: + * ```c++ + * auto outTensor = builder.insert(q0, tensor, 0); + * ``` + * ```mlir + * %outTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> + * ``` + */ Value insert(Value scalar, Value tensor, const std::variant& index); + /** + * @brief insert a qubit slice into a tensor + * @param scalar The slice that is inserted + * @param tensor The tensor where the slice is inserted + * @param offset The offset into where the slice is inserted + * @param size The size of the inserted slice + * @param strides The strides into where the values are inserted + * @return The output tensor + * + * @par Example: + * ```c++ + * auto outTensor = builder.insert_slice(slicedTensor, tensor, 0, 2, 1); + * ``` + * ```mlir + * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] + * : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + * ``` + */ Value insertSlice(Value sourceTensor, Value destTensor, const std::variant& offset, const std::variant& sizes, const std::variant& strides); + /** + * @brief Explicitly deallocate a tensor + * @param tensor Tensor to deallocate (must be valid/unconsumed) + * @return Reference to this builder for method chaining + * + * @par Example: + * ```c++ + * builder.deallocTensor(tensor); + * ``` + * ```mlir + * qtensor.dealloc %tensor : tensor<3x!qco.qubit> + * ``` + */ QCOProgramBuilder& deallocTensor(Value tensor); //===--------------------------------------------------------------------===// From bb093a2b3ef419835d1273679dcf755e879a95cc Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 17:59:35 +0100 Subject: [PATCH 011/108] adjust builder functions in QCOProgramBuilder --- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 119 ++++++++++++++---- 1 file changed, 94 insertions(+), 25 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index a4a7e32686..fbcb5c31ba 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -136,6 +136,36 @@ QCOProgramBuilder::allocClassicalBitRegister(const int64_t size, return {.name = std::move(name), .size = size}; } +//===----------------------------------------------------------------------===// +// Linear Type Tracking Helpers +//===----------------------------------------------------------------------===// + +void QCOProgramBuilder::validateQubitValue(Value qubit) const { + if (!validQubits.contains(qubit)) { + llvm::errs() << "Attempting to use an invalid qubit SSA value. " + << "The value may have been consumed by a previous operation " + << "or was never created through this builder.\n"; + llvm::reportFatalUsageError( + "Invalid qubit value used (either consumed or not tracked)"); + } +} + +void QCOProgramBuilder::updateQubitTracking(Value inputQubit, + Value outputQubit) { + // Validate the input qubit + validateQubitValue(inputQubit); + + // Remove the input (consumed) value from tracking + validQubits.erase(inputQubit); + + // Add the output (new) value to tracking + validQubits.insert(outputQubit); +} + +//===----------------------------------------------------------------------===// +// QTensor Operations +//===----------------------------------------------------------------------===// + Value QCOProgramBuilder::allocateTensor(int64_t size) { checkFinalized(); @@ -149,11 +179,21 @@ Value QCOProgramBuilder::allocateTensor(int64_t size) { auto allocOp = AllocOp::create(*this); qubits.emplace_back(allocOp); } + auto fromElementsOp = qtensor::FromElementsOp::create(*this, qubits); validQubits.insert(fromElementsOp); return fromElementsOp.getResult(); } Value QCOProgramBuilder::fromElements(ValueRange elements) { + checkFinalized(); + + for (auto element : elements) { + if (!llvm::isa(element.getType())) { + llvm::reportFatalUsageError("Elements must be QubitType!"); + } + validQubits.erase(element); + } + auto fromElementsOp = qtensor::FromElementsOp::create(*this, elements); validQubits.insert(fromElementsOp); return fromElementsOp.getResult(); @@ -164,6 +204,15 @@ QCOProgramBuilder::extract(Value tensor, const std::variant& index) { checkFinalized(); + auto rankedTensorType = llvm::dyn_cast(tensor.getType()); + + if (!rankedTensorType) { + llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + } + if (!llvm::isa(rankedTensorType.getElementType())) { + llvm::reportFatalUsageError("Elements must be of QubitType!"); + } + auto indexValue = utils::variantToValue(*this, getLoc(), index); auto extractOp = qtensor::ExtractOp::create(*this, tensor, indexValue); auto qubit = extractOp.getResult(); @@ -182,6 +231,15 @@ QCOProgramBuilder::extractSlice(Value tensor, const std::variant& strides) { checkFinalized(); + auto tensorType = llvm::dyn_cast(tensor.getType()); + + if (!tensorType) { + llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + } + if (!llvm::isa(tensorType.getElementType())) { + llvm::reportFatalUsageError("Elements must be of QubitType!"); + } + auto offsetValue = utils::variantToValue(*this, getLoc(), offset); auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); auto stridesValue = utils::variantToValue(*this, getLoc(), strides); @@ -200,6 +258,15 @@ Value QCOProgramBuilder::insert(Value scalar, Value tensor, const std::variant& index) { checkFinalized(); + auto tensorType = llvm::dyn_cast(tensor.getType()); + + if (!tensorType) { + llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + } + if (!llvm::isa(tensorType.getElementType())) { + llvm::reportFatalUsageError("Elements must be of QubitType!"); + } + auto indexValue = utils::variantToValue(*this, getLoc(), index); auto insertOp = qtensor::InsertOp::create(*this, scalar, tensor, indexValue); @@ -216,6 +283,24 @@ Value QCOProgramBuilder::insertSlice( const std::variant& strides) { checkFinalized(); + auto sourceTensorType = llvm::dyn_cast(source.getType()); + + if (!sourceTensorType) { + llvm::reportFatalUsageError("Source must be of RankedTensorType!"); + } + if (!llvm::isa(sourceTensorType.getElementType())) { + llvm::reportFatalUsageError("Source elements must be of QubitType!"); + } + + auto destTensorType = llvm::dyn_cast(source.getType()); + + if (!destTensorType) { + llvm::reportFatalUsageError("Dest must be of RankedTensorType!"); + } + if (!llvm::isa(destTensorType.getElementType())) { + llvm::reportFatalUsageError("Dest elements must be of QubitType!"); + } + auto offsetValue = utils::variantToValue(*this, getLoc(), offset); auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); auto stridesValue = utils::variantToValue(*this, getLoc(), strides); @@ -233,35 +318,19 @@ Value QCOProgramBuilder::insertSlice( QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { checkFinalized(); - qtensor::DeallocOp::create(*this, tensor); - - validQubits.erase(tensor); - return *this; -} -//===----------------------------------------------------------------------===// -// Linear Type Tracking Helpers -//===----------------------------------------------------------------------===// + auto tensorType = llvm::dyn_cast(tensor.getType()); -void QCOProgramBuilder::validateQubitValue(Value qubit) const { - if (!validQubits.contains(qubit)) { - llvm::errs() << "Attempting to use an invalid qubit SSA value. " - << "The value may have been consumed by a previous operation " - << "or was never created through this builder.\n"; - llvm::reportFatalUsageError( - "Invalid qubit value used (either consumed or not tracked)"); + if (!tensorType) { + llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + } + if (!llvm::isa(tensorType.getElementType())) { + llvm::reportFatalUsageError("Elements must be of QubitType!"); } -} - -void QCOProgramBuilder::updateQubitTracking(Value inputQubit, - Value outputQubit) { - // Validate the input qubit - validateQubitValue(inputQubit); - // Remove the input (consumed) value from tracking - validQubits.erase(inputQubit); + qtensor::DeallocOp::create(*this, tensor); - // Add the output (new) value to tracking - validQubits.insert(outputQubit); + validQubits.erase(tensor); + return *this; } //===----------------------------------------------------------------------===// From b705471dbc4bbb26e73d933a77eed5c64bd42560 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 18:14:36 +0100 Subject: [PATCH 012/108] add general canonicalization --- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 119 ++++++++++++++++----- 1 file changed, 95 insertions(+), 24 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 1f1242fd0c..1e336c35f9 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -21,6 +21,7 @@ #include // for arith::ArithDialect #include #include // for complex::ComplexDialect +#include #include #include #include @@ -176,6 +177,7 @@ void ExtractOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { results.add(context); } + //===----------------------------------------------------------------------===// // ExtractSliceOp //===----------------------------------------------------------------------===// @@ -189,8 +191,8 @@ void ExtractSliceOp::getAsmResultNames( /// rank-reduced, from the source type and the static representation of /// offsets, sizes and strides. Special sentinels encode the dynamic case. RankedTensorType ExtractSliceOp::inferResultType( - RankedTensorType sourceTensorType, ArrayRef staticOffsets, - ArrayRef staticSizes, ArrayRef staticStrides) { + RankedTensorType sourceTensorType, ArrayRef /*staticOffsets*/, + ArrayRef staticSizes, ArrayRef /*staticStrides*/) { // An extract_slice op may specify only a leading subset of offset/sizes/ // strides in which case we complete with offset=0, sizes from memref type // and strides=1. @@ -202,8 +204,8 @@ RankedTensorType ExtractSliceOp::inferResultType( } RankedTensorType ExtractSliceOp::inferResultType( - RankedTensorType sourceTensorType, ArrayRef offsets, - ArrayRef sizes, ArrayRef strides) { + RankedTensorType sourceTensorType, ArrayRef /*offsets*/, + ArrayRef sizes, ArrayRef /*strides*/) { SmallVector staticSizes; std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes); assert(static_cast(staticSizes.size()) == @@ -441,21 +443,31 @@ LogicalResult ExtractSliceOp::reifyResultShapes( } namespace { -/// Pattern to rewrite an extract_slice op with tensor::Cast arguments. -/// This essentially pushes memref_cast past its consuming slice when -/// `canFoldIntoConsumerOp` is true. -/// -/// Example: -/// ``` -/// %0 = tensor.cast %V : tensor<16x16xf32> to tensor -/// %1 = tensor.extract_slice %0[0, 0][3, 4][1, 1] : tensor to -/// tensor<3x4xf32> -/// ``` -/// is rewritten into: -/// ``` -/// %0 = tensor.extract_slice %V[0, 0][3, 4][1, 1] : tensor<16x16xf32> to -/// tensor<3x4xf32> %1 = tensor.cast %0: tensor<3x4xf32> to tensor<3x4xf32> -/// ``` +/** + * @brief Rewrite pattern that pushes tensor.cast past tensor.extract_slice. + * + * @details + * This pattern rewrites a `qtensor.extract_slice` operation whose source + * operand is produced by a `tensor.cast`. When `canFoldIntoConsumerOp` + * evaluates to true, the cast operation is moved after the slice operation. + * + * Conceptually, the slice is applied to the original tensor before the + * cast, avoiding unnecessary intermediate casts. + * + * Example: + * %0 = tensor.cast %V : tensor<3x!qco.qubit> to tensor + * %1, %2 = tensor.extract_slice %0[0][2][1] + * : tensor to tensor<2x!qco.qubit> + * + * is rewritten into: + * + * %0, %1 = tensor.extract_slice %V[0][2][1] + * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + * %2 = tensor.cast %0 : tensor<2x!qco.qubit> to tensor<2x!qco.qubit> + * + * This effectively folds the cast into the consumer operation and enables + * further canonicalization opportunities. + */ class ExtractSliceOpCastFolder final : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; @@ -639,11 +651,6 @@ Value mlir::tensor::createCanonicalRankReducingExtractSliceOp( offsets, sizes, strides); } -void QTensorDialect::getCanonicalizationPatterns( - RewritePatternSet& results) const { - // results.add(getContext()); -} - //===----------------------------------------------------------------------===// // FromElementsOp //===----------------------------------------------------------------------===// @@ -655,6 +662,7 @@ void FromElementsOp::build(OpBuilder& builder, OperationState& result, {static_cast(elements.size())}, elements.front().getType()); build(builder, result, resultType, elements); } + namespace { struct ConvertFromElementsOpToTensorOp @@ -1157,6 +1165,69 @@ Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder& b, sizes, strides); } +//===----------------------------------------------------------------------===// +// Common Canonicalizers and Folders. +//===----------------------------------------------------------------------===// + +namespace { + +bool foldTensorCastPrecondition(DestinationStyleOpInterface op) { + // 1. InsertSliceOp has its own logic about folding tensor.cast ops. + // 2. Exclude DPS ops that are also LoopLike from this interface as they + // might need special handling of attached regions. + if (isa(op.getOperation()) || + isa(op.getOperation())) { + return false; + } + + return mlir::tensor::hasFoldableTensorCastOperand(op); +} +} // namespace + +struct FoldTensorCastProducerOp + : public OpInterfaceRewritePattern { + using OpInterfaceRewritePattern< + DestinationStyleOpInterface>::OpInterfaceRewritePattern; + + LogicalResult matchAndRewrite(DestinationStyleOpInterface op, + PatternRewriter& rewriter) const override { + + // Reject PackOp/UnpackOp (i.e. RelayoutOps) - there are dedicated patterns + // for that instead. + if (!foldTensorCastPrecondition(op) || + isa(*op)) { + return failure(); + } + + SmallVector newResultTypes(op->getResultTypes()); + SmallVector newOperands = + mlir::tensor::getUpdatedOperandsAfterCastOpFolding(op, newResultTypes); + + // Clone op + auto newOp = clone(rewriter, op, newResultTypes, newOperands); + + SmallVector replacements; + replacements.reserve(newOp->getNumResults()); + for (auto [oldResult, newResult] : + llvm::zip(op->getResults(), newOp->getResults())) { + if (newResult.getType() != oldResult.getType()) { + replacements.push_back(rewriter.create( + op->getLoc(), oldResult.getType(), newResult)); + } else { + replacements.push_back(newResult); + } + } + rewriter.replaceOp(op, replacements); + + return success(); + } +}; + +void QTensorDialect::getCanonicalizationPatterns( + RewritePatternSet& results) const { + results.add(getContext()); +} + //===----------------------------------------------------------------------===// // Dialect //===----------------------------------------------------------------------===// From 4b310cdb935ae496dbf3d2ef4cc69afb441d4968 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 18:18:34 +0100 Subject: [PATCH 013/108] fix wrong operation in folExtractAfterInsertSlice --- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 1e336c35f9..63b3ddde88 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -611,7 +611,7 @@ foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, // TODO: This only checks the immediate producer; extend to go up the // insert/extract chain if the slices are disjoint. static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { - auto insertOp = extractOp.getSource().getDefiningOp(); + auto insertOp = extractOp.getSource().getDefiningOp(); auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; if (insertOp && insertOp.getSource().getType() == extractOp.getType() && From b8cf6442b9cade6084d1a211f470832244047946 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 2 Mar 2026 19:24:59 +0100 Subject: [PATCH 014/108] adjust structure of qtensor dialect --- .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 89 ++ mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 3 + .../QTensor/IR/Operations/DeallocTensorOp.cpp | 9 + .../QTensor/IR/Operations/ExtractOp.cpp | 114 ++ .../QTensor/IR/Operations/ExtractSliceOp.cpp | 460 +++++++ .../QTensor/IR/Operations/FromTensorOp.cpp | 68 + .../QTensor/IR/Operations/InsertOp.cpp | 80 ++ .../QTensor/IR/Operations/InsertSliceOp.cpp | 466 +++++++ mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 1110 ----------------- 9 files changed, 1289 insertions(+), 1110 deletions(-) create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h index 10f630e3c5..28b735bcab 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -18,6 +18,95 @@ #define DIALECT_NAME_QTensor "qtensor" +namespace mlir::qtensor { +/// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or +/// rank-extending tensor.insert_slice op. +inline llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, + ArrayRef mixedSizes) { + llvm::SmallBitVector droppedDims(mixedSizes.size()); + int64_t shapePos = reducedShape.size() - 1; + + for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { + size_t idx = mixedSizes.size() - size.index() - 1; + // Rank-reduced dims must have a static unit dimension. + bool isStaticUnitSize = + isa(size.value()) && + llvm::cast(cast(size.value())).getInt() == 1; + + if (shapePos < 0) { + // There are no more dims in the reduced shape. All remaining sizes must + // be rank-reduced dims. + assert(isStaticUnitSize && "expected unit dim"); + droppedDims.set(idx); + continue; + } + + // Dim is preserved if the size is not a static 1. + if (!isStaticUnitSize) { + --shapePos; + continue; + } + + // Dim is preserved if the reduced shape dim is also 1. + if (reducedShape[shapePos] == 1) { + --shapePos; + continue; + } + + // Otherwise: Dim is dropped. + droppedDims.set(idx); + } + + assert(shapePos < 0 && "dimension mismatch"); + return droppedDims; +} + +inline LogicalResult produceSliceErrorMsg(SliceVerificationResult result, + Operation* op, + RankedTensorType expectedType) { + switch (result) { + case SliceVerificationResult::Success: + return success(); + case SliceVerificationResult::RankTooLarge: + return op->emitError("expected rank to be smaller or equal to ") + << "the other rank. "; + case SliceVerificationResult::SizeMismatch: + return op->emitError("expected type to be ") + << expectedType << " or a rank-reduced version. (size mismatch) "; + case SliceVerificationResult::ElemTypeMismatch: + return op->emitError("expected element type to be ") + << expectedType.getElementType(); + default: + llvm_unreachable("unexpected extract_slice op verification result"); + } +} + +inline LogicalResult +foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, + ShapedType shapedType) { + OpBuilder b(op.getContext()); + for (OpFoldResult opFold : op.getMixedOffsets()) { + if (getConstantIntValue(opFold) != static_cast(0)) { + return failure(); + } + } + // Rank-reducing noops only need to inspect the leading dimensions: + // llvm::zip is appropriate. + auto shape = shapedType.getShape(); + for (auto it : llvm::zip(op.getMixedSizes(), shape)) { + if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) { + return failure(); + } + } + for (OpFoldResult opFold : op.getMixedStrides()) { + if (getConstantIntValue(opFold) != static_cast(1)) { + return failure(); + } + } + return success(); +} +} // namespace mlir::qtensor + //===----------------------------------------------------------------------===// // QTensor Dialect //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt index 97a3088c2d..7ed18e3bb7 100644 --- a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -6,9 +6,12 @@ # # 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 ${PROJECT_SOURCE_DIR}/mlir/include/mlir/Dialect/QTensor DEPENDS diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp new file mode 100644 index 0000000000..9650bbe91c --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp @@ -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 + */ 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..b80e6874ec --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -0,0 +1,114 @@ +/* + * 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 // for affine::AffineDialect +#include // for arith::ArithDialect +#include +#include // for complex::ComplexDialect +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +//===----------------------------------------------------------------------===// +// ExtractOp +//===----------------------------------------------------------------------===// + +void ExtractOp::getAsmResultNames( + function_ref setNameFn) { + setNameFn(getResult(), "q_extracted"); +} + +LogicalResult ExtractOp::verify() { + // Verify the # indices match if we have a ranked type. + auto tensorType = llvm::cast(getTensor().getType()); + if (tensorType.getRank() != static_cast(getIndices().size())) { + return emitOpError("incorrect number of indices for extract_element"); + } + return success(); +} +struct ExtractFromTensorCast : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(ExtractOp extract, + PatternRewriter& rewriter) const final { + auto tensorCast = extract.getTensor().getDefiningOp(); + if (!tensorCast) { + return failure(); + } + if (!llvm::isa(tensorCast.getSource().getType())) { + return failure(); + } + rewriter.replaceOpWithNewOp(extract, tensorCast.getSource(), + extract.getIndices()); + return success(); + } +}; + +/// If we have an ExtractOp consuming an InsertOp with the same +/// indices, we can return the InsertOp's scalar directly. +// TODO: This only checks the immediate producer; extend to go up the +// insert/extract chain if the slices are disjoint. +static Value foldExtractAfterInsert(ExtractOp extractOp) { + auto insertOp = extractOp.getTensor().getDefiningOp(); + + auto isSame = [](Value a, Value b) { + return getAsOpFoldResult(a) == getAsOpFoldResult(b); + }; + if (insertOp && insertOp.getScalar().getType() == extractOp.getType(0) && + llvm::equal(insertOp.getIndices(), extractOp.getIndices(), isSame)) { + return insertOp.getScalar(); + } + return {}; +} + +LogicalResult ExtractOp::fold(FoldAdaptor adaptor, + SmallVectorImpl& results) { + // Collect the constant indices into the tensor. + SmallVector indices; + for (Attribute indice : adaptor.getIndices()) { + if (!indice || !llvm::isa(indice)) { + return failure(); + } + indices.push_back(llvm::cast(indice).getInt()); + } + + if (Value result = foldExtractAfterInsert(*this)) { + results.push_back(result); + results.push_back(getTensor()); + return success(); + } + + return failure(); +} + +void ExtractOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} 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..340dae6dd5 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -0,0 +1,460 @@ +/* + * 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 // for affine::AffineDialect +#include // for arith::ArithDialect +#include +#include // for complex::ComplexDialect +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +void ExtractSliceOp::getAsmResultNames( + function_ref setNameFn) { + setNameFn(getResult(), "q_extracted_slice"); +} + +/// An extract_slice result type can be inferred, when it is not +/// rank-reduced, from the source type and the static representation of +/// offsets, sizes and strides. Special sentinels encode the dynamic case. +RankedTensorType ExtractSliceOp::inferResultType( + RankedTensorType sourceTensorType, ArrayRef /*staticOffsets*/, + ArrayRef staticSizes, ArrayRef /*staticStrides*/) { + // An extract_slice op may specify only a leading subset of offset/sizes/ + // strides in which case we complete with offset=0, sizes from memref type + // and strides=1. + assert(static_cast(staticSizes.size()) == + sourceTensorType.getRank() && + "unexpected staticSizes not equal to rank of source"); + return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), + sourceTensorType.getEncoding()); +} + +RankedTensorType ExtractSliceOp::inferResultType( + RankedTensorType sourceTensorType, ArrayRef /*offsets*/, + ArrayRef sizes, ArrayRef /*strides*/) { + SmallVector staticSizes; + std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes); + assert(static_cast(staticSizes.size()) == + sourceTensorType.getRank() && + "unexpected staticSizes not equal to rank of source"); + return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), + sourceTensorType.getEncoding()); +} + +/// If the rank is reduced (i.e. the desiredResultRank is smaller than the +/// number of sizes), drop as many size 1 as needed to produce an inferred +/// type with the desired rank. +/// +/// Note that there may be multiple ways to compute this rank-reduced type: +/// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. +/// +/// To disambiguate, this function always drops the first 1 sizes occurrences. +RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( + unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides) { + // Type inferred in the absence of rank-reducing behavior. + auto inferredType = llvm::cast( + inferResultType(sourceRankedTensorType, offsets, sizes, strides)); + int rankDiff = inferredType.getRank() - desiredResultRank; + if (rankDiff > 0) { + auto shape = inferredType.getShape(); + llvm::SmallBitVector dimsToProject = + getPositionsOfShapeOne(rankDiff, shape); + SmallVector projectedShape; + // Best effort rank-reducing: drop 1s in order. + for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) { + if (!dimsToProject.test(pos)) { + projectedShape.push_back(shape[pos]); + } + } + inferredType = + RankedTensorType::get(projectedShape, inferredType.getElementType()); + } + return inferredType; +} + +RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( + unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides) { + SmallVector staticOffsets; + SmallVector staticSizes; + SmallVector staticStrides; + SmallVector dynamicOffsets; + SmallVector dynamicSizes; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + return ExtractSliceOp::inferCanonicalRankReducedResultType( + desiredResultRank, sourceRankedTensorType, staticOffsets, staticSizes, + staticStrides); +} + +/// Build an ExtractSliceOp with mixed static and dynamic entries and custom +/// result type. If the type passed is nullptr, it is inferred. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, + RankedTensorType outSourceType, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + SmallVector staticOffsets; + SmallVector staticSizes; + SmallVector staticStrides; + SmallVector dynamicOffsets; + SmallVector dynamicSizes; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + auto sourceRankedTensorType = llvm::cast(source.getType()); + // Structuring implementation this way avoids duplication between builders. + if (!resultType) { + resultType = llvm::cast(ExtractSliceOp::inferResultType( + sourceRankedTensorType, staticOffsets, staticSizes, staticStrides)); + } + result.addAttributes(attrs); + build(b, result, {resultType, outSourceType}, source, dynamicOffsets, + dynamicSizes, dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), + b.getDenseI64ArrayAttr(staticSizes), + b.getDenseI64ArrayAttr(staticStrides)); +} +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + build(b, result, resultType, cast(source.getType()), source, + offsets, sizes, strides, attrs); +} + +/// Build an ExtractSliceOp with mixed static and dynamic entries and inferred +/// result type. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + build(b, result, RankedTensorType(), cast(source.getType()), + source, offsets, sizes, strides, attrs); +} + +/// Build an ExtractSliceOp with mixed static and dynamic entries packed into +/// a Range vector. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, + ArrayRef ranges, + ArrayRef attrs) { + auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); + build(b, result, RankedTensorType(), cast(source.getType()), + source, offsets, sizes, strides, attrs); +} + +/// Build an ExtractSliceOp with dynamic entries and custom result type. If +/// the type passed is nullptr, it is inferred. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, + RankedTensorType outSourceType, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, + ArrayRef /*attrs*/) { + SmallVector offsetValues = llvm::to_vector<4>( + llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); + SmallVector sizeValues = llvm::to_vector<4>( + llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = llvm::to_vector<4>( + llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + build(b, result, resultType, outSourceType, source, offsetValues, sizeValues, + strideValues); +} +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + build(b, result, resultType, resultType, source, offsets, sizes, strides, + attrs); +} + +/// Build an ExtractSliceOp with dynamic entries and inferred result type. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + build(b, result, RankedTensorType(), cast(source.getType()), + source, offsets, sizes, strides, attrs); +} + +/// Verifier for ExtractSliceOp. +LogicalResult ExtractSliceOp::verify() { + RankedTensorType sourceType = getSourceType(); + + // Verify result type against inferred type. + RankedTensorType expectedType = ExtractSliceOp::inferResultType( + sourceType, getMixedOffsets(), getMixedSizes(), getMixedStrides()); + SliceVerificationResult result = isRankReducedType(expectedType, getType()); + if (result != SliceVerificationResult::Success) { + return produceSliceErrorMsg(result, *this, expectedType); + } + + // Verify that offsets, sizes, strides do not run out-of-bounds with respect + // to the source tensor. + SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( + sourceType.getShape(), getStaticOffsets(), getStaticSizes(), + getStaticStrides(), /*generateErrorMessage=*/true); + if (!boundsResult.isValid) { + return getOperation()->emitError(boundsResult.errorMessage); + } + + return success(); +} + +llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { + return ::getDroppedDims(getType().getShape(), getMixedSizes()); +} + +FailureOr +ExtractSliceOp::rankReduceIfNeeded(OpBuilder& b, Location loc, Value value, + ArrayRef desiredShape) { + auto sourceTensorType = llvm::dyn_cast(value.getType()); + assert(sourceTensorType && "not a ranked tensor type"); + auto sourceShape = sourceTensorType.getShape(); + if (sourceShape.equals(desiredShape)) { + return value; + } + auto maybeRankReductionMask = + mlir::computeRankReductionMask(sourceShape, desiredShape); + if (!maybeRankReductionMask) { + return failure(); + } + return tensor::createCanonicalRankReducingExtractSliceOp( + b, loc, value, + RankedTensorType::Builder(sourceTensorType).setShape(desiredShape)); +} + +LogicalResult ExtractSliceOp::reifyResultShapes( + OpBuilder& /*builder*/, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { + reifiedReturnShapes.resize(1); + reifiedReturnShapes[0].reserve(getType().getRank()); + SmallVector mixedSizes = getMixedSizes(); + llvm::SmallBitVector droppedDims = getDroppedDims(); + for (const auto& size : enumerate(mixedSizes)) { + if (droppedDims.test(size.index())) { + continue; + } + reifiedReturnShapes[0].push_back(size.value()); + } + return success(); +} + +namespace { +/** + * @brief Rewrite pattern that pushes tensor.cast past tensor.extract_slice. + * + * @details + * This pattern rewrites a `qtensor.extract_slice` operation whose source + * operand is produced by a `tensor.cast`. When `canFoldIntoConsumerOp` + * evaluates to true, the cast operation is moved after the slice operation. + * + * Conceptually, the slice is applied to the original tensor before the + * cast, avoiding unnecessary intermediate casts. + * + * Example: + * %0 = tensor.cast %V : tensor<3x!qco.qubit> to tensor + * %1, %2 = tensor.extract_slice %0[0][2][1] + * : tensor to tensor<2x!qco.qubit> + * + * is rewritten into: + * + * %0, %1 = tensor.extract_slice %V[0][2][1] + * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + * %2 = tensor.cast %0 : tensor<2x!qco.qubit> to tensor<2x!qco.qubit> + * + * This effectively folds the cast into the consumer operation and enables + * further canonicalization opportunities. + */ +class ExtractSliceOpCastFolder final : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(ExtractSliceOp sliceOp, + PatternRewriter& rewriter) const override { + // Any constant operand, just return to let the constant folder kick in. + if (llvm::any_of(sliceOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) { + return failure(); + } + + auto castOp = sliceOp.getSource().getDefiningOp(); + if (!castOp) { + return failure(); + } + + if (!canFoldIntoConsumerOp(castOp)) { + return failure(); + } + + // Pattern does not apply if the produced op would not verify. + SliceBoundsVerificationResult sliceResult = verifyInBoundsSlice( + cast(castOp.getSource().getType()).getShape(), + sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), + sliceOp.getStaticStrides()); + if (!sliceResult.isValid) { + return failure(); + } + + // Create folded extract. + Location loc = sliceOp.getLoc(); + auto newResult = rewriter.create( + loc, sliceOp.getType(), castOp.getSource(), sliceOp.getOffsets(), + sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(), + sliceOp.getStaticSizes(), sliceOp.getStaticStrides()); + rewriter.replaceOp(sliceOp, newResult->getResult(0)); + rewriter.replaceOp(castOp, newResult->getResult(1)); + return success(); + } +}; + +/// Slice elements from `values` into `outValues`. `counts` represents the +/// numbers of elements to stride in the original values for each dimension. +/// The output values can be used to construct a DenseElementsAttr. +template +void sliceElements(IterTy values, ArrayRef counts, + ArrayRef offsets, ArrayRef sizes, + ArrayRef strides, + llvm::SmallVectorImpl* outValues) { + assert(offsets.size() == sizes.size()); + assert(offsets.size() == strides.size()); + if (offsets.empty()) { + return; + } + + int64_t offset = offsets.front(); + int64_t size = sizes.front(); + int64_t stride = strides.front(); + if (offsets.size() == 1) { + for (int64_t i = 0; i < size; ++i, offset += stride) { + outValues->push_back(*(values + offset)); + } + + return; + } + + for (int64_t i = 0; i < size; ++i, offset += stride) { + auto begin = values + offset * counts.front(); + sliceElements(begin, counts.drop_front(), + offsets.drop_front(), sizes.drop_front(), + strides.drop_front(), outValues); + } +} + +} // namespace + +/// Return the canonical type of the result of an extract_slice op. +struct SliceReturnTypeCanonicalizer { + RankedTensorType operator()(ExtractSliceOp op, + ArrayRef mixedOffsets, + ArrayRef mixedSizes, + ArrayRef mixedStrides) { + return ExtractSliceOp::inferCanonicalRankReducedResultType( + op.getType().getRank(), op.getSourceType(), mixedOffsets, mixedSizes, + mixedStrides); + } +}; + +/// A canonicalizer wrapper to replace ExtractSliceOps. +struct SliceCanonicalizer { + void operator()(PatternRewriter& rewriter, ExtractSliceOp op, + ExtractSliceOp newOp) { + Value replacement = newOp.getResult(); + Value outSource = newOp.getOutSource(); + if (replacement.getType() != op.getType()) { + replacement = rewriter.create(op.getLoc(), op.getType(), + replacement); + } + rewriter.replaceOp(op, {replacement, outSource}); + } +}; + +void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add< + OpWithOffsetSizesAndStridesConstantArgumentFolder< + ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>, + ExtractSliceOpCastFolder>(context); +} + +/// If we have an ExtractSliceOp consuming an InsertSliceOp with the same +/// slice, we can return the InsertSliceOp's source directly. +// TODO: This only checks the immediate producer; extend to go up the +// insert/extract chain if the slices are disjoint. +static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { + auto insertOp = extractOp.getSource().getDefiningOp(); + + auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; + if (insertOp && insertOp.getSource().getType() == extractOp.getType() && + insertOp.isSameAs(extractOp, isSame)) { + return insertOp.getSource(); + } + + return {}; +} + +LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, + SmallVectorImpl& results) { + + if (getSourceType() == getType() && + succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { + results.push_back(this->getSource()); + results.push_back(getSource()); + return success(); + } + if (Value slice = foldExtractAfterInsertSlice(*this)) { + results.push_back(slice); + results.push_back(getSource()); + return success(); + } + + return failure(); +} + +Value mlir::tensor::createCanonicalRankReducingExtractSliceOp( + OpBuilder& b, Location loc, Value tensor, RankedTensorType targetType) { + auto rankedTensorType = llvm::cast(tensor.getType()); + unsigned rank = rankedTensorType.getRank(); + SmallVector offsets(rank, b.getIndexAttr(0)); + SmallVector sizes = getMixedSizes(b, loc, tensor); + SmallVector strides(rank, b.getIndexAttr(1)); + return b.createOrFold(loc, targetType, tensor, + offsets, sizes, strides); +} 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..3524133ac4 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.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 +#include // for affine::AffineDialect +#include // for arith::ArithDialect +#include +#include // for complex::ComplexDialect +#include +#include +#include +#include +#include +#include +#include +#include +#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"); + Type resultType = RankedTensorType::get( + {static_cast(elements.size())}, elements.front().getType()); + build(builder, result, resultType, elements); +} + +namespace { + +struct ConvertFromElementsOpToTensorOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(qtensor::FromElementsOp fromElementsOp, + PatternRewriter& rewriter) const final { + + rewriter.replaceOpWithNewOp( + fromElementsOp, fromElementsOp.getElements()); + + return success(); + } +}; + +} // namespace + +void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} 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..a74812a646 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -0,0 +1,80 @@ +/* + * 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 // for affine::AffineDialect +#include // for arith::ArithDialect +#include +#include // for complex::ComplexDialect +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +namespace { + +struct ConvertInsertOpToTensorOp : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(qtensor::InsertOp insertOp, + PatternRewriter& rewriter) const final { + rewriter.replaceOpWithNewOp( + insertOp, insertOp.getScalar(), insertOp.getDest(), + insertOp.getIndices()); + return success(); + } +}; + +struct InsertFromExtractOp : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOp insertOp, + PatternRewriter& rewriter) const final { + auto extractOp = insertOp.getScalar().getDefiningOp(); + if (!extractOp) { + return failure(); + } + if (insertOp.getDest() != extractOp.getOutTensor()) { + return failure(); + } + if (insertOp.getIndices() != extractOp.getIndices()) { + return failure(); + } + + rewriter.replaceOp(insertOp, extractOp.getTensor()); + rewriter.eraseOp(extractOp); + return success(); + } +}; + +} // namespace + +void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} 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..61a9832690 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -0,0 +1,466 @@ +/* + * 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 // for affine::AffineDialect +#include // for arith::ArithDialect +#include +#include // for complex::ComplexDialect +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qtensor; + +void InsertSliceOp::getAsmResultNames( + function_ref setNameFn) { + setNameFn(getResult(), "inserted_slice"); +} + +// Build a InsertSliceOp with mixed static and dynamic entries. +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + SmallVector staticOffsets; + SmallVector staticSizes; + SmallVector staticStrides; + SmallVector dynamicOffsets; + SmallVector dynamicSizes; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + result.addAttributes(attrs); + build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, + dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), + b.getDenseI64ArrayAttr(staticSizes), + b.getDenseI64ArrayAttr(staticStrides)); +} + +/// Build an InsertSliceOp with mixed static and dynamic entries packed into a +/// Range vector. +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ArrayRef ranges, + ArrayRef attrs) { + auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); + build(b, result, source, dest, offsets, sizes, strides, attrs); +} + +// Build a InsertSliceOp with dynamic entries. +void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, + Value dest, ValueRange offsets, ValueRange sizes, + ValueRange strides, + ArrayRef /*attrs*/) { + SmallVector offsetValues = llvm::to_vector<4>( + llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); + SmallVector sizeValues = llvm::to_vector<4>( + llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = llvm::to_vector<4>( + llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + build(b, result, source, dest, offsetValues, sizeValues, strideValues); +} + +/// Rank-reducing type verification for both InsertSliceOp and +/// ParallelInsertSliceOp. +static SliceVerificationResult verifyInsertSliceOp( + RankedTensorType srcType, RankedTensorType dstType, + ArrayRef staticOffsets, ArrayRef staticSizes, + ArrayRef staticStrides, RankedTensorType* expectedType = nullptr) { + // insert_slice is the inverse of extract_slice, use the same type + // inference. + RankedTensorType expected = ExtractSliceOp::inferResultType( + dstType, staticOffsets, staticSizes, staticStrides); + if (expectedType != nullptr) { + *expectedType = expected; + } + return isRankReducedType(expected, srcType); +} + +/// Verifier for InsertSliceOp. +LogicalResult InsertSliceOp::verify() { + // Verify result type against inferred type. + RankedTensorType expectedType; + SliceVerificationResult result = + verifyInsertSliceOp(getSourceType(), getType(), getStaticOffsets(), + getStaticSizes(), getStaticStrides(), &expectedType); + if (result != SliceVerificationResult::Success) { + return produceSliceErrorMsg(result, *this, expectedType); + } + + // Verify that offsets, sizes, strides do not run out-of-bounds with respect + // to the destination tensor. + SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( + getDestType().getShape(), getStaticOffsets(), getStaticSizes(), + getStaticStrides(), /*generateErrorMessage=*/true); + if (!boundsResult.isValid) { + return getOperation()->emitError(boundsResult.errorMessage); + } + + return success(); +} + +/// If we have two consecutive InsertSliceOp writing to the same slice, we +/// can mutate the second InsertSliceOp's destination to the first one's. +/// +/// Example: +/// +/// ```mlir +/// %0 = tensor.insert_slice %slice0 into %input[0, 0] [64, 64] [1, 1] +/// %1 = tensor.insert_slice %slice1 into %0[0, 0] [64, 64] [1, 1] +/// ``` +/// +/// folds into: +/// +/// ```mlir +/// %1 = tensor.insert_slice %slice1 into %input[0, 0] [64, 64] [1, 1] +/// ``` +/// +/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { + auto prevInsertOp = insertOp.getDest().getDefiningOp(); + + auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; + if (!prevInsertOp || + prevInsertOp.getSource().getType() != insertOp.getSource().getType() || + !prevInsertOp.isSameAs(insertOp, isSame)) { + return failure(); + } + + insertOp.getDestMutable().assign(prevInsertOp.getDest()); + return success(); +} + +/// Folds round-trip extract/insert slice op pairs. +/// Example: +/// ```mlir +/// %0 = tensor.extract_slice %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] +/// %1 = tensor.insert_slice %0 into %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] +/// ``` +/// can be folded into %val. +static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { + auto extractOp = insertOp.getSource().getDefiningOp(); + auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; + if (!extractOp || extractOp.getOutSource() != insertOp.getDest() || + !extractOp.isSameAs(insertOp, isSame)) { + return nullptr; + } + return extractOp.getSource(); +} + +OpFoldResult InsertSliceOp::fold(FoldAdaptor) { + if (getSourceType().hasStaticShape() && getType().hasStaticShape() && + getSourceType() == getType() && + succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { + return this->getSource(); + } + if (succeeded(foldInsertAfterInsertSlice(*this))) { + return getResult(); + } + if (auto result = foldInsertAfterExtractSlice(*this)) { + return result; + } + if (llvm::any_of(getMixedSizes(), isZeroInteger)) { + return getDest(); + } + return {}; +} + +LogicalResult InsertSliceOp::reifyResultShapes( + OpBuilder& builder, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { + reifiedReturnShapes.resize(1, SmallVector(getType().getRank())); + reifiedReturnShapes[0] = tensor::getMixedSizes(builder, getLoc(), getDest()); + return success(); +} + +namespace { +/// Pattern to rewrite a insert_slice op with constant arguments. +/// +/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +template +class InsertSliceOpConstantArgumentFolder final + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + PatternRewriter& rewriter) const override { + SmallVector mixedOffsets(insertSliceOp.getMixedOffsets()); + SmallVector mixedSizes(insertSliceOp.getMixedSizes()); + SmallVector mixedStrides(insertSliceOp.getMixedStrides()); + + // No constant operands were folded, just return; + if (failed(foldDynamicOffsetSizeList(mixedOffsets)) && + failed(foldDynamicOffsetSizeList(mixedSizes)) && + failed(foldDynamicStrideList(mixedStrides))) { + return failure(); + } + + // Pattern does not apply if the produced op would not verify. + SliceBoundsVerificationResult sliceResult = + verifyInBoundsSlice(insertSliceOp.getDest().getType().getShape(), + mixedOffsets, mixedSizes, mixedStrides); + if (!sliceResult.isValid) { + return failure(); + } + + // Create the new op in canonical form. + auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType( + insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(), + mixedOffsets, mixedSizes, mixedStrides); + Value toInsert = insertSliceOp.getSource(); + if (sourceType != insertSliceOp.getSourceType()) { + OpBuilder::InsertionGuard g(rewriter); + // The only difference between InsertSliceOp and ParallelInsertSliceOp + // is that the insertion point is just before the ParallelCombiningOp in + // the parallel case. + if (std::is_same_v) { + rewriter.setInsertionPoint(insertSliceOp->getParentOp()); + } + toInsert = rewriter.create(insertSliceOp.getLoc(), + sourceType, toInsert); + } + rewriter.replaceOpWithNewOp( + insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets, + mixedSizes, mixedStrides); + return success(); + } +}; + +/// Fold tensor_casts with insert_slice operations. If the source or +/// destination tensor is a tensor_cast that removes static type information, +/// the cast is folded into the insert_slice operation. E.g.: +/// +/// ```mlir +/// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor +/// %2 = tensor.insert_slice %1 into ... : tensor into ... +/// ``` +/// +/// folds into: +/// +/// ```mlir +/// %2 = tensor.insert_slice %0 into ... : tensor<8x16xf32> into ... +/// ``` +/// +/// Note: When folding a cast on the destination tensor, the result of the +/// insert_slice operation is casted to ensure that the type of the result did +/// not change. +/// +/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +template +struct InsertSliceOpCastFolder final : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + PatternRewriter& rewriter) const override { + if (llvm::any_of(insertSliceOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) { + return failure(); + } + + auto getSourceOfCastOp = [](Value v) -> std::optional { + auto castOp = v.getDefiningOp(); + if (!castOp || !canFoldIntoConsumerOp(castOp)) { + return std::nullopt; + } + return castOp.getSource(); + }; + std::optional sourceCastSource = + getSourceOfCastOp(insertSliceOp.getSource()); + std::optional destCastSource = + getSourceOfCastOp(insertSliceOp.getDest()); + if (!sourceCastSource && !destCastSource) { + return failure(); + } + + auto src = + (sourceCastSource ? *sourceCastSource : insertSliceOp.getSource()); + auto dst = (destCastSource ? *destCastSource : insertSliceOp.getDest()); + auto srcType = llvm::dyn_cast(src.getType()); + auto dstType = llvm::dyn_cast(dst.getType()); + if (!srcType || !dstType) { + return failure(); + } + + // The tensor.cast source could have additional static information not seen + // in the insert slice op static sizes, so we ignore dynamic dims when + // computing the rank reduction mask. + SmallVector staticSizes(insertSliceOp.getStaticSizes()); + auto rankReductionMask = computeRankReductionMask( + staticSizes, srcType.getShape(), /*matchDynamic=*/true); + if (!rankReductionMask.has_value()) { + return failure(); + } + // Replace dimensions in the insert slice op with corresponding static dims + // from the cast source type. If the insert slice sizes have static dims + // that are not static in the tensor.cast source (i.e., when the cast op + // casts a dynamic dim to static), the dim should not be replaced, and the + // pattern will fail later in `verifyInsertSliceOp`. + SmallVector mixedSizes(insertSliceOp.getMixedSizes()); + int64_t rankReducedIdx = 0; + for (auto [idx, size] : enumerate(staticSizes)) { + if (!rankReductionMask.value().contains(idx) && + !srcType.isDynamicDim(rankReducedIdx)) { + mixedSizes[idx] = getAsIndexOpFoldResult( + rewriter.getContext(), srcType.getDimSize(rankReducedIdx)); + size = srcType.getDimSize(rankReducedIdx++); + } + } + + // Pattern does not apply if the produced op would not verify. + if (verifyInsertSliceOp(srcType, dstType, insertSliceOp.getStaticOffsets(), + staticSizes, insertSliceOp.getStaticStrides()) != + SliceVerificationResult::Success) { + return failure(); + } + SliceBoundsVerificationResult sliceResult = + verifyInBoundsSlice(dstType.getShape(), insertSliceOp.getMixedOffsets(), + mixedSizes, insertSliceOp.getMixedStrides()); + if (!sliceResult.isValid) { + return failure(); + } + + Operation* replacement = rewriter.create( + insertSliceOp.getLoc(), src, dst, insertSliceOp.getMixedOffsets(), + mixedSizes, insertSliceOp.getMixedStrides()); + + // In the parallel case there is no result and so nothing to cast. + bool isParallelInsert = + std::is_same::value; + if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { + replacement = rewriter.create(insertSliceOp.getLoc(), + insertSliceOp.getDestType(), + replacement->getResult(0)); + } + rewriter.replaceOp(insertSliceOp, replacement->getResults()); + return success(); + } +}; + +/// If additional static type information can be deduced from a insert_slice's +/// size operands, insert an explicit cast of the op's source operand. This +/// enables other canonicalization patterns that are matching for tensor_cast +/// ops such as `ForOpTensorCastFolder` in SCF. +/// +/// Example: +/// +/// ```mlir +/// %r = tensor.insert_slice %0 into %1[...] [64, 64] [1, 1] +/// : tensor into ... +/// ``` +/// +/// folds into: +/// +/// ```mlir +/// %tmp = tensor.cast %0 : tensor to tensor<64x64xf32> +/// %r = tensor.insert_slice %tmp into %1[...] [64, 64] [1, 1] +/// : tensor<64x64xf32> into ... +/// ``` +/// +/// This patterns works with both InsertSliceOp and ParallelInsertSliceOp. +template +struct InsertSliceOpSourceCastInserter final + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + PatternRewriter& rewriter) const override { + RankedTensorType srcType = insertSliceOp.getSourceType(); + if (srcType.getRank() != insertSliceOp.getDestType().getRank()) { + return failure(); + } + SmallVector newSrcShape(srcType.getShape()); + for (int64_t i = 0; i < srcType.getRank(); ++i) { + if (std::optional constInt = + getConstantIntValue(insertSliceOp.getMixedSizes()[i])) { + // Bail on invalid IR. + if (*constInt < 0) { + return failure(); + } + newSrcShape[i] = *constInt; + } + } + if (!hasValidSizesOffsets(newSrcShape)) { + return failure(); + } + + RankedTensorType newSrcType = RankedTensorType::get( + newSrcShape, srcType.getElementType(), srcType.getEncoding()); + if (srcType == newSrcType || + !mlir::tensor::preservesStaticInformation(srcType, newSrcType) || + !tensor::CastOp::areCastCompatible(srcType, newSrcType)) { + return failure(); + } + + // newSrcType is: + // 1) Different from srcType. + // 2) "More static" than srcType. + // 3) Cast-compatible with srcType. + // Insert the cast. + OpBuilder::InsertionGuard g(rewriter); + // The only difference between InsertSliceOp and ParallelInsertSliceOp is + // that the insertion point is just before the ParallelCombiningOp in the + // parallel case. + if (std::is_same_v) { + rewriter.setInsertionPoint(insertSliceOp->getParentOp()); + } + Value cast = rewriter.create( + insertSliceOp.getLoc(), newSrcType, insertSliceOp.getSource()); + rewriter.replaceOpWithNewOp( + insertSliceOp, cast, insertSliceOp.getDest(), + insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(), + insertSliceOp.getMixedStrides()); + return success(); + } +}; +} // namespace + +llvm::SmallBitVector InsertSliceOp::getDroppedDims() { + return ::getDroppedDims(getSourceType().getShape(), getMixedSizes()); +} + +void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add, + InsertSliceOpCastFolder, + InsertSliceOpSourceCastInserter>(context); +} + +Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder& b, + Location loc, + Value tensor, + Value dest) { + auto rankedTensorType = llvm::cast(dest.getType()); + unsigned rank = rankedTensorType.getRank(); + SmallVector offsets(rank, b.getIndexAttr(0)); + SmallVector sizes = getMixedSizes(b, loc, dest); + SmallVector strides(rank, b.getIndexAttr(1)); + return b.createOrFold(loc, tensor, dest, offsets, + sizes, strides); +} diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 63b3ddde88..cd14e7ac2b 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -13,9 +13,6 @@ #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated #include -#include -#include -#include #include #include // for affine::AffineDialect #include // for arith::ArithDialect @@ -33,7 +30,6 @@ #include #include #include -#include // The following headers are needed for some template instantiations. // IWYU pragma: begin_keep @@ -59,1112 +55,6 @@ Operation* QTensorDialect::materializeConstant(OpBuilder& builder, return nullptr; } -/// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or -/// rank-extending tensor.insert_slice op. -static llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, - ArrayRef mixedSizes) { - llvm::SmallBitVector droppedDims(mixedSizes.size()); - int64_t shapePos = reducedShape.size() - 1; - - for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { - size_t idx = mixedSizes.size() - size.index() - 1; - // Rank-reduced dims must have a static unit dimension. - bool isStaticUnitSize = - isa(size.value()) && - llvm::cast(cast(size.value())).getInt() == 1; - - if (shapePos < 0) { - // There are no more dims in the reduced shape. All remaining sizes must - // be rank-reduced dims. - assert(isStaticUnitSize && "expected unit dim"); - droppedDims.set(idx); - continue; - } - - // Dim is preserved if the size is not a static 1. - if (!isStaticUnitSize) { - --shapePos; - continue; - } - - // Dim is preserved if the reduced shape dim is also 1. - if (reducedShape[shapePos] == 1) { - --shapePos; - continue; - } - - // Otherwise: Dim is dropped. - droppedDims.set(idx); - } - - assert(shapePos < 0 && "dimension mismatch"); - return droppedDims; -} - -//===----------------------------------------------------------------------===// -// ExtractOp -//===----------------------------------------------------------------------===// - -void ExtractOp::getAsmResultNames( - function_ref setNameFn) { - setNameFn(getResult(), "q_extracted"); -} - -LogicalResult ExtractOp::verify() { - // Verify the # indices match if we have a ranked type. - auto tensorType = llvm::cast(getTensor().getType()); - if (tensorType.getRank() != static_cast(getIndices().size())) { - return emitOpError("incorrect number of indices for extract_element"); - } - return success(); -} -struct ExtractFromTensorCast : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ExtractOp extract, - PatternRewriter& rewriter) const final { - auto tensorCast = extract.getTensor().getDefiningOp(); - if (!tensorCast) { - return failure(); - } - if (!llvm::isa(tensorCast.getSource().getType())) { - return failure(); - } - rewriter.replaceOpWithNewOp(extract, tensorCast.getSource(), - extract.getIndices()); - return success(); - } -}; - -/// If we have an ExtractOp consuming an InsertOp with the same -/// indices, we can return the InsertOp's scalar directly. -// TODO: This only checks the immediate producer; extend to go up the -// insert/extract chain if the slices are disjoint. -static Value foldExtractAfterInsert(ExtractOp extractOp) { - auto insertOp = extractOp.getTensor().getDefiningOp(); - - auto isSame = [](Value a, Value b) { - return getAsOpFoldResult(a) == getAsOpFoldResult(b); - }; - if (insertOp && insertOp.getScalar().getType() == extractOp.getType(0) && - llvm::equal(insertOp.getIndices(), extractOp.getIndices(), isSame)) { - return insertOp.getScalar(); - } - return {}; -} - -LogicalResult ExtractOp::fold(FoldAdaptor adaptor, - SmallVectorImpl& results) { - // Collect the constant indices into the tensor. - SmallVector indices; - for (Attribute indice : adaptor.getIndices()) { - if (!indice || !llvm::isa(indice)) { - return failure(); - } - indices.push_back(llvm::cast(indice).getInt()); - } - - if (Value result = foldExtractAfterInsert(*this)) { - results.push_back(result); - results.push_back(getTensor()); - return success(); - } - - return failure(); -} - -void ExtractOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} - -//===----------------------------------------------------------------------===// -// ExtractSliceOp -//===----------------------------------------------------------------------===// - -void ExtractSliceOp::getAsmResultNames( - function_ref setNameFn) { - setNameFn(getResult(), "q_extracted_slice"); -} - -/// An extract_slice result type can be inferred, when it is not -/// rank-reduced, from the source type and the static representation of -/// offsets, sizes and strides. Special sentinels encode the dynamic case. -RankedTensorType ExtractSliceOp::inferResultType( - RankedTensorType sourceTensorType, ArrayRef /*staticOffsets*/, - ArrayRef staticSizes, ArrayRef /*staticStrides*/) { - // An extract_slice op may specify only a leading subset of offset/sizes/ - // strides in which case we complete with offset=0, sizes from memref type - // and strides=1. - assert(static_cast(staticSizes.size()) == - sourceTensorType.getRank() && - "unexpected staticSizes not equal to rank of source"); - return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), - sourceTensorType.getEncoding()); -} - -RankedTensorType ExtractSliceOp::inferResultType( - RankedTensorType sourceTensorType, ArrayRef /*offsets*/, - ArrayRef sizes, ArrayRef /*strides*/) { - SmallVector staticSizes; - std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes); - assert(static_cast(staticSizes.size()) == - sourceTensorType.getRank() && - "unexpected staticSizes not equal to rank of source"); - return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), - sourceTensorType.getEncoding()); -} - -/// If the rank is reduced (i.e. the desiredResultRank is smaller than the -/// number of sizes), drop as many size 1 as needed to produce an inferred -/// type with the desired rank. -/// -/// Note that there may be multiple ways to compute this rank-reduced type: -/// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. -/// -/// To disambiguate, this function always drops the first 1 sizes occurrences. -RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( - unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides) { - // Type inferred in the absence of rank-reducing behavior. - auto inferredType = llvm::cast( - inferResultType(sourceRankedTensorType, offsets, sizes, strides)); - int rankDiff = inferredType.getRank() - desiredResultRank; - if (rankDiff > 0) { - auto shape = inferredType.getShape(); - llvm::SmallBitVector dimsToProject = - getPositionsOfShapeOne(rankDiff, shape); - SmallVector projectedShape; - // Best effort rank-reducing: drop 1s in order. - for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) { - if (!dimsToProject.test(pos)) { - projectedShape.push_back(shape[pos]); - } - } - inferredType = - RankedTensorType::get(projectedShape, inferredType.getElementType()); - } - return inferredType; -} - -RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( - unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides) { - SmallVector staticOffsets; - SmallVector staticSizes; - SmallVector staticStrides; - SmallVector dynamicOffsets; - SmallVector dynamicSizes; - SmallVector dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); - return ExtractSliceOp::inferCanonicalRankReducedResultType( - desiredResultRank, sourceRankedTensorType, staticOffsets, staticSizes, - staticStrides); -} - -/// Build an ExtractSliceOp with mixed static and dynamic entries and custom -/// result type. If the type passed is nullptr, it is inferred. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, - RankedTensorType outSourceType, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - SmallVector staticOffsets; - SmallVector staticSizes; - SmallVector staticStrides; - SmallVector dynamicOffsets; - SmallVector dynamicSizes; - SmallVector dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); - auto sourceRankedTensorType = llvm::cast(source.getType()); - // Structuring implementation this way avoids duplication between builders. - if (!resultType) { - resultType = llvm::cast(ExtractSliceOp::inferResultType( - sourceRankedTensorType, staticOffsets, staticSizes, staticStrides)); - } - result.addAttributes(attrs); - build(b, result, {resultType, outSourceType}, source, dynamicOffsets, - dynamicSizes, dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), - b.getDenseI64ArrayAttr(staticSizes), - b.getDenseI64ArrayAttr(staticStrides)); -} -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - build(b, result, resultType, cast(source.getType()), source, - offsets, sizes, strides, attrs); -} - -/// Build an ExtractSliceOp with mixed static and dynamic entries and inferred -/// result type. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, strides, attrs); -} - -/// Build an ExtractSliceOp with mixed static and dynamic entries packed into -/// a Range vector. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, - ArrayRef ranges, - ArrayRef attrs) { - auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); - build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, strides, attrs); -} - -/// Build an ExtractSliceOp with dynamic entries and custom result type. If -/// the type passed is nullptr, it is inferred. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, - RankedTensorType outSourceType, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, - ArrayRef /*attrs*/) { - SmallVector offsetValues = llvm::to_vector<4>( - llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); - SmallVector sizeValues = llvm::to_vector<4>( - llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = llvm::to_vector<4>( - llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, resultType, outSourceType, source, offsetValues, sizeValues, - strideValues); -} -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { - build(b, result, resultType, resultType, source, offsets, sizes, strides, - attrs); -} - -/// Build an ExtractSliceOp with dynamic entries and inferred result type. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { - build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, strides, attrs); -} - -static LogicalResult produceSliceErrorMsg(SliceVerificationResult result, - Operation* op, - RankedTensorType expectedType) { - switch (result) { - case SliceVerificationResult::Success: - return success(); - case SliceVerificationResult::RankTooLarge: - return op->emitError("expected rank to be smaller or equal to ") - << "the other rank. "; - case SliceVerificationResult::SizeMismatch: - return op->emitError("expected type to be ") - << expectedType << " or a rank-reduced version. (size mismatch) "; - case SliceVerificationResult::ElemTypeMismatch: - return op->emitError("expected element type to be ") - << expectedType.getElementType(); - default: - llvm_unreachable("unexpected extract_slice op verification result"); - } -} - -/// Verifier for ExtractSliceOp. -LogicalResult ExtractSliceOp::verify() { - RankedTensorType sourceType = getSourceType(); - - // Verify result type against inferred type. - RankedTensorType expectedType = ExtractSliceOp::inferResultType( - sourceType, getMixedOffsets(), getMixedSizes(), getMixedStrides()); - SliceVerificationResult result = isRankReducedType(expectedType, getType()); - if (result != SliceVerificationResult::Success) { - return produceSliceErrorMsg(result, *this, expectedType); - } - - // Verify that offsets, sizes, strides do not run out-of-bounds with respect - // to the source tensor. - SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( - sourceType.getShape(), getStaticOffsets(), getStaticSizes(), - getStaticStrides(), /*generateErrorMessage=*/true); - if (!boundsResult.isValid) { - return getOperation()->emitError(boundsResult.errorMessage); - } - - return success(); -} - -llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { - return ::getDroppedDims(getType().getShape(), getMixedSizes()); -} - -FailureOr -ExtractSliceOp::rankReduceIfNeeded(OpBuilder& b, Location loc, Value value, - ArrayRef desiredShape) { - auto sourceTensorType = llvm::dyn_cast(value.getType()); - assert(sourceTensorType && "not a ranked tensor type"); - auto sourceShape = sourceTensorType.getShape(); - if (sourceShape.equals(desiredShape)) { - return value; - } - auto maybeRankReductionMask = - mlir::computeRankReductionMask(sourceShape, desiredShape); - if (!maybeRankReductionMask) { - return failure(); - } - return tensor::createCanonicalRankReducingExtractSliceOp( - b, loc, value, - RankedTensorType::Builder(sourceTensorType).setShape(desiredShape)); -} - -LogicalResult ExtractSliceOp::reifyResultShapes( - OpBuilder& /*builder*/, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { - reifiedReturnShapes.resize(1); - reifiedReturnShapes[0].reserve(getType().getRank()); - SmallVector mixedSizes = getMixedSizes(); - llvm::SmallBitVector droppedDims = getDroppedDims(); - for (const auto& size : enumerate(mixedSizes)) { - if (droppedDims.test(size.index())) { - continue; - } - reifiedReturnShapes[0].push_back(size.value()); - } - return success(); -} - -namespace { -/** - * @brief Rewrite pattern that pushes tensor.cast past tensor.extract_slice. - * - * @details - * This pattern rewrites a `qtensor.extract_slice` operation whose source - * operand is produced by a `tensor.cast`. When `canFoldIntoConsumerOp` - * evaluates to true, the cast operation is moved after the slice operation. - * - * Conceptually, the slice is applied to the original tensor before the - * cast, avoiding unnecessary intermediate casts. - * - * Example: - * %0 = tensor.cast %V : tensor<3x!qco.qubit> to tensor - * %1, %2 = tensor.extract_slice %0[0][2][1] - * : tensor to tensor<2x!qco.qubit> - * - * is rewritten into: - * - * %0, %1 = tensor.extract_slice %V[0][2][1] - * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> - * %2 = tensor.cast %0 : tensor<2x!qco.qubit> to tensor<2x!qco.qubit> - * - * This effectively folds the cast into the consumer operation and enables - * further canonicalization opportunities. - */ -class ExtractSliceOpCastFolder final : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ExtractSliceOp sliceOp, - PatternRewriter& rewriter) const override { - // Any constant operand, just return to let the constant folder kick in. - if (llvm::any_of(sliceOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) { - return failure(); - } - - auto castOp = sliceOp.getSource().getDefiningOp(); - if (!castOp) { - return failure(); - } - - if (!canFoldIntoConsumerOp(castOp)) { - return failure(); - } - - // Pattern does not apply if the produced op would not verify. - SliceBoundsVerificationResult sliceResult = verifyInBoundsSlice( - cast(castOp.getSource().getType()).getShape(), - sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), - sliceOp.getStaticStrides()); - if (!sliceResult.isValid) { - return failure(); - } - - // Create folded extract. - Location loc = sliceOp.getLoc(); - auto newResult = rewriter.create( - loc, sliceOp.getType(), castOp.getSource(), sliceOp.getOffsets(), - sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(), - sliceOp.getStaticSizes(), sliceOp.getStaticStrides()); - rewriter.replaceOp(sliceOp, newResult->getResult(0)); - rewriter.replaceOp(castOp, newResult->getResult(1)); - return success(); - } -}; - -/// Slice elements from `values` into `outValues`. `counts` represents the -/// numbers of elements to stride in the original values for each dimension. -/// The output values can be used to construct a DenseElementsAttr. -template -void sliceElements(IterTy values, ArrayRef counts, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides, - llvm::SmallVectorImpl* outValues) { - assert(offsets.size() == sizes.size()); - assert(offsets.size() == strides.size()); - if (offsets.empty()) { - return; - } - - int64_t offset = offsets.front(); - int64_t size = sizes.front(); - int64_t stride = strides.front(); - if (offsets.size() == 1) { - for (int64_t i = 0; i < size; ++i, offset += stride) { - outValues->push_back(*(values + offset)); - } - - return; - } - - for (int64_t i = 0; i < size; ++i, offset += stride) { - auto begin = values + offset * counts.front(); - sliceElements(begin, counts.drop_front(), - offsets.drop_front(), sizes.drop_front(), - strides.drop_front(), outValues); - } -} - -} // namespace - -/// Return the canonical type of the result of an extract_slice op. -struct SliceReturnTypeCanonicalizer { - RankedTensorType operator()(ExtractSliceOp op, - ArrayRef mixedOffsets, - ArrayRef mixedSizes, - ArrayRef mixedStrides) { - return ExtractSliceOp::inferCanonicalRankReducedResultType( - op.getType().getRank(), op.getSourceType(), mixedOffsets, mixedSizes, - mixedStrides); - } -}; - -/// A canonicalizer wrapper to replace ExtractSliceOps. -struct SliceCanonicalizer { - void operator()(PatternRewriter& rewriter, ExtractSliceOp op, - ExtractSliceOp newOp) { - Value replacement = newOp.getResult(); - Value outSource = newOp.getOutSource(); - if (replacement.getType() != op.getType()) { - replacement = rewriter.create(op.getLoc(), op.getType(), - replacement); - } - rewriter.replaceOp(op, {replacement, outSource}); - } -}; - -void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add< - OpWithOffsetSizesAndStridesConstantArgumentFolder< - ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>, - ExtractSliceOpCastFolder>(context); -} - -// -static LogicalResult -foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, - ShapedType shapedType) { - OpBuilder b(op.getContext()); - for (OpFoldResult opFold : op.getMixedOffsets()) { - if (getConstantIntValue(opFold) != static_cast(0)) { - return failure(); - } - } - // Rank-reducing noops only need to inspect the leading dimensions: - // llvm::zip is appropriate. - auto shape = shapedType.getShape(); - for (auto it : llvm::zip(op.getMixedSizes(), shape)) { - if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) { - return failure(); - } - } - for (OpFoldResult opFold : op.getMixedStrides()) { - if (getConstantIntValue(opFold) != static_cast(1)) { - return failure(); - } - } - return success(); -} - -/// If we have an ExtractSliceOp consuming an InsertSliceOp with the same -/// slice, we can return the InsertSliceOp's source directly. -// TODO: This only checks the immediate producer; extend to go up the -// insert/extract chain if the slices are disjoint. -static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { - auto insertOp = extractOp.getSource().getDefiningOp(); - - auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; - if (insertOp && insertOp.getSource().getType() == extractOp.getType() && - insertOp.isSameAs(extractOp, isSame)) { - return insertOp.getSource(); - } - - return {}; -} - -LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, - SmallVectorImpl& results) { - - if (getSourceType() == getType() && - succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { - results.push_back(this->getSource()); - results.push_back(getSource()); - return success(); - } - if (Value slice = foldExtractAfterInsertSlice(*this)) { - results.push_back(slice); - results.push_back(getSource()); - return success(); - } - - return failure(); -} - -Value mlir::tensor::createCanonicalRankReducingExtractSliceOp( - OpBuilder& b, Location loc, Value tensor, RankedTensorType targetType) { - auto rankedTensorType = llvm::cast(tensor.getType()); - unsigned rank = rankedTensorType.getRank(); - SmallVector offsets(rank, b.getIndexAttr(0)); - SmallVector sizes = getMixedSizes(b, loc, tensor); - SmallVector strides(rank, b.getIndexAttr(1)); - return b.createOrFold(loc, targetType, tensor, - offsets, sizes, strides); -} - -//===----------------------------------------------------------------------===// -// FromElementsOp -//===----------------------------------------------------------------------===// - -void FromElementsOp::build(OpBuilder& builder, OperationState& result, - ValueRange elements) { - assert(!elements.empty() && "expected at least one element"); - Type resultType = RankedTensorType::get( - {static_cast(elements.size())}, elements.front().getType()); - build(builder, result, resultType, elements); -} - -namespace { - -struct ConvertFromElementsOpToTensorOp - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(qtensor::FromElementsOp fromElementsOp, - PatternRewriter& rewriter) const final { - - rewriter.replaceOpWithNewOp( - fromElementsOp, fromElementsOp.getElements()); - - return success(); - } -}; - -} // namespace - -void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} - -//===----------------------------------------------------------------------===// -// InsertOp -//===----------------------------------------------------------------------===// - -namespace { - -struct ConvertInsertOpToTensorOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(qtensor::InsertOp insertOp, - PatternRewriter& rewriter) const final { - rewriter.replaceOpWithNewOp( - insertOp, insertOp.getScalar(), insertOp.getDest(), - insertOp.getIndices()); - return success(); - } -}; - -struct InsertFromExtractOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOp insertOp, - PatternRewriter& rewriter) const final { - auto extractOp = insertOp.getScalar().getDefiningOp(); - if (!extractOp) { - return failure(); - } - if (insertOp.getDest() != extractOp.getOutTensor()) { - return failure(); - } - if (insertOp.getIndices() != extractOp.getIndices()) { - return failure(); - } - - rewriter.replaceOp(insertOp, extractOp.getTensor()); - rewriter.eraseOp(extractOp); - return success(); - } -}; - -} // namespace - -void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} - -//===----------------------------------------------------------------------===// -// InsertSliceOp -//===----------------------------------------------------------------------===// -void InsertSliceOp::getAsmResultNames( - function_ref setNameFn) { - setNameFn(getResult(), "inserted_slice"); -} - -// Build a InsertSliceOp with mixed static and dynamic entries. -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - SmallVector staticOffsets; - SmallVector staticSizes; - SmallVector staticStrides; - SmallVector dynamicOffsets; - SmallVector dynamicSizes; - SmallVector dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); - result.addAttributes(attrs); - build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, - dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), - b.getDenseI64ArrayAttr(staticSizes), - b.getDenseI64ArrayAttr(staticStrides)); -} - -/// Build an InsertSliceOp with mixed static and dynamic entries packed into a -/// Range vector. -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ArrayRef ranges, - ArrayRef attrs) { - auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); - build(b, result, source, dest, offsets, sizes, strides, attrs); -} - -// Build a InsertSliceOp with dynamic entries. -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ValueRange offsets, ValueRange sizes, - ValueRange strides, - ArrayRef /*attrs*/) { - SmallVector offsetValues = llvm::to_vector<4>( - llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); - SmallVector sizeValues = llvm::to_vector<4>( - llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = llvm::to_vector<4>( - llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, source, dest, offsetValues, sizeValues, strideValues); -} - -/// Rank-reducing type verification for both InsertSliceOp and -/// ParallelInsertSliceOp. -static SliceVerificationResult verifyInsertSliceOp( - RankedTensorType srcType, RankedTensorType dstType, - ArrayRef staticOffsets, ArrayRef staticSizes, - ArrayRef staticStrides, RankedTensorType* expectedType = nullptr) { - // insert_slice is the inverse of extract_slice, use the same type - // inference. - RankedTensorType expected = ExtractSliceOp::inferResultType( - dstType, staticOffsets, staticSizes, staticStrides); - if (expectedType != nullptr) { - *expectedType = expected; - } - return isRankReducedType(expected, srcType); -} - -/// Verifier for InsertSliceOp. -LogicalResult InsertSliceOp::verify() { - // Verify result type against inferred type. - RankedTensorType expectedType; - SliceVerificationResult result = - verifyInsertSliceOp(getSourceType(), getType(), getStaticOffsets(), - getStaticSizes(), getStaticStrides(), &expectedType); - if (result != SliceVerificationResult::Success) { - return produceSliceErrorMsg(result, *this, expectedType); - } - - // Verify that offsets, sizes, strides do not run out-of-bounds with respect - // to the destination tensor. - SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( - getDestType().getShape(), getStaticOffsets(), getStaticSizes(), - getStaticStrides(), /*generateErrorMessage=*/true); - if (!boundsResult.isValid) { - return getOperation()->emitError(boundsResult.errorMessage); - } - - return success(); -} - -/// If we have two consecutive InsertSliceOp writing to the same slice, we -/// can mutate the second InsertSliceOp's destination to the first one's. -/// -/// Example: -/// -/// ```mlir -/// %0 = tensor.insert_slice %slice0 into %input[0, 0] [64, 64] [1, 1] -/// %1 = tensor.insert_slice %slice1 into %0[0, 0] [64, 64] [1, 1] -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %1 = tensor.insert_slice %slice1 into %input[0, 0] [64, 64] [1, 1] -/// ``` -/// -/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. -static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { - auto prevInsertOp = insertOp.getDest().getDefiningOp(); - - auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; - if (!prevInsertOp || - prevInsertOp.getSource().getType() != insertOp.getSource().getType() || - !prevInsertOp.isSameAs(insertOp, isSame)) { - return failure(); - } - - insertOp.getDestMutable().assign(prevInsertOp.getDest()); - return success(); -} - -/// Folds round-trip extract/insert slice op pairs. -/// Example: -/// ```mlir -/// %0 = tensor.extract_slice %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] -/// %1 = tensor.insert_slice %0 into %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] -/// ``` -/// can be folded into %val. -static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { - auto extractOp = insertOp.getSource().getDefiningOp(); - auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; - if (!extractOp || extractOp.getOutSource() != insertOp.getDest() || - !extractOp.isSameAs(insertOp, isSame)) { - return nullptr; - } - return extractOp.getSource(); -} - -OpFoldResult InsertSliceOp::fold(FoldAdaptor) { - if (getSourceType().hasStaticShape() && getType().hasStaticShape() && - getSourceType() == getType() && - succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { - return this->getSource(); - } - if (succeeded(foldInsertAfterInsertSlice(*this))) { - return getResult(); - } - if (auto result = foldInsertAfterExtractSlice(*this)) { - return result; - } - if (llvm::any_of(getMixedSizes(), isZeroInteger)) { - return getDest(); - } - return {}; -} - -LogicalResult InsertSliceOp::reifyResultShapes( - OpBuilder& builder, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { - reifiedReturnShapes.resize(1, SmallVector(getType().getRank())); - reifiedReturnShapes[0] = tensor::getMixedSizes(builder, getLoc(), getDest()); - return success(); -} - -namespace { -/// Pattern to rewrite a insert_slice op with constant arguments. -/// -/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. -template -class InsertSliceOpConstantArgumentFolder final - : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, - PatternRewriter& rewriter) const override { - SmallVector mixedOffsets(insertSliceOp.getMixedOffsets()); - SmallVector mixedSizes(insertSliceOp.getMixedSizes()); - SmallVector mixedStrides(insertSliceOp.getMixedStrides()); - - // No constant operands were folded, just return; - if (failed(foldDynamicOffsetSizeList(mixedOffsets)) && - failed(foldDynamicOffsetSizeList(mixedSizes)) && - failed(foldDynamicStrideList(mixedStrides))) { - return failure(); - } - - // Pattern does not apply if the produced op would not verify. - SliceBoundsVerificationResult sliceResult = - verifyInBoundsSlice(insertSliceOp.getDest().getType().getShape(), - mixedOffsets, mixedSizes, mixedStrides); - if (!sliceResult.isValid) { - return failure(); - } - - // Create the new op in canonical form. - auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType( - insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(), - mixedOffsets, mixedSizes, mixedStrides); - Value toInsert = insertSliceOp.getSource(); - if (sourceType != insertSliceOp.getSourceType()) { - OpBuilder::InsertionGuard g(rewriter); - // The only difference between InsertSliceOp and ParallelInsertSliceOp - // is that the insertion point is just before the ParallelCombiningOp in - // the parallel case. - if (std::is_same_v) { - rewriter.setInsertionPoint(insertSliceOp->getParentOp()); - } - toInsert = rewriter.create(insertSliceOp.getLoc(), - sourceType, toInsert); - } - rewriter.replaceOpWithNewOp( - insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets, - mixedSizes, mixedStrides); - return success(); - } -}; - -/// Fold tensor_casts with insert_slice operations. If the source or -/// destination tensor is a tensor_cast that removes static type information, -/// the cast is folded into the insert_slice operation. E.g.: -/// -/// ```mlir -/// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor -/// %2 = tensor.insert_slice %1 into ... : tensor into ... -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %2 = tensor.insert_slice %0 into ... : tensor<8x16xf32> into ... -/// ``` -/// -/// Note: When folding a cast on the destination tensor, the result of the -/// insert_slice operation is casted to ensure that the type of the result did -/// not change. -/// -/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. -template -struct InsertSliceOpCastFolder final : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, - PatternRewriter& rewriter) const override { - if (llvm::any_of(insertSliceOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) { - return failure(); - } - - auto getSourceOfCastOp = [](Value v) -> std::optional { - auto castOp = v.getDefiningOp(); - if (!castOp || !canFoldIntoConsumerOp(castOp)) { - return std::nullopt; - } - return castOp.getSource(); - }; - std::optional sourceCastSource = - getSourceOfCastOp(insertSliceOp.getSource()); - std::optional destCastSource = - getSourceOfCastOp(insertSliceOp.getDest()); - if (!sourceCastSource && !destCastSource) { - return failure(); - } - - auto src = - (sourceCastSource ? *sourceCastSource : insertSliceOp.getSource()); - auto dst = (destCastSource ? *destCastSource : insertSliceOp.getDest()); - auto srcType = llvm::dyn_cast(src.getType()); - auto dstType = llvm::dyn_cast(dst.getType()); - if (!srcType || !dstType) { - return failure(); - } - - // The tensor.cast source could have additional static information not seen - // in the insert slice op static sizes, so we ignore dynamic dims when - // computing the rank reduction mask. - SmallVector staticSizes(insertSliceOp.getStaticSizes()); - auto rankReductionMask = computeRankReductionMask( - staticSizes, srcType.getShape(), /*matchDynamic=*/true); - if (!rankReductionMask.has_value()) { - return failure(); - } - // Replace dimensions in the insert slice op with corresponding static dims - // from the cast source type. If the insert slice sizes have static dims - // that are not static in the tensor.cast source (i.e., when the cast op - // casts a dynamic dim to static), the dim should not be replaced, and the - // pattern will fail later in `verifyInsertSliceOp`. - SmallVector mixedSizes(insertSliceOp.getMixedSizes()); - int64_t rankReducedIdx = 0; - for (auto [idx, size] : enumerate(staticSizes)) { - if (!rankReductionMask.value().contains(idx) && - !srcType.isDynamicDim(rankReducedIdx)) { - mixedSizes[idx] = getAsIndexOpFoldResult( - rewriter.getContext(), srcType.getDimSize(rankReducedIdx)); - size = srcType.getDimSize(rankReducedIdx++); - } - } - - // Pattern does not apply if the produced op would not verify. - if (verifyInsertSliceOp(srcType, dstType, insertSliceOp.getStaticOffsets(), - staticSizes, insertSliceOp.getStaticStrides()) != - SliceVerificationResult::Success) { - return failure(); - } - SliceBoundsVerificationResult sliceResult = - verifyInBoundsSlice(dstType.getShape(), insertSliceOp.getMixedOffsets(), - mixedSizes, insertSliceOp.getMixedStrides()); - if (!sliceResult.isValid) { - return failure(); - } - - Operation* replacement = rewriter.create( - insertSliceOp.getLoc(), src, dst, insertSliceOp.getMixedOffsets(), - mixedSizes, insertSliceOp.getMixedStrides()); - - // In the parallel case there is no result and so nothing to cast. - bool isParallelInsert = - std::is_same::value; - if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { - replacement = rewriter.create(insertSliceOp.getLoc(), - insertSliceOp.getDestType(), - replacement->getResult(0)); - } - rewriter.replaceOp(insertSliceOp, replacement->getResults()); - return success(); - } -}; - -/// If additional static type information can be deduced from a insert_slice's -/// size operands, insert an explicit cast of the op's source operand. This -/// enables other canonicalization patterns that are matching for tensor_cast -/// ops such as `ForOpTensorCastFolder` in SCF. -/// -/// Example: -/// -/// ```mlir -/// %r = tensor.insert_slice %0 into %1[...] [64, 64] [1, 1] -/// : tensor into ... -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %tmp = tensor.cast %0 : tensor to tensor<64x64xf32> -/// %r = tensor.insert_slice %tmp into %1[...] [64, 64] [1, 1] -/// : tensor<64x64xf32> into ... -/// ``` -/// -/// This patterns works with both InsertSliceOp and ParallelInsertSliceOp. -template -struct InsertSliceOpSourceCastInserter final - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, - PatternRewriter& rewriter) const override { - RankedTensorType srcType = insertSliceOp.getSourceType(); - if (srcType.getRank() != insertSliceOp.getDestType().getRank()) { - return failure(); - } - SmallVector newSrcShape(srcType.getShape()); - for (int64_t i = 0; i < srcType.getRank(); ++i) { - if (std::optional constInt = - getConstantIntValue(insertSliceOp.getMixedSizes()[i])) { - // Bail on invalid IR. - if (*constInt < 0) { - return failure(); - } - newSrcShape[i] = *constInt; - } - } - if (!hasValidSizesOffsets(newSrcShape)) { - return failure(); - } - - RankedTensorType newSrcType = RankedTensorType::get( - newSrcShape, srcType.getElementType(), srcType.getEncoding()); - if (srcType == newSrcType || - !mlir::tensor::preservesStaticInformation(srcType, newSrcType) || - !tensor::CastOp::areCastCompatible(srcType, newSrcType)) { - return failure(); - } - - // newSrcType is: - // 1) Different from srcType. - // 2) "More static" than srcType. - // 3) Cast-compatible with srcType. - // Insert the cast. - OpBuilder::InsertionGuard g(rewriter); - // The only difference between InsertSliceOp and ParallelInsertSliceOp is - // that the insertion point is just before the ParallelCombiningOp in the - // parallel case. - if (std::is_same_v) { - rewriter.setInsertionPoint(insertSliceOp->getParentOp()); - } - Value cast = rewriter.create( - insertSliceOp.getLoc(), newSrcType, insertSliceOp.getSource()); - rewriter.replaceOpWithNewOp( - insertSliceOp, cast, insertSliceOp.getDest(), - insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(), - insertSliceOp.getMixedStrides()); - return success(); - } -}; -} // namespace - -llvm::SmallBitVector InsertSliceOp::getDroppedDims() { - return ::getDroppedDims(getSourceType().getShape(), getMixedSizes()); -} - -void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add, - InsertSliceOpCastFolder, - InsertSliceOpSourceCastInserter>(context); -} - -Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder& b, - Location loc, - Value tensor, - Value dest) { - auto rankedTensorType = llvm::cast(dest.getType()); - unsigned rank = rankedTensorType.getRank(); - SmallVector offsets(rank, b.getIndexAttr(0)); - SmallVector sizes = getMixedSizes(b, loc, dest); - SmallVector strides(rank, b.getIndexAttr(1)); - return b.createOrFold(loc, tensor, dest, offsets, - sizes, strides); -} - //===----------------------------------------------------------------------===// // Common Canonicalizers and Folders. //===----------------------------------------------------------------------===// From 34af39467ac6efa41b935b9e1e6d600c3098fc54 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 3 Mar 2026 15:38:10 +0100 Subject: [PATCH 015/108] adjust headers --- .../Dialect/QTensor/IR/Operations/ExtractOp.cpp | 10 ---------- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 6 ------ .../QTensor/IR/Operations/FromTensorOp.cpp | 15 --------------- .../Dialect/QTensor/IR/Operations/InsertOp.cpp | 16 ---------------- .../QTensor/IR/Operations/InsertSliceOp.cpp | 6 ------ 5 files changed, 53 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index b80e6874ec..ebe1a1107f 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -11,15 +11,7 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include -#include -#include -#include #include -#include // for affine::AffineDialect -#include // for arith::ArithDialect -#include -#include // for complex::ComplexDialect -#include #include #include #include @@ -28,8 +20,6 @@ #include #include #include -#include -#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 340dae6dd5..e964c9b5e2 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -13,13 +13,8 @@ #include #include #include -#include #include -#include // for affine::AffineDialect -#include // for arith::ArithDialect #include -#include // for complex::ComplexDialect -#include #include #include #include @@ -28,7 +23,6 @@ #include #include #include -#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 3524133ac4..0ddbfa7272 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -10,26 +10,11 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include -#include -#include -#include -#include -#include // for affine::AffineDialect -#include // for arith::ArithDialect -#include -#include // for complex::ComplexDialect -#include #include -#include #include #include #include -#include #include -#include -#include -#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index a74812a646..b9cc5a6aef 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -10,26 +10,10 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include -#include -#include -#include -#include -#include // for affine::AffineDialect -#include // for arith::ArithDialect -#include -#include // for complex::ComplexDialect #include #include -#include -#include #include -#include -#include #include -#include -#include -#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 61a9832690..26d4439e95 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -13,13 +13,8 @@ #include #include #include -#include #include -#include // for affine::AffineDialect -#include // for arith::ArithDialect #include -#include // for complex::ComplexDialect -#include #include #include #include @@ -28,7 +23,6 @@ #include #include #include -#include #include #include #include From 85906f81467581d43c0b751443c5e07c0f4480dd Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 3 Mar 2026 15:38:54 +0100 Subject: [PATCH 016/108] refactor dialect structure --- .../mlir/Dialect/QTensor/IR/QTensorBase.td | 39 -------- .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 89 ++----------------- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 32 ++++++- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 88 ++++++++++++++++++ 4 files changed, 127 insertions(+), 121 deletions(-) delete mode 100644 mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td deleted file mode 100644 index 04df4da1ad..0000000000 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorBase.td +++ /dev/null @@ -1,39 +0,0 @@ -// 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 - -//===- TensorBase.td - Base definitions for tensor dialect -*- tablegen -*-===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// - -#ifndef TENSOR_BASE -#define TENSOR_BASE - -include "mlir/IR/OpBase.td" - -def QTensorDialect : Dialect { - let name = "qtensor"; - let cppNamespace = "::mlir::qtensor"; - - let description = [{ - TODO - }]; - - let hasCanonicalizer = 1; - let hasConstantMaterializer = 1; - let dependentDialects = [ - "affine::AffineDialect", - "arith::ArithDialect", - "complex::ComplexDialect", - ]; -} - -#endif // TENSOR_BASE diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h index 28b735bcab..4f7e832057 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -19,92 +19,19 @@ #define DIALECT_NAME_QTensor "qtensor" namespace mlir::qtensor { + /// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or /// rank-extending tensor.insert_slice op. -inline llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, - ArrayRef mixedSizes) { - llvm::SmallBitVector droppedDims(mixedSizes.size()); - int64_t shapePos = reducedShape.size() - 1; - - for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { - size_t idx = mixedSizes.size() - size.index() - 1; - // Rank-reduced dims must have a static unit dimension. - bool isStaticUnitSize = - isa(size.value()) && - llvm::cast(cast(size.value())).getInt() == 1; - - if (shapePos < 0) { - // There are no more dims in the reduced shape. All remaining sizes must - // be rank-reduced dims. - assert(isStaticUnitSize && "expected unit dim"); - droppedDims.set(idx); - continue; - } - - // Dim is preserved if the size is not a static 1. - if (!isStaticUnitSize) { - --shapePos; - continue; - } - - // Dim is preserved if the reduced shape dim is also 1. - if (reducedShape[shapePos] == 1) { - --shapePos; - continue; - } - - // Otherwise: Dim is dropped. - droppedDims.set(idx); - } - - assert(shapePos < 0 && "dimension mismatch"); - return droppedDims; -} +llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, + ArrayRef mixedSizes); -inline LogicalResult produceSliceErrorMsg(SliceVerificationResult result, - Operation* op, - RankedTensorType expectedType) { - switch (result) { - case SliceVerificationResult::Success: - return success(); - case SliceVerificationResult::RankTooLarge: - return op->emitError("expected rank to be smaller or equal to ") - << "the other rank. "; - case SliceVerificationResult::SizeMismatch: - return op->emitError("expected type to be ") - << expectedType << " or a rank-reduced version. (size mismatch) "; - case SliceVerificationResult::ElemTypeMismatch: - return op->emitError("expected element type to be ") - << expectedType.getElementType(); - default: - llvm_unreachable("unexpected extract_slice op verification result"); - } -} +LogicalResult produceSliceErrorMsg(SliceVerificationResult result, + Operation* op, + RankedTensorType expectedType); -inline LogicalResult +LogicalResult foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, - ShapedType shapedType) { - OpBuilder b(op.getContext()); - for (OpFoldResult opFold : op.getMixedOffsets()) { - if (getConstantIntValue(opFold) != static_cast(0)) { - return failure(); - } - } - // Rank-reducing noops only need to inspect the leading dimensions: - // llvm::zip is appropriate. - auto shape = shapedType.getShape(); - for (auto it : llvm::zip(op.getMixedSizes(), shape)) { - if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) { - return failure(); - } - } - for (OpFoldResult opFold : op.getMixedStrides()) { - if (getConstantIntValue(opFold) != static_cast(1)) { - return failure(); - } - } - return success(); -} + ShapedType shapedType); } // namespace mlir::qtensor //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index bf99b86b9b..824d33d5ba 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -17,7 +17,7 @@ #ifndef QTENSOROPS #define QTENSOROPS -include "mlir/Dialect/QTensor/IR/QTensorBase.td" +include "mlir/IR/OpBase.td" include "mlir/Interfaces/CastInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/DestinationStyleOpInterface.td" @@ -30,6 +30,31 @@ include "mlir/Interfaces/TilingInterface.td" include "mlir/Interfaces/ViewLikeInterface.td" include "mlir/IR/OpAsmInterface.td" +//===----------------------------------------------------------------------===// +// Dialect +//===----------------------------------------------------------------------===// + +def QTensorDialect : Dialect { + let name = "qtensor"; + let cppNamespace = "::mlir::qtensor"; + + let description = [{ + TODO + }]; + + let hasCanonicalizer = 1; + let hasConstantMaterializer = 1; + let dependentDialects = [ + "affine::AffineDialect", + "arith::ArithDialect", + "complex::ComplexDialect", + ]; +} + +//===----------------------------------------------------------------------===// +// Base Operation Classes +//===----------------------------------------------------------------------===// + class QTensorOp traits = []> : Op; @@ -54,6 +79,11 @@ class QTensorOpWithOffsetSizesAndStrides, Pure, diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index cd14e7ac2b..38d9246af7 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -55,6 +55,94 @@ Operation* QTensorDialect::materializeConstant(OpBuilder& builder, return nullptr; } +namespace mlir::qtensor { + +llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, + ArrayRef mixedSizes) { + llvm::SmallBitVector droppedDims(mixedSizes.size()); + int64_t shapePos = reducedShape.size() - 1; + + for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { + size_t idx = mixedSizes.size() - size.index() - 1; + // Rank-reduced dims must have a static unit dimension. + bool isStaticUnitSize = + isa(size.value()) && + llvm::cast(cast(size.value())).getInt() == 1; + + if (shapePos < 0) { + // There are no more dims in the reduced shape. All remaining sizes must + // be rank-reduced dims. + assert(isStaticUnitSize && "expected unit dim"); + droppedDims.set(idx); + continue; + } + + // Dim is preserved if the size is not a static 1. + if (!isStaticUnitSize) { + --shapePos; + continue; + } + + // Dim is preserved if the reduced shape dim is also 1. + if (reducedShape[shapePos] == 1) { + --shapePos; + continue; + } + + // Otherwise: Dim is dropped. + droppedDims.set(idx); + } + + assert(shapePos < 0 && "dimension mismatch"); + return droppedDims; +} + +LogicalResult produceSliceErrorMsg(SliceVerificationResult result, + Operation* op, + RankedTensorType expectedType) { + switch (result) { + case SliceVerificationResult::Success: + return success(); + case SliceVerificationResult::RankTooLarge: + return op->emitError("expected rank to be smaller or equal to ") + << "the other rank. "; + case SliceVerificationResult::SizeMismatch: + return op->emitError("expected type to be ") + << expectedType << " or a rank-reduced version. (size mismatch) "; + case SliceVerificationResult::ElemTypeMismatch: + return op->emitError("expected element type to be ") + << expectedType.getElementType(); + default: + llvm_unreachable("unexpected extract_slice op verification result"); + } +} + +LogicalResult +foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, + ShapedType shapedType) { + OpBuilder b(op.getContext()); + for (OpFoldResult opFold : op.getMixedOffsets()) { + if (getConstantIntValue(opFold) != static_cast(0)) { + return failure(); + } + } + // Rank-reducing noops only need to inspect the leading dimensions: + // llvm::zip is appropriate. + auto shape = shapedType.getShape(); + for (auto it : llvm::zip(op.getMixedSizes(), shape)) { + if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) { + return failure(); + } + } + for (OpFoldResult opFold : op.getMixedStrides()) { + if (getConstantIntValue(opFold) != static_cast(1)) { + return failure(); + } + } + return success(); +} +} // namespace mlir::qtensor + //===----------------------------------------------------------------------===// // Common Canonicalizers and Folders. //===----------------------------------------------------------------------===// From bdd1753042dfc6aeeea4f0ca92eb4e5f6b98b363 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 3 Mar 2026 15:45:51 +0100 Subject: [PATCH 017/108] fix issues in copied code --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 8 -------- mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 2 +- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 8 ++++---- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 824d33d5ba..87dc6e434f 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -6,14 +6,6 @@ // // Licensed under the MIT License -//===- QTensorOps.td - QTensor op definitions ----------------*- tablegen -*-===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// - #ifndef QTENSOROPS #define QTENSOROPS diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index e964c9b5e2..4980de5bcc 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -78,7 +78,7 @@ RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( // Type inferred in the absence of rank-reducing behavior. auto inferredType = llvm::cast( inferResultType(sourceRankedTensorType, offsets, sizes, strides)); - int rankDiff = inferredType.getRank() - desiredResultRank; + int64_t rankDiff = inferredType.getRank() - desiredResultRank; if (rankDiff > 0) { auto shape = inferredType.getShape(); llvm::SmallBitVector dimsToProject = diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 38d9246af7..a2414ee269 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -14,10 +14,10 @@ #include #include -#include // for affine::AffineDialect -#include // for arith::ArithDialect +#include +#include #include -#include // for complex::ComplexDialect +#include #include #include #include @@ -60,7 +60,7 @@ namespace mlir::qtensor { llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, ArrayRef mixedSizes) { llvm::SmallBitVector droppedDims(mixedSizes.size()); - int64_t shapePos = reducedShape.size() - 1; + int64_t shapePos = static_cast(reducedShape.size()) - 1; for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { size_t idx = mixedSizes.size() - size.index() - 1; From bd06b9d2cc92e778ed01e59ec34b82c0ddfebfc5 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 4 Mar 2026 10:09:50 +0100 Subject: [PATCH 018/108] smaller fixes --- mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 1 - mlir/unittests/programs/qco_programs.h | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt index 7ed18e3bb7..203af124cf 100644 --- a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -23,7 +23,6 @@ add_mlir_dialect_library( MLIRViewLikeInterface MLIRInferTypeOpInterface MLIRSideEffectInterfaces - Eigen3::Eigen DISABLE_INSTALL) mqt_mlir_target_use_project_options(MLIRQTensorDialect) diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 2cba4e458f..d2a5c76ffd 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -960,4 +960,5 @@ void extractSliceInsertSliceTensor(QCOProgramBuilder& b); /// Extracts a slice of qubits, a qubit from the slice, insert the qubit back to /// the slice and the slice back to the tensor immediately. void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b); + } // namespace mlir::qco From 5eea7b9c6409f14cfc6c3f0548ec3d8f5b541e7f Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 4 Mar 2026 10:27:28 +0100 Subject: [PATCH 019/108] remove memref interface for now --- mlir/include/mlir/Dialect/QC/IR/QCOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QC/IR/QCOps.td b/mlir/include/mlir/Dialect/QC/IR/QCOps.td index 035067bf3a..5956e39a56 100644 --- a/mlir/include/mlir/Dialect/QC/IR/QCOps.td +++ b/mlir/include/mlir/Dialect/QC/IR/QCOps.td @@ -60,7 +60,7 @@ class QCType traits = []> let mnemonic = typeMnemonic; } -def QubitType : QCType<"Qubit", "qubit", [MemRefElementTypeInterface]> { +def QubitType : QCType<"Qubit", "qubit"> { let summary = "QC qubit reference type"; let description = [{ The `!qc.qubit` type represents a reference to a quantum bit in the From 9817a9bba8c531cf352f250ada7fb6876ae53338 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 4 Mar 2026 11:08:13 +0100 Subject: [PATCH 020/108] remove duplicated code --- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 30 ------------------- .../QTensor/IR/Operations/InsertSliceOp.cpp | 13 -------- 2 files changed, 43 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 4980de5bcc..fff0868e7b 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -236,25 +236,6 @@ llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { return ::getDroppedDims(getType().getShape(), getMixedSizes()); } -FailureOr -ExtractSliceOp::rankReduceIfNeeded(OpBuilder& b, Location loc, Value value, - ArrayRef desiredShape) { - auto sourceTensorType = llvm::dyn_cast(value.getType()); - assert(sourceTensorType && "not a ranked tensor type"); - auto sourceShape = sourceTensorType.getShape(); - if (sourceShape.equals(desiredShape)) { - return value; - } - auto maybeRankReductionMask = - mlir::computeRankReductionMask(sourceShape, desiredShape); - if (!maybeRankReductionMask) { - return failure(); - } - return tensor::createCanonicalRankReducingExtractSliceOp( - b, loc, value, - RankedTensorType::Builder(sourceTensorType).setShape(desiredShape)); -} - LogicalResult ExtractSliceOp::reifyResultShapes( OpBuilder& /*builder*/, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { reifiedReturnShapes.resize(1); @@ -441,14 +422,3 @@ LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, return failure(); } - -Value mlir::tensor::createCanonicalRankReducingExtractSliceOp( - OpBuilder& b, Location loc, Value tensor, RankedTensorType targetType) { - auto rankedTensorType = llvm::cast(tensor.getType()); - unsigned rank = rankedTensorType.getRank(); - SmallVector offsets(rank, b.getIndexAttr(0)); - SmallVector sizes = getMixedSizes(b, loc, tensor); - SmallVector strides(rank, b.getIndexAttr(1)); - return b.createOrFold(loc, targetType, tensor, - offsets, sizes, strides); -} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 26d4439e95..c26f5479b2 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -445,16 +445,3 @@ void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, InsertSliceOpCastFolder, InsertSliceOpSourceCastInserter>(context); } - -Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder& b, - Location loc, - Value tensor, - Value dest) { - auto rankedTensorType = llvm::cast(dest.getType()); - unsigned rank = rankedTensorType.getRank(); - SmallVector offsets(rank, b.getIndexAttr(0)); - SmallVector sizes = getMixedSizes(b, loc, dest); - SmallVector strides(rank, b.getIndexAttr(1)); - return b.createOrFold(loc, tensor, dest, offsets, - sizes, strides); -} From a4c9713598db201865adb3ba1b805de9cccc1141 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 5 Mar 2026 10:18:24 +0100 Subject: [PATCH 021/108] add allocTensorOp --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 2 +- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 16 ++++++++++++ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 22 ++++------------ .../QTensor/IR/Operations/AllocTensorOp.cpp | 26 +++++++++++++++++++ mlir/unittests/programs/qco_programs.cpp | 20 +++++++------- 5 files changed, 58 insertions(+), 28 deletions(-) create mode 100644 mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 3d1f3ef0ef..1e43beb101 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -222,7 +222,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> * ``` */ - Value allocateTensor(int64_t size); + Value allocTensor(int64_t size); /** * @brief Allocate a qubit tensor from a list of qubit values diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 87dc6e434f..533e1992ab 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -76,6 +76,22 @@ class QTensorOpWithOffsetSizesAndStrides { + let summary = "Tensor alloc operation"; + let description = [{ + TODO + }]; + + let arguments = (ins ConfinedAttr:$size); + let results = (outs AnyStaticShapeTensor:$result); + let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; + + let builders = [ + OpBuilder<(ins "int64_t":$size)> + ]; +} + def QTensor_ExtractOp : QTensorOp<"extract", [ DeclareOpInterfaceMethods, Pure, diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index fbcb5c31ba..fa5456ead9 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -12,7 +12,6 @@ #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" @@ -166,24 +165,13 @@ void QCOProgramBuilder::updateQubitTracking(Value inputQubit, // QTensor Operations //===----------------------------------------------------------------------===// -Value QCOProgramBuilder::allocateTensor(int64_t size) { +Value QCOProgramBuilder::allocTensor(int64_t size) { checkFinalized(); - - if (size <= 0) { - llvm::reportFatalUsageError("Size must be positive"); - } - - llvm::SmallVector qubits; - qubits.reserve(static_cast(size)); - for (int64_t i = 0; i < size; ++i) { - auto allocOp = AllocOp::create(*this); - qubits.emplace_back(allocOp); - } - - auto fromElementsOp = qtensor::FromElementsOp::create(*this, qubits); - validQubits.insert(fromElementsOp); - return fromElementsOp.getResult(); + auto allocOp = qtensor::AllocOp::create(*this, size); + validQubits.insert(allocOp); + return allocOp.getResult(); } + Value QCOProgramBuilder::fromElements(ValueRange elements) { checkFinalized(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp new file mode 100644 index 0000000000..72998f1db2 --- /dev/null +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -0,0 +1,26 @@ +/* + * 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 + +using namespace mlir; +using namespace mlir::qtensor; + +void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { + auto resultType = + RankedTensorType::get({size}, qco::QubitType::get(builder.getContext())); + build(builder, result, resultType, + IntegerAttr::get(builder.getIntegerType(64), size)); +} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 4e07bf7aa6..d03a8c1140 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2050,33 +2050,33 @@ void nestedFalseIf(QCOProgramBuilder& b) { }); } -void allocTensor(QCOProgramBuilder& b) { b.allocateTensor(3); } +void allocTensor(QCOProgramBuilder& b) { b.allocTensor(3); } void allocDeallocTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); b.deallocTensor(qtensor); } void extractTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); // does not work for now - // b.extract(qtensor, 0); + b.extract(qtensor, 0); } void insertTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); auto [q0, extractOutTensor] = b.extract(qtensor, 0); auto q1 = b.h(q0); b.insert(q1, extractOutTensor, 0); } void extractSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); b.extractSlice(qtensor, 0, 2, 1); } void insertSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); auto q1 = b.h(q0); @@ -2085,19 +2085,19 @@ void insertSliceTensor(QCOProgramBuilder& b) { } void extractInsertTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); auto [q0, extractOutTensor] = b.extract(qtensor, 0); b.insert(q0, extractOutTensor, 0); } void extractSliceInsertSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); b.insertSlice(slicedTensor, extractSliceOutTensor, 0, 2, 1); } void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocateTensor(3); + auto qtensor = b.allocTensor(3); auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); auto insertOutTensor = b.insert(q0, extractOutTensor, 0); From 1f767dee860fe21ceef0fae8cc2f05f9942a6d89 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 5 Mar 2026 10:18:43 +0100 Subject: [PATCH 022/108] fix some headers --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 1 + mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index ebe1a1107f..64b7f9ca76 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -10,6 +10,7 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index fff0868e7b..4758a14ad5 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -10,6 +10,8 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include +#include #include #include #include @@ -21,11 +23,13 @@ #include #include #include +#include #include #include #include #include #include +#include using namespace mlir; using namespace mlir::qtensor; From 6f2a2a57b76f8afb6418451cde2d019c5a6e2790 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 5 Mar 2026 10:39:00 +0100 Subject: [PATCH 023/108] add allocDeallocPair canonicalization --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 2 + .../QTensor/IR/Operations/DeallocTensorOp.cpp | 42 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 533e1992ab..2bd9caab37 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -433,5 +433,7 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { let arguments = (ins AnyRankedTensor:$tensor); let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; + + let hasCanonicalizer = 1; } #endif // TENSOR_OPS diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp index 9650bbe91c..dea23d5aa2 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp @@ -7,3 +7,45 @@ * * Licensed under the MIT License */ + +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#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 if the predecessor is an qtensor::AllocOp + auto* defOp = op.getTensor().getDefiningOp(); + if (!llvm::isa(defOp)) { + return failure(); + } + + // Remove the AllocOp and the DeallocOp + rewriter.eraseOp(op); + rewriter.eraseOp(defOp); + return success(); + } +}; + +} // namespace + +void DeallocOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} From b1f18e69774bfabc5d734b51f14e506473036c27 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 5 Mar 2026 10:53:28 +0100 Subject: [PATCH 024/108] fix linter issues --- .../QTensor/IR/Operations/AllocTensorOp.cpp | 3 ++- .../Dialect/QTensor/IR/Operations/ExtractOp.cpp | 3 ++- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 6 ++++-- .../Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 2 ++ .../QTensor/IR/Operations/InsertSliceOp.cpp | 9 ++++++++- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 14 +++++++++++--- 6 files changed, 29 insertions(+), 8 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index 72998f1db2..7ff27b7952 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -11,10 +11,11 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include +#include + using namespace mlir; using namespace mlir::qtensor; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 64b7f9ca76..386dc1ccf3 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -10,7 +10,6 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -24,6 +23,8 @@ #include #include +#include + using namespace mlir; using namespace mlir::qtensor; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 4758a14ad5..fa4e66e018 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -8,10 +8,9 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include -#include #include #include #include @@ -29,6 +28,9 @@ #include #include #include + +#include +#include #include using namespace mlir; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 0ddbfa7272..c4a5e5e0d3 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -18,6 +18,8 @@ #include #include +#include + using namespace mlir; using namespace mlir::qtensor; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index c26f5479b2..5e916a8995 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include @@ -20,13 +21,19 @@ #include #include #include +#include #include #include #include +#include #include #include #include +#include +#include +#include + using namespace mlir; using namespace mlir::qtensor; @@ -346,7 +353,7 @@ struct InsertSliceOpCastFolder final : public OpRewritePattern { // In the parallel case there is no result and so nothing to cast. bool isParallelInsert = - std::is_same::value; + std::is_same_v; if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { replacement = rewriter.create(insertSliceOp.getLoc(), insertSliceOp.getDestType(), diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index a2414ee269..70f0b1fdad 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -13,24 +13,32 @@ #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated #include +#include #include -#include +#include #include -#include #include #include #include #include +#include #include #include #include +#include #include +#include #include #include #include +#include #include #include +#include +#include +#include + // The following headers are needed for some template instantiations. // IWYU pragma: begin_keep #include @@ -160,7 +168,6 @@ bool foldTensorCastPrecondition(DestinationStyleOpInterface op) { return mlir::tensor::hasFoldableTensorCastOperand(op); } -} // namespace struct FoldTensorCastProducerOp : public OpInterfaceRewritePattern { @@ -200,6 +207,7 @@ struct FoldTensorCastProducerOp return success(); } }; +} // namespace void QTensorDialect::getCanonicalizationPatterns( RewritePatternSet& results) const { From 9dacf342bd6acbdf0b5370b92055c00e8a9cd6aa Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 5 Mar 2026 11:19:20 +0100 Subject: [PATCH 025/108] remove unused code --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 1 - .../QTensor/IR/Operations/ExtractSliceOp.cpp | 33 ------------------- 2 files changed, 34 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 2bd9caab37..551383865f 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -37,7 +37,6 @@ def QTensorDialect : Dialect { let hasCanonicalizer = 1; let hasConstantMaterializer = 1; let dependentDialects = [ - "affine::AffineDialect", "arith::ArithDialect", "complex::ComplexDialect", ]; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index fa4e66e018..7afec06158 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -326,39 +326,6 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern { } }; -/// Slice elements from `values` into `outValues`. `counts` represents the -/// numbers of elements to stride in the original values for each dimension. -/// The output values can be used to construct a DenseElementsAttr. -template -void sliceElements(IterTy values, ArrayRef counts, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides, - llvm::SmallVectorImpl* outValues) { - assert(offsets.size() == sizes.size()); - assert(offsets.size() == strides.size()); - if (offsets.empty()) { - return; - } - - int64_t offset = offsets.front(); - int64_t size = sizes.front(); - int64_t stride = strides.front(); - if (offsets.size() == 1) { - for (int64_t i = 0; i < size; ++i, offset += stride) { - outValues->push_back(*(values + offset)); - } - - return; - } - - for (int64_t i = 0; i < size; ++i, offset += stride) { - auto begin = values + offset * counts.front(); - sliceElements(begin, counts.drop_front(), - offsets.drop_front(), sizes.drop_front(), - strides.drop_front(), outValues); - } -} - } // namespace /// Return the canonical type of the result of an extract_slice op. From c30daaa8c2f477a0b2b9c615e318f4e8828bc62b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 5 Mar 2026 10:20:39 +0000 Subject: [PATCH 026/108] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h | 1 + 1 file changed, 1 insertion(+) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index 75fc520754..ee3bd13a44 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -26,6 +26,7 @@ #include #include + #include #define GET_OP_CLASSES From 511cfc0b5472cdf4293ee287929fa9699ce40d70 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 5 Mar 2026 11:35:49 +0100 Subject: [PATCH 027/108] more linter issue fixes --- mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 1 + mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp | 1 + mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 4 ++++ mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 2 ++ mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp | 1 - mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 2 ++ mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 5 ++--- 7 files changed, 12 insertions(+), 4 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index d1fde19a26..8f0bebb579 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index 7ff27b7952..eef9fedbd4 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -13,6 +13,7 @@ #include #include +#include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 7afec06158..cb020f4707 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -18,13 +18,17 @@ #include #include #include +#include #include #include #include +#include +#include #include #include #include #include +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index c4a5e5e0d3..25403cbef0 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -14,10 +14,12 @@ #include #include #include +#include #include #include #include +#include #include using namespace mlir; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index b9cc5a6aef..6cfa0f983a 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -10,7 +10,6 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 5e916a8995..eb5580c1c8 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -18,11 +18,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 70f0b1fdad..ca47d2ebc5 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -155,9 +155,7 @@ foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, // Common Canonicalizers and Folders. //===----------------------------------------------------------------------===// -namespace { - -bool foldTensorCastPrecondition(DestinationStyleOpInterface op) { +static bool foldTensorCastPrecondition(DestinationStyleOpInterface op) { // 1. InsertSliceOp has its own logic about folding tensor.cast ops. // 2. Exclude DPS ops that are also LoopLike from this interface as they // might need special handling of attached regions. @@ -169,6 +167,7 @@ bool foldTensorCastPrecondition(DestinationStyleOpInterface op) { return mlir::tensor::hasFoldableTensorCastOperand(op); } +namespace { struct FoldTensorCastProducerOp : public OpInterfaceRewritePattern { using OpInterfaceRewritePattern< From 03bc636a1ba74d189d44b94904328bc5971c8e4b Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sat, 7 Mar 2026 15:48:03 +0100 Subject: [PATCH 028/108] update slice operations to newer version --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 27 +++--- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 94 ++++++++++++------- .../QTensor/IR/Operations/InsertSliceOp.cpp | 41 ++++---- 3 files changed, 93 insertions(+), 69 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 551383865f..cb224ae0ed 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -165,7 +165,7 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", let builders = [ // Build an ExtractSliceOp with mixed static and dynamic entries and // inferred result type. - OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, + OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, "ArrayRef":$sizes, "ArrayRef":$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with mixed static and dynamic entries and custom @@ -187,13 +187,18 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, CArg<"ArrayRef", "{}">:$attrs)>, - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with mixed static and dynamic entries packed in // a Range vector. OpBuilder<(ins "Value":$source, "ArrayRef":$ranges, CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with mixed static and dynamic sizes, inferred + // result type, offsets set to 0 and strides set to 1. + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ArrayRef":$sizes, + CArg<"ArrayRef", "{}">:$attrs)>, ]; let extraClassDeclaration = extraBaseClassDeclaration # [{ @@ -213,17 +218,13 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", /// An extract_slice result type can be inferred, when it is not /// rank-reduced, from the source type and the static representation of - /// offsets, sizes and strides. Special sentinels encode the dynamic case. + /// sizes. Special sentinels encode the dynamic case. static RankedTensorType inferResultType( RankedTensorType sourceTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); + ArrayRef staticSizes); static RankedTensorType inferResultType( RankedTensorType sourceTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); + ArrayRef staticSizes); /// If the rank is reduced (i.e. the desiredResultRank is smaller than the /// number of sizes), drop as many size 1 as needed to produce an inferred type @@ -236,15 +237,11 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", static RankedTensorType inferCanonicalRankReducedResultType( unsigned resultRank, RankedTensorType sourceRankedTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); + ArrayRef staticSizes); static RankedTensorType inferCanonicalRankReducedResultType( unsigned resultRank, RankedTensorType sourceRankedTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); + ArrayRef staticSizes); /// Return the expected rank of each of the`static_offsets`, `static_sizes` /// and `static_strides` attributes. diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index cb020f4707..af898d511d 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -45,12 +45,19 @@ void ExtractSliceOp::getAsmResultNames( setNameFn(getResult(), "q_extracted_slice"); } -/// An extract_slice result type can be inferred, when it is not -/// rank-reduced, from the source type and the static representation of -/// offsets, sizes and strides. Special sentinels encode the dynamic case. -RankedTensorType ExtractSliceOp::inferResultType( - RankedTensorType sourceTensorType, ArrayRef /*staticOffsets*/, - ArrayRef staticSizes, ArrayRef /*staticStrides*/) { +/** + * @brief Infers the result type of an extract_slice operation when it is + * not rank-reduced. + * + * @param sourceTensorType The ranked source tensor type. + * @param staticOffsets The static offsets (sentinel values indicate dynamic). + * @param staticSizes The static sizes (sentinel values indicate dynamic). + * @param staticStrides The static strides (sentinel values indicate dynamic). + * @return The inferred RankedTensorType for the resulting slice. + */ +RankedTensorType +ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, + ArrayRef staticSizes) { // An extract_slice op may specify only a leading subset of offset/sizes/ // strides in which case we complete with offset=0, sizes from memref type // and strides=1. @@ -61,11 +68,26 @@ RankedTensorType ExtractSliceOp::inferResultType( sourceTensorType.getEncoding()); } -RankedTensorType ExtractSliceOp::inferResultType( - RankedTensorType sourceTensorType, ArrayRef /*offsets*/, - ArrayRef sizes, ArrayRef /*strides*/) { +/** + * @brief Infers the result type of an extract_slice operation when it is + * not rank-reduced, using mixed static/dynamic offsets, sizes, and strides. + * + * Static sizes are extracted from the `sizes` parameter. Dynamic values + * are represented using sentinels in `OpFoldResult`. The function asserts + * that the number of static sizes matches the rank of the source tensor. + * + * @param sourceTensorType The ranked source tensor type. + * @param offsets The mixed static/dynamic offsets. + * @param sizes The mixed static/dynamic sizes. + * @param strides The mixed static/dynamic strides. + * @return The inferred RankedTensorType for the resulting slice. + */ +RankedTensorType +ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, + ArrayRef sizes) { SmallVector staticSizes; std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes); + assert(static_cast(staticSizes.size()) == sourceTensorType.getRank() && "unexpected staticSizes not equal to rank of source"); @@ -83,11 +105,10 @@ RankedTensorType ExtractSliceOp::inferResultType( /// To disambiguate, this function always drops the first 1 sizes occurrences. RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides) { + ArrayRef sizes) { // Type inferred in the absence of rank-reducing behavior. auto inferredType = llvm::cast( - inferResultType(sourceRankedTensorType, offsets, sizes, strides)); + inferResultType(sourceRankedTensorType, sizes)); int64_t rankDiff = inferredType.getRank() - desiredResultRank; if (rankDiff > 0) { auto shape = inferredType.getShape(); @@ -108,20 +129,12 @@ RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, - ArrayRef offsets, ArrayRef sizes, - ArrayRef strides) { - SmallVector staticOffsets; + ArrayRef sizes) { SmallVector staticSizes; - SmallVector staticStrides; - SmallVector dynamicOffsets; SmallVector dynamicSizes; - SmallVector dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); return ExtractSliceOp::inferCanonicalRankReducedResultType( - desiredResultRank, sourceRankedTensorType, staticOffsets, staticSizes, - staticStrides); + desiredResultRank, sourceRankedTensorType, staticSizes); } /// Build an ExtractSliceOp with mixed static and dynamic entries and custom @@ -145,8 +158,8 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, auto sourceRankedTensorType = llvm::cast(source.getType()); // Structuring implementation this way avoids duplication between builders. if (!resultType) { - resultType = llvm::cast(ExtractSliceOp::inferResultType( - sourceRankedTensorType, staticOffsets, staticSizes, staticStrides)); + resultType = llvm::cast( + ExtractSliceOp::inferResultType(sourceRankedTensorType, staticSizes)); } result.addAttributes(attrs); build(b, result, {resultType, outSourceType}, source, dynamicOffsets, @@ -154,6 +167,7 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, b.getDenseI64ArrayAttr(staticSizes), b.getDenseI64ArrayAttr(staticStrides)); } + void ExtractSliceOp::build(OpBuilder& b, OperationState& result, RankedTensorType resultType, Value source, ArrayRef offsets, @@ -206,8 +220,8 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, RankedTensorType resultType, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef attrs) { - build(b, result, resultType, resultType, source, offsets, sizes, strides, - attrs); + build(b, result, resultType, cast(source.getType()), source, + offsets, sizes, strides, attrs); } /// Build an ExtractSliceOp with dynamic entries and inferred result type. @@ -218,13 +232,26 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, source, offsets, sizes, strides, attrs); } +/// Build an ExtractSliceOp with mixed static and dynamic sizes, inferred +/// result type, offsets set to 0 and strides set to 1. +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, + RankedTensorType resultType, Value source, + ArrayRef sizes, + ArrayRef attrs) { + Attribute zeroIdxAttr = b.getIndexAttr(0); + Attribute oneIdxAttr = b.getIndexAttr(1); + SmallVector readStrides(sizes.size(), oneIdxAttr); + SmallVector readOffsets(sizes.size(), zeroIdxAttr); + build(b, result, resultType, cast(source.getType()), source, + readOffsets, sizes, readStrides, attrs); +} /// Verifier for ExtractSliceOp. LogicalResult ExtractSliceOp::verify() { RankedTensorType sourceType = getSourceType(); // Verify result type against inferred type. - RankedTensorType expectedType = ExtractSliceOp::inferResultType( - sourceType, getMixedOffsets(), getMixedSizes(), getMixedStrides()); + RankedTensorType expectedType = + ExtractSliceOp::inferResultType(sourceType, getMixedSizes()); SliceVerificationResult result = isRankReducedType(expectedType, getType()); if (result != SliceVerificationResult::Success) { return produceSliceErrorMsg(result, *this, expectedType); @@ -275,12 +302,12 @@ namespace { * * Example: * %0 = tensor.cast %V : tensor<3x!qco.qubit> to tensor - * %1, %2 = tensor.extract_slice %0[0][2][1] + * %1, %2 = qtensor.extract_slice %0[0][2][1] * : tensor to tensor<2x!qco.qubit> * * is rewritten into: * - * %0, %1 = tensor.extract_slice %V[0][2][1] + * %0, %1 = qtensor.extract_slice %V[0][2][1] * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> * %2 = tensor.cast %0 : tensor<2x!qco.qubit> to tensor<2x!qco.qubit> * @@ -335,12 +362,11 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern { /// Return the canonical type of the result of an extract_slice op. struct SliceReturnTypeCanonicalizer { RankedTensorType operator()(ExtractSliceOp op, - ArrayRef mixedOffsets, + ArrayRef /*mixedOffsets*/, ArrayRef mixedSizes, - ArrayRef mixedStrides) { + ArrayRef /*mixedStrides*/) { return ExtractSliceOp::inferCanonicalRankReducedResultType( - op.getType().getRank(), op.getSourceType(), mixedOffsets, mixedSizes, - mixedStrides); + op.getType().getRank(), op.getSourceType(), mixedSizes); } }; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index eb5580c1c8..866fbfccdf 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -91,14 +91,16 @@ void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, /// Rank-reducing type verification for both InsertSliceOp and /// ParallelInsertSliceOp. -static SliceVerificationResult verifyInsertSliceOp( - RankedTensorType srcType, RankedTensorType dstType, - ArrayRef staticOffsets, ArrayRef staticSizes, - ArrayRef staticStrides, RankedTensorType* expectedType = nullptr) { +static SliceVerificationResult +verifyInsertSliceOp(RankedTensorType srcType, RankedTensorType dstType, + ArrayRef /*staticOffsets*/, + ArrayRef staticSizes, + ArrayRef /*staticStrides*/, + RankedTensorType* expectedType = nullptr) { // insert_slice is the inverse of extract_slice, use the same type // inference. - RankedTensorType expected = ExtractSliceOp::inferResultType( - dstType, staticOffsets, staticSizes, staticStrides); + RankedTensorType expected = + ExtractSliceOp::inferResultType(dstType, staticSizes); if (expectedType != nullptr) { *expectedType = expected; } @@ -235,18 +237,16 @@ class InsertSliceOpConstantArgumentFolder final // Create the new op in canonical form. auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType( insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(), - mixedOffsets, mixedSizes, mixedStrides); + mixedSizes); Value toInsert = insertSliceOp.getSource(); if (sourceType != insertSliceOp.getSourceType()) { OpBuilder::InsertionGuard g(rewriter); // The only difference between InsertSliceOp and ParallelInsertSliceOp - // is that the insertion point is just before the ParallelCombiningOp in + // is that the insertion point is just before the InParallelOp in // the parallel case. - if (std::is_same_v) { - rewriter.setInsertionPoint(insertSliceOp->getParentOp()); - } - toInsert = rewriter.create(insertSliceOp.getLoc(), - sourceType, toInsert); + + toInsert = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), + sourceType, toInsert); } rewriter.replaceOpWithNewOp( insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets, @@ -349,17 +349,18 @@ struct InsertSliceOpCastFolder final : public OpRewritePattern { return failure(); } - Operation* replacement = rewriter.create( - insertSliceOp.getLoc(), src, dst, insertSliceOp.getMixedOffsets(), - mixedSizes, insertSliceOp.getMixedStrides()); + Operation* replacement = + InsertOpTy::create(rewriter, insertSliceOp.getLoc(), src, dst, + insertSliceOp.getMixedOffsets(), mixedSizes, + insertSliceOp.getMixedStrides()); // In the parallel case there is no result and so nothing to cast. bool isParallelInsert = - std::is_same_v; + std::is_same::value; if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { - replacement = rewriter.create(insertSliceOp.getLoc(), - insertSliceOp.getDestType(), - replacement->getResult(0)); + replacement = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), + insertSliceOp.getDestType(), + replacement->getResult(0)); } rewriter.replaceOp(insertSliceOp, replacement->getResults()); return success(); From a47cee6686c84c82abd9d690fead365a1e04cf63 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sat, 7 Mar 2026 16:58:48 +0100 Subject: [PATCH 029/108] update descriptions and docs to fit project style --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 661 ++++++++++-------- .../QTensor/IR/Operations/ExtractOp.cpp | 12 +- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 63 +- .../QTensor/IR/Operations/InsertOp.cpp | 6 +- .../QTensor/IR/Operations/InsertSliceOp.cpp | 170 +++-- 5 files changed, 491 insertions(+), 421 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index cb224ae0ed..9ebb8d7bb4 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -28,14 +28,21 @@ include "mlir/IR/OpAsmInterface.td" def QTensorDialect : Dialect { let name = "qtensor"; - let cppNamespace = "::mlir::qtensor"; + + let summary = "The QTensor dialect for tensors with linear typing."; let description = [{ - TODO + The Qtensor dialect is an adjusted variant of the standard tensor dialect + of MLIR that supports linear typing for tensor types. The extract operations + of this dialect are modified so that they also return the updated input tensor + as return value. In addition, the alloc/dealloc operation is added to the dialect + to support the bulk allocation and deallocation of tensors with linear types. }]; let hasCanonicalizer = 1; let hasConstantMaterializer = 1; + + let cppNamespace = "::mlir::qtensor"; let dependentDialects = [ "arith::ArithDialect", "complex::ComplexDialect", @@ -77,18 +84,73 @@ class QTensorOpWithOffsetSizesAndStrides { - let summary = "Tensor alloc operation"; - let description = [{ - TODO - }]; + let summary = "Tensor alloc operation"; + let description = [{ + Allocates a qubit tensor with the given size and returns the allocated tensor. + The qubits are initialized to the |0> state. - let arguments = (ins ConfinedAttr:$size); - let results = (outs AnyStaticShapeTensor:$result); - let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; + Example : + ```mlir + %qtensor = qtensor.alloc(3) : tensor<3x!qco.qubit> + ``` + }]; - let builders = [ - OpBuilder<(ins "int64_t":$size)> - ]; + let arguments = (ins ConfinedAttr:$size); + let results = (outs AnyStaticShapeTensor:$result); + let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; + + let builders = [ + OpBuilder<(ins "int64_t":$size)> + ]; +} + +def QTensor_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. + + Example: + ```mlir + qtensor.dealloc %tensor : <3x!qco.qubit> + ``` + }]; + + let arguments = (ins AnyRankedTensor:$tensor); + let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; + + let hasCanonicalizer = 1; +} + +def QTensor_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 = "tensor from elements operation"; + let description = [{ + The `qtensor.from_elements`is a wrapper operation for the `tensor.from_elements` operation + of the tensor dialect. This operation creates a tensor using the operands as its values. + During the canonicalization this operation is converted into a `tensor.from_elements` operation. + + Example: + + ```mlir + %tensor = tensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins Variadic:$elements); + let results = (outs AnyStaticShapeTensor:$result); + let assemblyFormat = "$elements attr-dict `:` type($result)"; + + let builders = [ + // Special case builder for when `elements` has size >=1. + OpBuilder<(ins "ValueRange":$elements)> + ]; + + let hasCanonicalizer = 1; } def QTensor_ExtractOp : QTensorOp<"extract", [ @@ -98,37 +160,29 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ "tensor", "result", "::llvm::cast($_self).getElementType()">, TypesMatchWith< - "returned tensor type matches input tensor", - "tensor", "out_tensor", "$_self">]> { - let summary = "element extraction operation"; - let description = [{ - The `tensor.extract` op reads a ranked tensor and returns one element as - specified by the given indices. The result of the op is a value with the - same type as the elements of the tensor. The arity of indices must match - the rank of the accessed value. All indices should all be of `index` type. - - Example: - - ```mlir - %4 = tensor.extract %t[%1, %2] : tensor<4x4xi32> - %5 = tensor.extract %rt[%1, %2] : tensor - ``` - }]; + "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 element from the input tensor + at the given indices. In addition, it also returns the updated input tensor as a result. + The indices must be of `index` type. + + Example: + ```mlir + %q0, %outTensor = qtensor.extract %tensor[%c0] : tensor<4x!qco.qubit> + ``` + }]; - let arguments = (ins AnyRankedTensor:$tensor, Variadic:$indices); - let results = (outs AnyType:$result, AnyRankedTensor:$out_tensor); - let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)"; + let arguments = (ins AnyRankedTensor:$tensor, Variadic:$indices); + let results = (outs AnyType:$result, AnyRankedTensor:$out_tensor); + let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)"; - let hasFolder = 1; - let hasVerifier = 1; - let hasCanonicalizer = 1; + let hasFolder = 1; + let hasVerifier = 1; + let hasCanonicalizer = 1; } - -//===----------------------------------------------------------------------===// -// ExtractSliceOp -//===----------------------------------------------------------------------===// - def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", [ DeclareOpInterfaceMethods, DeclareOpInterfaceMethods, @@ -138,177 +192,162 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", TypesMatchWith<"returned tensor type matches input tensor", "source", "out_source", "$_self"> ]> { - let summary = "extract slice operation"; - let description = [{ - - }]; - - let arguments = (ins - AnyRankedTensor:$source, - Variadic:$offsets, - Variadic:$sizes, - Variadic:$strides, - DenseI64ArrayAttr:$static_offsets, - DenseI64ArrayAttr:$static_sizes, - DenseI64ArrayAttr:$static_strides - ); - let results = (outs AnyRankedTensor:$result, AnyRankedTensor:$out_source); - - let assemblyFormat = [{ - $source `` - custom($offsets, $static_offsets) - custom($sizes, $static_sizes) - custom($strides, $static_strides) - attr-dict `:` type($source) `to` type($result) - }]; - - let builders = [ - // Build an ExtractSliceOp with mixed static and dynamic entries and - // inferred result type. - OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, - "ArrayRef":$sizes, "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with mixed static and dynamic entries and custom - // result type. If the type passed is nullptr, it is inferred. - OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, - "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with dynamic entries and custom result type. If - // the type passed is nullptr, it is inferred. - OpBuilder<(ins "Value":$source, "ValueRange":$offsets, - "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with dynamic entries and inferred result type. - OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with mixed static and dynamic entries packed in - // a Range vector. - OpBuilder<(ins "Value":$source, "ArrayRef":$ranges, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with mixed static and dynamic sizes, inferred - // result type, offsets set to 0 and strides set to 1. - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ArrayRef":$sizes, - CArg<"ArrayRef", "{}">:$attrs)>, - ]; - - let extraClassDeclaration = extraBaseClassDeclaration # [{ - /// The result of an extract_slice is always a tensor. - // TODO: deprecate - RankedTensorType getType() { - return getResultType(); - } - - /// Compute the rank-reduction mask that can be applied to map the source - /// tensor type to the result tensor type by dropping unit dims. - std::optional> - computeRankReductionMask() { - return ::mlir::computeRankReductionMask(getSourceType().getShape(), - getType().getShape()); - }; - - /// An extract_slice result type can be inferred, when it is not - /// rank-reduced, from the source type and the static representation of - /// sizes. Special sentinels encode the dynamic case. - static RankedTensorType inferResultType( - RankedTensorType sourceTensorType, - ArrayRef staticSizes); - static RankedTensorType inferResultType( - RankedTensorType sourceTensorType, - ArrayRef staticSizes); - - /// If the rank is reduced (i.e. the desiredResultRank is smaller than the - /// number of sizes), drop as many size 1 as needed to produce an inferred type - /// with the desired rank. - /// - /// Note that there may be multiple ways to compute this rank-reduced type: - /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. - /// - /// To disambiguate, this function always drops the first 1 sizes occurrences. - static RankedTensorType inferCanonicalRankReducedResultType( - unsigned resultRank, - RankedTensorType sourceRankedTensorType, - ArrayRef staticSizes); - static RankedTensorType inferCanonicalRankReducedResultType( - unsigned resultRank, - RankedTensorType sourceRankedTensorType, - ArrayRef staticSizes); - - /// Return the expected rank of each of the`static_offsets`, `static_sizes` - /// and `static_strides` attributes. - std::array getArrayAttrMaxRanks() { - unsigned rank = getSourceType().getRank(); - return {rank, rank, rank}; - } - - /// Return the number of leading operands before the `offsets`, `sizes` and - /// and `strides` operands. - static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } - - /// Return the dimensions of the source that are dropped in the - /// result when the result is rank-reduced. - llvm::SmallBitVector getDroppedDims(); - - /// Given a `value`, asserted to be of RankedTensorType, build an - /// ExtractSliceOp that results in a rank-reducing extract to the desired - /// tensor shape and return the new value created. - /// If the shape of `value` is already the `desiredShape`, just return - /// `value`. - /// If the shape of `value` cannot be rank-reduced to `desiredShape`, fail. - static FailureOr rankReduceIfNeeded( - OpBuilder &b, Location loc, Value value, ArrayRef desiredShape); - }]; - - let hasCanonicalizer = 1; - let hasFolder = 1; - let hasVerifier = 1; -} - -def QTensor_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 = "tensor from elements operation."; - let description = [{ - Create a N-D tensor from a range of same-type arguments. The number of - provided `elements` should equal to the number of the elements in the - result type. The `elements` correspond to a flattened tensor. - - Example: - - ```mlir - tensor.from_elements %a, %b, %c, %d, %e, %f : tensor<2x3xindex> - ``` - - will result in a tensor - - [[%a, %b, %c] - [%d, %e, %f]] - }]; - - let arguments = (ins Variadic:$elements); - let results = (outs AnyStaticShapeTensor:$result); + 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 ranked tensor and returns the extracted tensor specified by the + offsets, sizes and strides arguments. In addition, it also returns the updated input tensor as result. + + The extract_slice operation supports the following arguments: + + * source: the "base" tensor from which to extract a slice. + * offsets: tensor-rank number of offsets into the "base" tensor from which + to extract the slice. + * sizes: tensor-rank number of sizes which specify the sizes of the result + tensor type. + * strides: tensor-rank number of strides specifying subsampling in each + dimension. + + The representation based on offsets, sizes and strides support a + partially-static specification via attributes specified through the + `static_offsets`, `static_sizes` and `static_strides` arguments. A special + sentinel value ShapedType::kDynamic encodes that the corresponding entry has + a dynamic value. + + Example: + ```mlir + %extractedSlice, %outSourcetensor = qtensor.extract_slice %sourceTensor[%c0][%c2][%c1] : + : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + ``` + }]; - let assemblyFormat = "$elements attr-dict `:` type($result)"; + let arguments = (ins + AnyRankedTensor:$source, + Variadic:$offsets, + Variadic:$sizes, + Variadic:$strides, + DenseI64ArrayAttr:$static_offsets, + DenseI64ArrayAttr:$static_sizes, + DenseI64ArrayAttr:$static_strides + ); + let results = (outs AnyRankedTensor:$result, AnyRankedTensor:$out_source); + + let assemblyFormat = [{ + $source `` + custom($offsets, $static_offsets) + custom($sizes, $static_sizes) + custom($strides, $static_strides) + attr-dict `:` type($source) `to` type($result) + }]; - let builders = [ - // Special case builder for when `elements` has size >=1. - OpBuilder<(ins "ValueRange":$elements)> - ]; + let builders = [ + // Build an ExtractSliceOp with mixed static and dynamic entries and + // inferred result type. + OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, + "ArrayRef":$sizes, "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with mixed static and dynamic entries and custom + // result type. If the type passed is nullptr, it is inferred. + OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with dynamic entries and custom result type. If + // the type passed is nullptr, it is inferred. + OpBuilder<(ins "Value":$source, "ValueRange":$offsets, + "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with dynamic entries and inferred result type. + OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with mixed static and dynamic entries packed in + // a Range vector. + OpBuilder<(ins "Value":$source, "ArrayRef":$ranges, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an ExtractSliceOp with mixed static and dynamic sizes, inferred + // result type, offsets set to 0 and strides set to 1. + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ArrayRef":$sizes, + CArg<"ArrayRef", "{}">:$attrs)>, + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + /// The result of an extract_slice is always a tensor. + RankedTensorType getType() { + return getResultType(); + } + + /// Compute the rank-reduction mask that can be applied to map the source + /// tensor type to the result tensor type by dropping unit dims. + std::optional> + computeRankReductionMask() { + return ::mlir::computeRankReductionMask(getSourceType().getShape(), + getType().getShape()); + }; + + /// An extract_slice result type can be inferred, when it is not + /// rank-reduced, from the source type and the static representation of + /// sizes. Special sentinels encode the dynamic case. + static RankedTensorType inferResultType( + RankedTensorType sourceTensorType, + ArrayRef staticSizes); + static RankedTensorType inferResultType( + RankedTensorType sourceTensorType, + ArrayRef staticSizes); + + /// If the rank is reduced (i.e. the desiredResultRank is smaller than the + /// number of sizes), drop as many size 1 as needed to produce an inferred type + /// with the desired rank. + /// + /// Note that there may be multiple ways to compute this rank-reduced type: + /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. + /// + /// To disambiguate, this function always drops the first 1 sizes occurrences. + static RankedTensorType inferCanonicalRankReducedResultType( + unsigned resultRank, + RankedTensorType sourceRankedTensorType, + ArrayRef staticSizes); + static RankedTensorType inferCanonicalRankReducedResultType( + unsigned resultRank, + RankedTensorType sourceRankedTensorType, + ArrayRef staticSizes); + + /// Return the expected rank of each of the`static_offsets`, `static_sizes` + /// and `static_strides` attributes. + std::array getArrayAttrMaxRanks() { + unsigned rank = getSourceType().getRank(); + return {rank, rank, rank}; + } + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } + + /// Return the dimensions of the source that are dropped in the + /// result when the result is rank-reduced. + llvm::SmallBitVector getDroppedDims(); + + /// Given a `value`, asserted to be of RankedTensorType, build an + /// ExtractSliceOp that results in a rank-reducing extract to the desired + /// tensor shape and return the new value created. + /// If the shape of `value` is already the `desiredShape`, just return + /// `value`. + /// If the shape of `value` cannot be rank-reduced to `desiredShape`, fail. + static FailureOr rankReduceIfNeeded( + OpBuilder &b, Location loc, Value value, ArrayRef desiredShape); + }]; - let hasCanonicalizer = 1; + let hasCanonicalizer = 1; + let hasFolder = 1; + let hasVerifier = 1; } def QTensor_InsertOp : QTensorOp<"insert", [ @@ -319,19 +358,28 @@ def QTensor_InsertOp : QTensorOp<"insert", [ TypesMatchWith<"scalar type matches element type of dest", "dest", "scalar", "::llvm::cast($_self).getElementType()">]> { - let summary = "element insertion operation"; - let description = [{ - }]; + 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 scalar into a ranked tensor as specified by the operation`s indices. + During the canonicalization this operation is converted into a `tensor.insert` operation. - let arguments = (ins AnyType:$scalar, - AnyRankedTensor:$dest, - Variadic:$indices); - let results = (outs AnyRankedTensor:$result); - let assemblyFormat = [{ - $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) - }]; + Example: - let hasCanonicalizer = 1; + ```mlir + %newTensor = tensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> + ``` + }]; + + let arguments = (ins AnyType:$scalar, + AnyRankedTensor:$dest, + Variadic:$indices); + let results = (outs AnyRankedTensor:$result); + let assemblyFormat = [{ + $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) + }]; + + let hasCanonicalizer = 1; } def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ @@ -344,92 +392,105 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ TypesMatchWith<"expected result type to match dest type", "dest", "result", "$_self"> ]> { - let summary = "insert_slice operation"; - let description = [{ - - }]; - - let arguments = (ins - AnyRankedTensor:$source, - AnyRankedTensor:$dest, - Variadic:$offsets, - Variadic:$sizes, - Variadic:$strides, - DenseI64ArrayAttr:$static_offsets, - DenseI64ArrayAttr:$static_sizes, - DenseI64ArrayAttr:$static_strides - ); - let results = (outs AnyRankedTensor:$result); - - let assemblyFormat = [{ - $source `into` $dest `` - custom($offsets, $static_offsets) - custom($sizes, $static_sizes) - custom($strides, $static_strides) - attr-dict `:` type($source) `into` type($dest) - }]; - - let builders = [ - // Build a InsertSliceOp with mixed static and dynamic entries and inferred - // result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build a InsertSliceOp with dynamic entries and inferred result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an InsertSliceOp with mixed static and dynamic entries packed in - // a Range vector and inferred result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ArrayRef":$ranges, - CArg<"ArrayRef", "{}">:$attrs)> - ]; - - let extraClassDeclaration = extraBaseClassDeclaration # [{ - /// The result of a insert_slice is always a tensor. - // TODO: Deprecate this method. - RankedTensorType getType() { - return getResultType(); - } - - /// The `dest` type is the same as the result type. - RankedTensorType getDestType() { - return getResultType(); - } - - /// Return the expected rank of each of the`static_offsets`, `static_sizes` - /// and `static_strides` attributes. - std::array getArrayAttrMaxRanks() { - unsigned rank = getResultType().getRank(); - return {rank, rank, rank}; - } - - /// Return the dimensions of the dest that are omitted to insert a source - /// when the result is rank-extended. - llvm::SmallBitVector getDroppedDims(); - - /// Return the number of leading operands before the `offsets`, `sizes` and - /// and `strides` operands. - static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; } - - MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } - }]; - - let hasCanonicalizer = 1; - let hasFolder = 1; - let hasVerifier = 1; -} - -def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { - let summary = "Deallocate a tensor"; + let summary = "Insert slice into tensor"; let description = [{ + The `qtensor.insert_slice` operation is a minimally adjusted 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, sizes and strides arguments. + + The insert_slice operation supports the following arguments: + + * source: the tensor that is inserted. + * dest: the tensor into which the source tensor is inserted. + * offsets: tensor-rank number of offsets into the `dest` tensor into which + the slice is inserted. + * sizes: tensor-rank number of sizes which specify the sizes of the source + tensor type. + * strides: tensor-rank number of strides that specify subsampling in each + dimension. + + The representation based on offsets, sizes and strides support a + partially-static specification via attributes specified through the + `static_offsets`, `static_sizes` and `static_strides` arguments. A special + sentinel value ShapedType::kDynamic encodes that the corresponding entry has + a dynamic value. + + Example: + ```mlir + %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] + : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + ``` }]; - let arguments = (ins AnyRankedTensor:$tensor); - let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; + let arguments = (ins + AnyRankedTensor:$source, + AnyRankedTensor:$dest, + Variadic:$offsets, + Variadic:$sizes, + Variadic:$strides, + DenseI64ArrayAttr:$static_offsets, + DenseI64ArrayAttr:$static_sizes, + DenseI64ArrayAttr:$static_strides + ); + let results = (outs AnyRankedTensor:$result); + let assemblyFormat = [{ + $source `into` $dest `` + custom($offsets, $static_offsets) + custom($sizes, $static_sizes) + custom($strides, $static_strides) + attr-dict `:` type($source) `into` type($dest) + }]; + + let builders = [ + // Build a InsertSliceOp with mixed static and dynamic entries and inferred + // result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a InsertSliceOp with dynamic entries and inferred result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build an InsertSliceOp with mixed static and dynamic entries packed in + // a Range vector and inferred result type. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$ranges, + CArg<"ArrayRef", "{}">:$attrs)> + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + /// The result of a insert_slice is always a tensor. + RankedTensorType getType() { + return getResultType(); + } + + /// The `dest` type is the same as the result type. + RankedTensorType getDestType() { + return getResultType(); + } + + /// Return the expected rank of each of the`static_offsets`, `static_sizes` + /// and `static_strides` attributes. + std::array getArrayAttrMaxRanks() { + unsigned rank = getResultType().getRank(); + return {rank, rank, rank}; + } + + /// Return the dimensions of the dest that are omitted to insert a source + /// when the result is rank-extended. + llvm::SmallBitVector getDroppedDims(); + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; } + + MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } + }]; let hasCanonicalizer = 1; + let hasFolder = 1; + let hasVerifier = 1; } + #endif // TENSOR_OPS diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 386dc1ccf3..1a8c7b25d2 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -28,10 +28,6 @@ using namespace mlir; using namespace mlir::qtensor; -//===----------------------------------------------------------------------===// -// ExtractOp -//===----------------------------------------------------------------------===// - void ExtractOp::getAsmResultNames( function_ref setNameFn) { setNameFn(getResult(), "q_extracted"); @@ -63,10 +59,10 @@ struct ExtractFromTensorCast : public OpRewritePattern { } }; -/// If we have an ExtractOp consuming an InsertOp with the same -/// indices, we can return the InsertOp's scalar directly. -// TODO: This only checks the immediate producer; extend to go up the -// insert/extract chain if the slices are disjoint. +/** + * @brief If an ExtractOp consumes an InsertOp with identical indices, + * return the scalar from the InsertOp directly. + */ static Value foldExtractAfterInsert(ExtractOp extractOp) { auto insertOp = extractOp.getTensor().getDefiningOp(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index af898d511d..d99c57865c 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -49,10 +49,14 @@ void ExtractSliceOp::getAsmResultNames( * @brief Infers the result type of an extract_slice operation when it is * not rank-reduced. * + * @details + * The result type can be inferred from the source tensor type and the + * static representation of offsets, sizes, and strides. Special sentinel + * values are used to encode dynamic entries. + * * @param sourceTensorType The ranked source tensor type. - * @param staticOffsets The static offsets (sentinel values indicate dynamic). - * @param staticSizes The static sizes (sentinel values indicate dynamic). - * @param staticStrides The static strides (sentinel values indicate dynamic). + * @param staticSizes The static sizes of the slice (sentinel values + * indicate dynamic sizes). * @return The inferred RankedTensorType for the resulting slice. */ RankedTensorType @@ -68,20 +72,6 @@ ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, sourceTensorType.getEncoding()); } -/** - * @brief Infers the result type of an extract_slice operation when it is - * not rank-reduced, using mixed static/dynamic offsets, sizes, and strides. - * - * Static sizes are extracted from the `sizes` parameter. Dynamic values - * are represented using sentinels in `OpFoldResult`. The function asserts - * that the number of static sizes matches the rank of the source tensor. - * - * @param sourceTensorType The ranked source tensor type. - * @param offsets The mixed static/dynamic offsets. - * @param sizes The mixed static/dynamic sizes. - * @param strides The mixed static/dynamic strides. - * @return The inferred RankedTensorType for the resulting slice. - */ RankedTensorType ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, ArrayRef sizes) { @@ -95,14 +85,19 @@ ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, sourceTensorType.getEncoding()); } -/// If the rank is reduced (i.e. the desiredResultRank is smaller than the -/// number of sizes), drop as many size 1 as needed to produce an inferred -/// type with the desired rank. -/// -/// Note that there may be multiple ways to compute this rank-reduced type: -/// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. -/// -/// To disambiguate, this function always drops the first 1 sizes occurrences. +/** + * @brief Computes the rank-reduced result type. + * + * @details + * If the desired result rank is smaller than the number of slice sizes, + * rank reduction is performed by dropping dimensions of size 1 until the + * desired rank is reached. + * + * Multiple rank-reduced shapes may be possible. For example, a tensor of + * shape 1x6x1 can be reduced to either 1x6 or 6x1. To ensure deterministic + * behavior, this function always drops the first occurrences of size-1 + * dimensions. + */ RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, ArrayRef sizes) { @@ -245,6 +240,7 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, build(b, result, resultType, cast(source.getType()), source, readOffsets, sizes, readStrides, attrs); } + /// Verifier for ExtractSliceOp. LogicalResult ExtractSliceOp::verify() { RankedTensorType sourceType = getSourceType(); @@ -347,10 +343,11 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern { // Create folded extract. Location loc = sliceOp.getLoc(); - auto newResult = rewriter.create( - loc, sliceOp.getType(), castOp.getSource(), sliceOp.getOffsets(), - sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(), - sliceOp.getStaticSizes(), sliceOp.getStaticStrides()); + auto newResult = ExtractSliceOp::create( + rewriter, loc, sliceOp.getType(), castOp.getSource(), + sliceOp.getOffsets(), sliceOp.getSizes(), sliceOp.getStrides(), + sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), + sliceOp.getStaticStrides()); rewriter.replaceOp(sliceOp, newResult->getResult(0)); rewriter.replaceOp(castOp, newResult->getResult(1)); return success(); @@ -377,8 +374,8 @@ struct SliceCanonicalizer { Value replacement = newOp.getResult(); Value outSource = newOp.getOutSource(); if (replacement.getType() != op.getType()) { - replacement = rewriter.create(op.getLoc(), op.getType(), - replacement); + replacement = tensor::CastOp::create(rewriter, op.getLoc(), op.getType(), + replacement); } rewriter.replaceOp(op, {replacement, outSource}); } @@ -392,10 +389,6 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, ExtractSliceOpCastFolder>(context); } -/// If we have an ExtractSliceOp consuming an InsertSliceOp with the same -/// slice, we can return the InsertSliceOp's source directly. -// TODO: This only checks the immediate producer; extend to go up the -// insert/extract chain if the slices are disjoint. static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { auto insertOp = extractOp.getSource().getDefiningOp(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 6cfa0f983a..0171b499bb 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -33,12 +33,16 @@ struct ConvertInsertOpToTensorOp : public OpRewritePattern { } }; +/** + * @brief If an InsertOp consumes an ExtractOp with identical indices, + * return the tensor from the extractOp directly. + */ struct InsertFromExtractOp : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(InsertOp insertOp, PatternRewriter& rewriter) const final { - auto extractOp = insertOp.getScalar().getDefiningOp(); + auto extractOp = insertOp.getScalar().getDefiningOp(); if (!extractOp) { return failure(); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 866fbfccdf..6712963cd3 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -130,23 +130,27 @@ LogicalResult InsertSliceOp::verify() { return success(); } -/// If we have two consecutive InsertSliceOp writing to the same slice, we -/// can mutate the second InsertSliceOp's destination to the first one's. -/// -/// Example: -/// -/// ```mlir -/// %0 = tensor.insert_slice %slice0 into %input[0, 0] [64, 64] [1, 1] -/// %1 = tensor.insert_slice %slice1 into %0[0, 0] [64, 64] [1, 1] -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %1 = tensor.insert_slice %slice1 into %input[0, 0] [64, 64] [1, 1] -/// ``` -/// -/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +/** + * @brief Folds consecutive InsertSliceOp operations writing to the same slice. + * + * @details + * If two consecutive InsertSliceOp operations write to the same slice, + * the destination of the second InsertSliceOp can be updated to the + * destination of the first one, eliminating the intermediate operation. + * + * Example: + * + * ```mlir + * %0 = qtensor.insert_slice %slice0 into %input[0][2][1] + * %1 = qtensor.insert_slice %slice1 into %0[0][2][1] + * ``` + * + * This folds into: + * + * ```mlir + * %1 = qtensor.insert_slice %slice1 into %input[0][2][1] + * ``` + */ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { auto prevInsertOp = insertOp.getDest().getDefiningOp(); @@ -161,13 +165,24 @@ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { return success(); } -/// Folds round-trip extract/insert slice op pairs. -/// Example: -/// ```mlir -/// %0 = tensor.extract_slice %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] -/// %1 = tensor.insert_slice %0 into %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] -/// ``` -/// can be folded into %val. +/** + * @brief Folds round-trip extract/insert slice operation pairs. + * + * @details + * Detects patterns where a slice is extracted from a tensor and then + * inserted back into the same tensor at the same offsets, sizes, and + * strides. In such cases, the pair of operations forms a no-op and can + * be folded to the original tensor value. + * + * Example: + * + * ```mlir + * %0 = qtensor.extract_slice %val[0][2][1] + * %1 = qtensor.insert_slice %0 into %val[0][2][1] + * ``` + * + * This can be folded into `%val`. + */ static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { auto extractOp = insertOp.getSource().getDefiningOp(); auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; @@ -204,9 +219,7 @@ LogicalResult InsertSliceOp::reifyResultShapes( } namespace { -/// Pattern to rewrite a insert_slice op with constant arguments. -/// -/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. + template class InsertSliceOpConstantArgumentFolder final : public OpRewritePattern { @@ -255,26 +268,30 @@ class InsertSliceOpConstantArgumentFolder final } }; -/// Fold tensor_casts with insert_slice operations. If the source or -/// destination tensor is a tensor_cast that removes static type information, -/// the cast is folded into the insert_slice operation. E.g.: -/// -/// ```mlir -/// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor -/// %2 = tensor.insert_slice %1 into ... : tensor into ... -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %2 = tensor.insert_slice %0 into ... : tensor<8x16xf32> into ... -/// ``` -/// -/// Note: When folding a cast on the destination tensor, the result of the -/// insert_slice operation is casted to ensure that the type of the result did -/// not change. -/// -/// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. +/** + * @brief Folds tensor.cast operations with insert_slice. + * + * @details + * If the source or destination tensor of an insert_slice operation is + * produced by a tensor.cast that removes static type information, the + * cast can be folded into the insert_slice operation. + * + * Example: + * + * ```mlir + * %1 = tensor.cast %0 : tensor<3!qco.qubit> to tensor + * %2 = qtensor.insert_slice %1 into ... : tensor into ... + * ``` + * + * This folds into: + * + * ```mlir + * %2 = qtensor.insert_slice %0 into ... : tensor<3!qco.qubit> into ... + * ``` + * + * When folding a cast on the destination tensor, the result of the + * insert_slice operation is cast to preserve the original result type. + */ template struct InsertSliceOpCastFolder final : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -356,7 +373,7 @@ struct InsertSliceOpCastFolder final : public OpRewritePattern { // In the parallel case there is no result and so nothing to cast. bool isParallelInsert = - std::is_same::value; + std::is_same_v; if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { replacement = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), insertSliceOp.getDestType(), @@ -367,27 +384,32 @@ struct InsertSliceOpCastFolder final : public OpRewritePattern { } }; -/// If additional static type information can be deduced from a insert_slice's -/// size operands, insert an explicit cast of the op's source operand. This -/// enables other canonicalization patterns that are matching for tensor_cast -/// ops such as `ForOpTensorCastFolder` in SCF. -/// -/// Example: -/// -/// ```mlir -/// %r = tensor.insert_slice %0 into %1[...] [64, 64] [1, 1] -/// : tensor into ... -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %tmp = tensor.cast %0 : tensor to tensor<64x64xf32> -/// %r = tensor.insert_slice %tmp into %1[...] [64, 64] [1, 1] -/// : tensor<64x64xf32> into ... -/// ``` -/// -/// This patterns works with both InsertSliceOp and ParallelInsertSliceOp. +/** + * @brief Inserts a tensor.cast before insert_slice when additional static + * type information can be inferred from the slice sizes. + * + * @details + * If the size operands of an insert_slice operation provide additional + * static shape information, an explicit tensor.cast is inserted on the + * source operand to refine its type. This enables further canonicalization + * patterns that match tensor.cast operations, such as + * `ForOpTensorCastFolder` in the SCF dialect. + * + * Example: + * + * ```mlir + * %r = qtensor.insert_slice %0 into %1[...] [3] [1] + * : tensor into ... + * ``` + * + * This folds into: + * + * ```mlir + * %tmp = tensor.cast %0 : tensor to tensor<3x!qco.qubit> + * %r = qtensor.insert_slice %tmp into %1[...] [3] [1] + * : tensor<3x!qco.qubit> into ... + * ``` + */ template struct InsertSliceOpSourceCastInserter final : public OpRewritePattern { @@ -428,14 +450,8 @@ struct InsertSliceOpSourceCastInserter final // 3) Cast-compatible with srcType. // Insert the cast. OpBuilder::InsertionGuard g(rewriter); - // The only difference between InsertSliceOp and ParallelInsertSliceOp is - // that the insertion point is just before the ParallelCombiningOp in the - // parallel case. - if (std::is_same_v) { - rewriter.setInsertionPoint(insertSliceOp->getParentOp()); - } - Value cast = rewriter.create( - insertSliceOp.getLoc(), newSrcType, insertSliceOp.getSource()); + Value cast = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), + newSrcType, insertSliceOp.getSource()); rewriter.replaceOpWithNewOp( insertSliceOp, cast, insertSliceOp.getDest(), insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(), From 6f8dcb6d9a8e28cbbc980b1d324869067f537a27 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sat, 7 Mar 2026 17:04:48 +0100 Subject: [PATCH 030/108] remove useless functions --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 3 --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 5 ----- mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 5 ----- mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 5 ----- 4 files changed, 18 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 9ebb8d7bb4..01eaca3cdb 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -154,7 +154,6 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ } def QTensor_ExtractOp : QTensorOp<"extract", [ - DeclareOpInterfaceMethods, Pure, TypesMatchWith<"result type matches element type of tensor", "tensor", "result", @@ -184,7 +183,6 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ } def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", [ - DeclareOpInterfaceMethods, DeclareOpInterfaceMethods, AttrSizedOperandSegments, Pure, @@ -383,7 +381,6 @@ def QTensor_InsertOp : QTensorOp<"insert", [ } def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ - DeclareOpInterfaceMethods, DeclareOpInterfaceMethods, AttrSizedOperandSegments, DestinationStyleOpInterface, diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 1a8c7b25d2..8d5b7a5bdd 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -28,11 +28,6 @@ using namespace mlir; using namespace mlir::qtensor; -void ExtractOp::getAsmResultNames( - function_ref setNameFn) { - setNameFn(getResult(), "q_extracted"); -} - LogicalResult ExtractOp::verify() { // Verify the # indices match if we have a ranked type. auto tensorType = llvm::cast(getTensor().getType()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index d99c57865c..8027a777ac 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -40,11 +40,6 @@ using namespace mlir; using namespace mlir::qtensor; -void ExtractSliceOp::getAsmResultNames( - function_ref setNameFn) { - setNameFn(getResult(), "q_extracted_slice"); -} - /** * @brief Infers the result type of an extract_slice operation when it is * not rank-reduced. diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 6712963cd3..3430f44f77 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -39,11 +39,6 @@ using namespace mlir; using namespace mlir::qtensor; -void InsertSliceOp::getAsmResultNames( - function_ref setNameFn) { - setNameFn(getResult(), "inserted_slice"); -} - // Build a InsertSliceOp with mixed static and dynamic entries. void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value dest, ArrayRef offsets, From 0750cc881ff5fc49a4a119851a7ae2dd3e2a6ad7 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 09:45:48 +0100 Subject: [PATCH 031/108] add fromTensor test --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 2 ++ mlir/unittests/programs/qco_programs.cpp | 8 +++++++- mlir/unittests/programs/qco_programs.h | 3 +++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index bdf9dfdb28..5e6d7279d2 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1034,6 +1034,8 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(allocTensor)}, QCOTestCase{"AllocDeallocTensor", MQT_NAMED_BUILDER(allocDeallocTensor), MQT_NAMED_BUILDER(allocTensor)}, + QCOTestCase{"FromElements", MQT_NAMED_BUILDER(fromElements), + MQT_NAMED_BUILDER(fromElements)}, QCOTestCase{"ExtractTensor", MQT_NAMED_BUILDER(extractTensor), MQT_NAMED_BUILDER(extractTensor)}, QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(insertTensor), diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index d2275689d1..e9970191c9 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2058,9 +2058,15 @@ void allocDeallocTensor(QCOProgramBuilder& b) { b.deallocTensor(qtensor); } +void fromElements(QCOProgramBuilder& b) { + auto q0 = b.allocQubit(); + auto q1 = b.allocQubit(); + auto q2 = b.allocQubit(); + b.fromElements({q0, q1, q2}); +} + void extractTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); - // does not work for now b.extract(qtensor, 0); } diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index d2a5c76ffd..aa54baeae9 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -939,6 +939,9 @@ void allocTensor(QCOProgramBuilder& b); /// Allocates and explicitly deallocates a tensor. void allocDeallocTensor(QCOProgramBuilder& b); +/// Constructs a tensor with from_elements. +void fromElements(QCOProgramBuilder& b); + /// Extracts a qubit from a tensor. void extractTensor(QCOProgramBuilder& b); From 5865cc2b931f246215173088c9f183c317333979 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 09:56:31 +0100 Subject: [PATCH 032/108] fix linter issue --- mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 3430f44f77..4a54931127 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -188,7 +188,7 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { return extractOp.getSource(); } -OpFoldResult InsertSliceOp::fold(FoldAdaptor) { +OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { if (getSourceType().hasStaticShape() && getType().hasStaticShape() && getSourceType() == getType() && succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { From 092024cc7e307b4c971d9a575adabea279816eca Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 10:23:02 +0100 Subject: [PATCH 033/108] smaller fixes --- .../include/mlir/Dialect/QTensor/IR/QTensorOps.td | 4 ++-- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 15 --------------- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 01eaca3cdb..0efb120f91 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -137,7 +137,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ Example: ```mlir - %tensor = tensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> ``` }]; @@ -365,7 +365,7 @@ def QTensor_InsertOp : QTensorOp<"insert", [ Example: ```mlir - %newTensor = tensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> + %newTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> ``` }]; diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index ca47d2ebc5..d2ad168ec7 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -48,21 +48,6 @@ using namespace mlir; using namespace mlir::qtensor; -/// Materialize a single constant operation from a given attribute value with -/// the desired resultant type. -Operation* QTensorDialect::materializeConstant(OpBuilder& builder, - Attribute value, Type type, - Location loc) { - if (auto op = arith::ConstantOp::materialize(builder, value, type, loc)) { - return op; - } - if (complex::ConstantOp::isBuildableWith(value, type)) { - return builder.create(loc, type, - llvm::cast(value)); - } - return nullptr; -} - namespace mlir::qtensor { llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, From 496425a446c0fd5f5230cf6b771bc73ce6c18df4 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 10:23:27 +0100 Subject: [PATCH 034/108] add source of copied parts --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 3 +++ mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 3 +++ mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 3 +++ mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 3 +++ 4 files changed, 12 insertions(+) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 8d5b7a5bdd..0abfcd9ee5 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -28,6 +28,9 @@ using namespace mlir; using namespace mlir::qtensor; +// Adjusted from +// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp + LogicalResult ExtractOp::verify() { // Verify the # indices match if we have a ranked type. auto tensorType = llvm::cast(getTensor().getType()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 8027a777ac..9e82ad330e 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -40,6 +40,9 @@ using namespace mlir; using namespace mlir::qtensor; +// Adjusted from +// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp + /** * @brief Infers the result type of an extract_slice operation when it is * not rank-reduced. diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 25403cbef0..ee8950e021 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -25,6 +25,9 @@ using namespace mlir; using namespace mlir::qtensor; +// Adjusted from +// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp + void FromElementsOp::build(OpBuilder& builder, OperationState& result, ValueRange elements) { assert(!elements.empty() && "expected at least one element"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 4a54931127..c0d680ab39 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -39,6 +39,9 @@ using namespace mlir; using namespace mlir::qtensor; +// Adjusted from +// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp + // Build a InsertSliceOp with mixed static and dynamic entries. void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value dest, ArrayRef offsets, From aca212d3721c339425d0304ad9d7e774f05284c8 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 10:40:51 +0100 Subject: [PATCH 035/108] adjust tablegen file --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 0efb120f91..585fa50791 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -35,12 +35,11 @@ def QTensorDialect : Dialect { The Qtensor dialect is an adjusted variant of the standard tensor dialect of MLIR that supports linear typing for tensor types. The extract operations of this dialect are modified so that they also return the updated input tensor - as return value. In addition, the alloc/dealloc operation is added to the dialect + as return value. In addition, the alloc/dealloc operations are added to the dialect to support the bulk allocation and deallocation of tensors with linear types. }]; let hasCanonicalizer = 1; - let hasConstantMaterializer = 1; let cppNamespace = "::mlir::qtensor"; let dependentDialects = [ From aa91580590ce9ed8643d11719e44ae9f9369a155 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 10:41:34 +0100 Subject: [PATCH 036/108] add additional test for dynamic indices --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 30 +++++++++++++------ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 5 ++++ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 5 ++++ mlir/unittests/programs/qco_programs.cpp | 21 +++++++++++++ mlir/unittests/programs/qco_programs.h | 6 ++++ 5 files changed, 58 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 1f34134a1b..263e039d3f 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -98,6 +98,21 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { */ Value intConstant(int64_t value); + /** + * @brief Create a constant index value + * @param value The value to store in the constant + * @return The value produced by the constant operation + * + * @par Example: + * ```c++ + * auto c = builder.indexConstant(1); + * ``` + * ```mlir + * %c = arith.constant 1 : index + * ``` + */ + Value indexConstant(int64_t value); + //===--------------------------------------------------------------------===// // Memory Management //===--------------------------------------------------------------------===// @@ -214,13 +229,10 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto tensor = builder.allocateTensor(3); + * auto tensor = builder.allocTensor(3); * ``` * ```mlir - * %q0 = qco.alloc : !qco.qubit - * %q1 = qco.alloc : !qco.qubit - * %q2 = qco.alloc : !qco.qubit - * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + * %tensor = qtensor.alloc : tensor<3x!qco.qubit> * ``` */ Value allocTensor(int64_t size); @@ -280,7 +292,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { const std::variant& strides); /** - * @brief insert a qubit into a tensor + * @brief Insert a qubit into a tensor * @param scalar The scalar qubit that is inserted * @param tensor The tensor where the qubit is inserted * @param index The index into where the qubit is inserted @@ -298,12 +310,12 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { const std::variant& index); /** - * @brief insert a qubit slice into a tensor + * @brief Insert a qubit slice into a tensor * @param scalar The slice that is inserted * @param tensor The tensor where the slice is inserted * @param offset The offset into where the slice is inserted * @param size The size of the inserted slice - * @param strides The strides into where the values are inserted + * @param strides The strides of where the qubits are inserted * @return The output tensor * * @par Example: @@ -311,7 +323,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * auto outTensor = builder.insert_slice(slicedTensor, tensor, 0, 2, 1); * ``` * ```mlir - * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] + * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%0][%c2][%c1] * : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> * ``` */ diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 8f0bebb579..9ed5654acb 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -70,6 +70,11 @@ Value QCOProgramBuilder::intConstant(const int64_t value) { return arith::ConstantOp::create(*this, getI64IntegerAttr(value)).getResult(); } +Value QCOProgramBuilder::indexConstant(const int64_t value) { + checkFinalized(); + return arith::ConstantOp::create(*this, getIndexAttr(value)).getResult(); +} + Value QCOProgramBuilder::allocQubit() { checkFinalized(); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 5e6d7279d2..1d8881c7a7 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1040,10 +1040,15 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(extractTensor)}, QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(insertTensor), MQT_NAMED_BUILDER(insertTensor)}, + QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(dynamicInsertTensor), + MQT_NAMED_BUILDER(dynamicInsertTensor)}, QCOTestCase{"ExtractSliceTensor", MQT_NAMED_BUILDER(extractSliceTensor), MQT_NAMED_BUILDER(extractSliceTensor)}, QCOTestCase{"InsertSliceTensor", MQT_NAMED_BUILDER(insertSliceTensor), MQT_NAMED_BUILDER(insertSliceTensor)}, + QCOTestCase{"InsertSliceTensor", + MQT_NAMED_BUILDER(dynamicInsertSliceTensor), + MQT_NAMED_BUILDER(dynamicInsertSliceTensor)}, QCOTestCase{"ExtractInsert", MQT_NAMED_BUILDER(extractInsertTensor), MQT_NAMED_BUILDER(allocTensor)}, QCOTestCase{"ExtractSliceInsertSlice", diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index e9970191c9..a8b470bd80 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2077,6 +2077,14 @@ void insertTensor(QCOProgramBuilder& b) { b.insert(q1, extractOutTensor, 0); } +void dynamicInsertTensor(QCOProgramBuilder& b) { + auto c0 = b.indexConstant(0); + auto qtensor = b.allocTensor(3); + auto [q0, extractOutTensor] = b.extract(qtensor, c0); + auto q1 = b.h(q0); + b.insert(q1, extractOutTensor, c0); +} + void extractSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); b.extractSlice(qtensor, 0, 2, 1); @@ -2091,6 +2099,19 @@ void insertSliceTensor(QCOProgramBuilder& b) { b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2, 1); } +void dynamicInsertSliceTensor(QCOProgramBuilder& b) { + auto c0 = b.indexConstant(0); + auto c1 = b.indexConstant(1); + auto c2 = b.indexConstant(2); + auto qtensor = b.allocTensor(3); + auto [slicedTensor, extractSliceOutTensor] = + b.extractSlice(qtensor, c0, c2, c1); + auto [q0, extractOutTensor] = b.extract(slicedTensor, c0); + auto q1 = b.h(q0); + auto insertOutTensor = b.insert(q1, extractOutTensor, c0); + b.insertSlice(insertOutTensor, extractSliceOutTensor, c0, c2, c1); +} + void extractInsertTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); auto [q0, extractOutTensor] = b.extract(qtensor, 0); diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index aa54baeae9..c0f8a09e04 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -948,12 +948,18 @@ void extractTensor(QCOProgramBuilder& b); /// Inserts a qubit in a tensor. void insertTensor(QCOProgramBuilder& b); +/// Inserts a qubit in a tensor with dynamic index. +void dynamicInsertTensor(QCOProgramBuilder& b); + /// Extracts a slice from a tensor. void extractSliceTensor(QCOProgramBuilder& b); /// Inserts a slice from a tensor. void insertSliceTensor(QCOProgramBuilder& b); +/// Inserts a slice from a tensor with dynamic indices. +void dynamicInsertSliceTensor(QCOProgramBuilder& b); + /// Extracts a qubit from a tensor and insert it immediately. void extractInsertTensor(QCOProgramBuilder& b); From edd2151ea6c5097148416c935ec363cd2c5cca8a Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 10:54:37 +0100 Subject: [PATCH 037/108] remove redundant dialect includes --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 4 ---- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 3 --- 2 files changed, 7 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 585fa50791..dd394eeffe 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -42,10 +42,6 @@ def QTensorDialect : Dialect { let hasCanonicalizer = 1; let cppNamespace = "::mlir::qtensor"; - let dependentDialects = [ - "arith::ArithDialect", - "complex::ComplexDialect", - ]; } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index d2ad168ec7..712750b8e4 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -16,8 +16,6 @@ #include #include #include -#include -#include #include #include #include @@ -25,7 +23,6 @@ #include #include #include -#include #include #include #include From 8147bb1909ac200b6eceed7f2be9b9b4303c11cf Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 12:52:23 +0100 Subject: [PATCH 038/108] address code rabbit suggestions --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 6 +-- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 5 ++- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 2 +- .../QTensor/IR/Operations/DeallocTensorOp.cpp | 15 +++++-- .../QTensor/IR/Operations/ExtractOp.cpp | 16 +++----- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 41 +++++++------------ .../QTensor/IR/Operations/FromTensorOp.cpp | 13 +++++- .../QTensor/IR/Operations/InsertOp.cpp | 12 ++++++ .../QTensor/IR/Operations/InsertSliceOp.cpp | 21 +++++----- 9 files changed, 74 insertions(+), 57 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 263e039d3f..b7e531bfca 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -279,7 +279,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto [extractedSlice, outTensor] = builder.extract_slice(tensor, 0, 2, 1); + * auto [extractedSlice, outTensor] = builder.extractSlice(tensor, 0, 2, 1); * ``` * ```mlir * %extractedSlice, %outTensor = qtensor.extract_slice %tensor[%c0][%c2][%c1] @@ -311,8 +311,8 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Insert a qubit slice into a tensor - * @param scalar The slice that is inserted - * @param tensor The tensor where the slice is inserted + * @param sourceTensor The slice that is inserted + * @param destTensor The tensor where the slice is inserted * @param offset The offset into where the slice is inserted * @param size The size of the inserted slice * @param strides The strides of where the qubits are inserted diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index dd394eeffe..8e8930934a 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -113,6 +113,7 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { let arguments = (ins AnyRankedTensor:$tensor); let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; + let hasVerifier = 1; let hasCanonicalizer = 1; } @@ -145,6 +146,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ OpBuilder<(ins "ValueRange":$elements)> ]; + let hasVerifier = 1; let hasCanonicalizer = 1; } @@ -178,7 +180,6 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ } def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", [ - DeclareOpInterfaceMethods, AttrSizedOperandSegments, Pure, OffsetSizeAndStrideOpInterface, @@ -372,11 +373,11 @@ def QTensor_InsertOp : QTensorOp<"insert", [ $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) }]; + let hasVerifier = 1; let hasCanonicalizer = 1; } def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ - DeclareOpInterfaceMethods, AttrSizedOperandSegments, DestinationStyleOpInterface, Pure, diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 9ed5654acb..aa04ee5bab 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -287,7 +287,7 @@ Value QCOProgramBuilder::insertSlice( llvm::reportFatalUsageError("Source elements must be of QubitType!"); } - auto destTensorType = llvm::dyn_cast(source.getType()); + auto destTensorType = llvm::dyn_cast(dest.getType()); if (!destTensorType) { llvm::reportFatalUsageError("Dest must be of RankedTensorType!"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp index dea23d5aa2..17c0775b4b 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include @@ -19,6 +20,13 @@ using namespace mlir; using namespace mlir::qtensor; +LogicalResult DeallocOp::verify() { + if (!llvm::isa(getTensor().getType().getElementType())) { + return emitOpError("Elements of tensor must be of qubit type"); + } + return success(); +} + namespace { /** @@ -31,14 +39,15 @@ struct RemoveAllocDeallocPair final : OpRewritePattern { LogicalResult matchAndRewrite(DeallocOp op, PatternRewriter& rewriter) const override { // Check if the predecessor is an qtensor::AllocOp - auto* defOp = op.getTensor().getDefiningOp(); - if (!llvm::isa(defOp)) { + auto tensor = op.getTensor(); + auto allocOp = tensor.getDefiningOp(); + if (!allocOp || !tensor.hasOneUse()) { return failure(); } // Remove the AllocOp and the DeallocOp rewriter.eraseOp(op); - rewriter.eraseOp(defOp); + rewriter.eraseOp(allocOp); return success(); } }; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 0abfcd9ee5..d60d9a5fbc 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include @@ -34,11 +35,15 @@ using namespace mlir::qtensor; LogicalResult ExtractOp::verify() { // Verify the # indices match if we have a ranked type. auto tensorType = llvm::cast(getTensor().getType()); + if (!llvm::isa(tensorType.getElementType())) { + return emitOpError("Elements of tensor must be of qubit type"); + } if (tensorType.getRank() != static_cast(getIndices().size())) { return emitOpError("incorrect number of indices for extract_element"); } return success(); } + struct ExtractFromTensorCast : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -74,17 +79,8 @@ static Value foldExtractAfterInsert(ExtractOp extractOp) { return {}; } -LogicalResult ExtractOp::fold(FoldAdaptor adaptor, +LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { - // Collect the constant indices into the tensor. - SmallVector indices; - for (Attribute indice : adaptor.getIndices()) { - if (!indice || !llvm::isa(indice)) { - return failure(); - } - indices.push_back(llvm::cast(indice).getInt()); - } - if (Value result = foldExtractAfterInsert(*this)) { results.push_back(result); results.push_back(getTensor()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 9e82ad330e..f1d47b9307 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" @@ -198,8 +199,7 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, RankedTensorType resultType, RankedTensorType outSourceType, Value source, ValueRange offsets, ValueRange sizes, - ValueRange strides, - ArrayRef /*attrs*/) { + ValueRange strides, ArrayRef attrs) { SmallVector offsetValues = llvm::to_vector<4>( llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); SmallVector sizeValues = llvm::to_vector<4>( @@ -207,7 +207,7 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, SmallVector strideValues = llvm::to_vector<4>( llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); build(b, result, resultType, outSourceType, source, offsetValues, sizeValues, - strideValues); + strideValues, attrs); } void ExtractSliceOp::build(OpBuilder& b, OperationState& result, RankedTensorType resultType, Value source, @@ -243,6 +243,9 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, LogicalResult ExtractSliceOp::verify() { RankedTensorType sourceType = getSourceType(); + if (!llvm::isa(sourceType.getElementType())) { + return emitOpError("Elements of tensor must be of qubit type"); + } // Verify result type against inferred type. RankedTensorType expectedType = ExtractSliceOp::inferResultType(sourceType, getMixedSizes()); @@ -267,21 +270,6 @@ llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { return ::getDroppedDims(getType().getShape(), getMixedSizes()); } -LogicalResult ExtractSliceOp::reifyResultShapes( - OpBuilder& /*builder*/, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { - reifiedReturnShapes.resize(1); - reifiedReturnShapes[0].reserve(getType().getRank()); - SmallVector mixedSizes = getMixedSizes(); - llvm::SmallBitVector droppedDims = getDroppedDims(); - for (const auto& size : enumerate(mixedSizes)) { - if (droppedDims.test(size.index())) { - continue; - } - reifiedReturnShapes[0].push_back(size.value()); - } - return success(); -} - namespace { /** * @brief Rewrite pattern that pushes tensor.cast past tensor.extract_slice. @@ -346,8 +334,15 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern { sliceOp.getOffsets(), sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), sliceOp.getStaticStrides()); - rewriter.replaceOp(sliceOp, newResult->getResult(0)); - rewriter.replaceOp(castOp, newResult->getResult(1)); + Value newOutSource = newResult->getResult(1); + if (newOutSource.getType() != sliceOp.getOutSource().getType()) { + newOutSource = tensor::CastOp::create( + rewriter, loc, sliceOp.getOutSource().getType(), newOutSource); + } + rewriter.replaceOp(sliceOp, {newResult->getResult(0), newOutSource}); + if (castOp->use_empty()) { + rewriter.eraseOp(castOp); + } return success(); } }; @@ -402,12 +397,6 @@ static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { - if (getSourceType() == getType() && - succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { - results.push_back(this->getSource()); - results.push_back(getSource()); - return success(); - } if (Value slice = foldExtractAfterInsertSlice(*this)) { results.push_back(slice); results.push_back(getSource()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index ee8950e021..0620d88fe3 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include @@ -36,6 +37,15 @@ void FromElementsOp::build(OpBuilder& builder, OperationState& result, build(builder, result, resultType, elements); } +LogicalResult FromElementsOp::verify() { + for (auto type : getElements().getTypes()) { + if (!llvm::isa(type)) { + return emitOpError("Elements of ValueRange must be of qubit type"); + } + } + return success(); +} + namespace { struct ConvertFromElementsOpToTensorOp @@ -44,9 +54,8 @@ struct ConvertFromElementsOpToTensorOp LogicalResult matchAndRewrite(qtensor::FromElementsOp fromElementsOp, PatternRewriter& rewriter) const final { - rewriter.replaceOpWithNewOp( - fromElementsOp, fromElementsOp.getElements()); + fromElementsOp, fromElementsOp.getType(), fromElementsOp.getElements()); return success(); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 0171b499bb..b60bfc9d92 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include @@ -19,6 +20,16 @@ using namespace mlir; using namespace mlir::qtensor; +LogicalResult InsertOp::verify() { + if (!llvm::isa(getScalar().getType())) { + return emitOpError("Scalar must be of qubit type"); + } + if (!llvm::isa(getDest().getType().getElementType())) { + return emitOpError("Elements of dest tensor must be of qubit type"); + } + return success(); +} + namespace { struct ConvertInsertOpToTensorOp : public OpRewritePattern { @@ -46,6 +57,7 @@ struct InsertFromExtractOp : public OpRewritePattern { if (!extractOp) { return failure(); } + if (insertOp.getDest() != extractOp.getOutTensor()) { return failure(); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index c0d680ab39..a7b7e42274 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" @@ -76,15 +77,14 @@ void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, // Build a InsertSliceOp with dynamic entries. void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value dest, ValueRange offsets, ValueRange sizes, - ValueRange strides, - ArrayRef /*attrs*/) { + ValueRange strides, ArrayRef attrs) { SmallVector offsetValues = llvm::to_vector<4>( llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); SmallVector sizeValues = llvm::to_vector<4>( llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); SmallVector strideValues = llvm::to_vector<4>( llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, source, dest, offsetValues, sizeValues, strideValues); + build(b, result, source, dest, offsetValues, sizeValues, strideValues, attrs); } /// Rank-reducing type verification for both InsertSliceOp and @@ -109,6 +109,14 @@ verifyInsertSliceOp(RankedTensorType srcType, RankedTensorType dstType, LogicalResult InsertSliceOp::verify() { // Verify result type against inferred type. RankedTensorType expectedType; + + if (!llvm::isa(getSourceType().getElementType())) { + return emitOpError("Elements of source tensor must be of qubit type"); + } + if (!llvm::isa(getDestType().getElementType())) { + return emitOpError("Elements of dest tensor must be of qubit type"); + } + SliceVerificationResult result = verifyInsertSliceOp(getSourceType(), getType(), getStaticOffsets(), getStaticSizes(), getStaticStrides(), &expectedType); @@ -209,13 +217,6 @@ OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { return {}; } -LogicalResult InsertSliceOp::reifyResultShapes( - OpBuilder& builder, ReifiedRankedShapedTypeDims& reifiedReturnShapes) { - reifiedReturnShapes.resize(1, SmallVector(getType().getRank())); - reifiedReturnShapes[0] = tensor::getMixedSizes(builder, getLoc(), getDest()); - return success(); -} - namespace { template From 509205d5457ad5ef7b620d296658a36efdfcf1a3 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 13:16:13 +0100 Subject: [PATCH 039/108] fix linter issues --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 1 - mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 1 + mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp | 1 + mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 1 - 4 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index f1d47b9307..2797457f4d 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -29,7 +29,6 @@ #include #include #include -#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 0620d88fe3..3de0e48efb 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -11,6 +11,7 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index b60bfc9d92..6278f369e1 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -11,6 +11,7 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index a7b7e42274..7d57024142 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -28,7 +28,6 @@ #include #include #include -#include #include #include #include From 45a457a9ad394705db7f4ad16ae651bf2f6dac19 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 13:49:56 +0100 Subject: [PATCH 040/108] address more coderabbit suggestions --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 2 +- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 2 ++ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 9 +++++++ .../QTensor/IR/Operations/AllocTensorOp.cpp | 27 +++++++++++++++++++ .../QTensor/IR/Operations/FromTensorOp.cpp | 4 +++ .../QTensor/IR/Operations/InsertOp.cpp | 6 ++++- 6 files changed, 48 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index b7e531bfca..32c4812106 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -320,7 +320,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto outTensor = builder.insert_slice(slicedTensor, tensor, 0, 2, 1); + * auto outTensor = builder.insertSlice(slicedTensor, tensor, 0, 2, 1); * ``` * ```mlir * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%0][%c2][%c1] diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 8e8930934a..e321f9ba65 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -97,6 +97,8 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ let builders = [ OpBuilder<(ins "int64_t":$size)> ]; + + let hasVerifier = 1; } def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index aa04ee5bab..d272d0e309 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -174,6 +174,11 @@ void QCOProgramBuilder::updateQubitTracking(Value inputQubit, Value QCOProgramBuilder::allocTensor(int64_t size) { checkFinalized(); + + if (size <= 0) { + llvm::reportFatalUsageError("Size must be positive"); + } + auto allocOp = qtensor::AllocOp::create(*this, size); validQubits.insert(allocOp); return allocOp.getResult(); @@ -186,6 +191,7 @@ Value QCOProgramBuilder::fromElements(ValueRange elements) { if (!llvm::isa(element.getType())) { llvm::reportFatalUsageError("Elements must be QubitType!"); } + validateQubitValue(element); validQubits.erase(element); } @@ -267,6 +273,7 @@ Value QCOProgramBuilder::insert(Value scalar, Value tensor, auto outTensor = insertOp.getResult(); + validateQubitValue(scalar); validQubits.erase(scalar); updateQubitTracking(tensor, outTensor); return outTensor; @@ -304,6 +311,7 @@ Value QCOProgramBuilder::insertSlice( auto outTensor = insertSliceOp.getResult(); + validateQubitValue(source); validQubits.erase(source); updateQubitTracking(dest, outTensor); @@ -324,6 +332,7 @@ QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { qtensor::DeallocOp::create(*this, tensor); + validateQubitValue(tensor); validQubits.erase(tensor); return *this; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index eef9fedbd4..3b8c16d746 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -11,9 +11,12 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include #include +#include #include +#include #include @@ -26,3 +29,27 @@ void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { build(builder, result, resultType, IntegerAttr::get(builder.getIntegerType(64), size)); } + +LogicalResult AllocOp::verify() { + auto resultType = dyn_cast(getResult().getType()); + if (!resultType) { + return emitOpError("Result must be a ranked tensor"); + } + + if (!llvm::isa(resultType.getElementType())) { + return emitOpError("Result element type must be of qubit type"); + } + + auto size = static_cast(getSize()); + + if (resultType.getRank() != 1) { + return emitOpError("Result must be a 1-D tensor"); + } + + if (resultType.getShape()[0] != size) { + return emitOpError("Tensor length must match size operand (") + << size << "), but got " << resultType.getShape()[0]; + } + + return success(); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 3de0e48efb..70cbb88ddd 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -39,6 +39,10 @@ void FromElementsOp::build(OpBuilder& builder, OperationState& result, } LogicalResult FromElementsOp::verify() { + if (!llvm::isa(getResult().getType().getElementType())) { + return emitOpError("result tensor must have qubit element type"); + } + for (auto type : getElements().getTypes()) { if (!llvm::isa(type)) { return emitOpError("Elements of ValueRange must be of qubit type"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 6278f369e1..1ce9a43748 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -22,12 +22,16 @@ using namespace mlir; using namespace mlir::qtensor; LogicalResult InsertOp::verify() { + auto destType = getDest().getType(); if (!llvm::isa(getScalar().getType())) { return emitOpError("Scalar must be of qubit type"); } - if (!llvm::isa(getDest().getType().getElementType())) { + if (!llvm::isa(destType.getElementType())) { return emitOpError("Elements of dest tensor must be of qubit type"); } + if (destType.getRank() != static_cast(getIndices().size())) { + return emitOpError("incorrect number of indices for insert"); + } return success(); } From b9824b25720c6aece471840500d1539b7e52acc1 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 13:51:24 +0100 Subject: [PATCH 041/108] fix typo --- mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 32c4812106..f650d30b3f 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -323,7 +323,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * auto outTensor = builder.insertSlice(slicedTensor, tensor, 0, 2, 1); * ``` * ```mlir - * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%0][%c2][%c1] + * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] * : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> * ``` */ From 892443adbf75ead967b5c1d2113cc799b584950f Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 14:15:07 +0100 Subject: [PATCH 042/108] rename test name --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 1d8881c7a7..df9c2961b4 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1040,13 +1040,14 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(extractTensor)}, QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(insertTensor), MQT_NAMED_BUILDER(insertTensor)}, - QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(dynamicInsertTensor), + QCOTestCase{"InsertTensorDynamicIndex", + MQT_NAMED_BUILDER(dynamicInsertTensor), MQT_NAMED_BUILDER(dynamicInsertTensor)}, QCOTestCase{"ExtractSliceTensor", MQT_NAMED_BUILDER(extractSliceTensor), MQT_NAMED_BUILDER(extractSliceTensor)}, QCOTestCase{"InsertSliceTensor", MQT_NAMED_BUILDER(insertSliceTensor), MQT_NAMED_BUILDER(insertSliceTensor)}, - QCOTestCase{"InsertSliceTensor", + QCOTestCase{"InsertSliceTensorDynamicIndex", MQT_NAMED_BUILDER(dynamicInsertSliceTensor), MQT_NAMED_BUILDER(dynamicInsertSliceTensor)}, QCOTestCase{"ExtractInsert", MQT_NAMED_BUILDER(extractInsertTensor), From 16541bbbba27c02718ba59c6f4565eb7e8a689be Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 14:17:29 +0100 Subject: [PATCH 043/108] separate tensor and qubit tracking --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 20 ++++++ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 65 ++++++++++++++----- 2 files changed, 70 insertions(+), 15 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index f650d30b3f..1a1eb9f52e 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1320,5 +1320,25 @@ 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 + * @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/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index d272d0e309..1836646dcc 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -168,6 +168,28 @@ 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)"); + } +} + +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 //===----------------------------------------------------------------------===// @@ -180,7 +202,7 @@ Value QCOProgramBuilder::allocTensor(int64_t size) { } auto allocOp = qtensor::AllocOp::create(*this, size); - validQubits.insert(allocOp); + validTensors.insert(allocOp); return allocOp.getResult(); } @@ -196,7 +218,7 @@ Value QCOProgramBuilder::fromElements(ValueRange elements) { } auto fromElementsOp = qtensor::FromElementsOp::create(*this, elements); - validQubits.insert(fromElementsOp); + validTensors.insert(fromElementsOp); return fromElementsOp.getResult(); } @@ -220,7 +242,7 @@ QCOProgramBuilder::extract(Value tensor, auto outTensor = extractOp.getOutTensor(); validQubits.insert(qubit); - updateQubitTracking(tensor, outTensor); + updateTensorTracking(tensor, outTensor); return {qubit, outTensor}; } @@ -249,8 +271,8 @@ QCOProgramBuilder::extractSlice(Value tensor, auto slicedTensor = extractSliceOp.getResult(); auto outTensor = extractSliceOp.getOutSource(); - validQubits.insert(slicedTensor); - updateQubitTracking(tensor, outTensor); + validTensors.insert(slicedTensor); + updateTensorTracking(tensor, outTensor); return {slicedTensor, outTensor}; } @@ -275,7 +297,7 @@ Value QCOProgramBuilder::insert(Value scalar, Value tensor, validateQubitValue(scalar); validQubits.erase(scalar); - updateQubitTracking(tensor, outTensor); + updateTensorTracking(tensor, outTensor); return outTensor; } @@ -311,9 +333,9 @@ Value QCOProgramBuilder::insertSlice( auto outTensor = insertSliceOp.getResult(); - validateQubitValue(source); - validQubits.erase(source); - updateQubitTracking(dest, outTensor); + validateTensorValue(source); + validTensors.erase(source); + updateTensorTracking(dest, outTensor); return outTensor; } @@ -332,8 +354,8 @@ QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { qtensor::DeallocOp::create(*this, tensor); - validateQubitValue(tensor); - validQubits.erase(tensor); + validateTensorValue(tensor); + validTensors.erase(tensor); return *this; } @@ -964,14 +986,27 @@ OwningOpRef QCOProgramBuilder::finalize() { return opA->isBeforeInBlock(opB); }); for (auto qubit : sortedQubits) { - if (llvm::isa(qubit.getType())) { - DeallocOp::create(*this, qubit); - } else { - qtensor::DeallocOp::create(*this, qubit); + 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, [](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); + }); + for (auto tensor : sortedTensors) { + qtensor::DeallocOp::create(*this, tensor); } validQubits.clear(); + validTensors.clear(); // Create constant 0 for successful exit code auto exitCode = intConstant(0); From 4c1554212f49a24a1f8d3c07115df9b51577c136 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 14:39:00 +0100 Subject: [PATCH 044/108] fix linter issues --- mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp | 2 +- mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index 3b8c16d746..8566112cbe 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -31,7 +31,7 @@ void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { } LogicalResult AllocOp::verify() { - auto resultType = dyn_cast(getResult().getType()); + auto resultType = getResult().getType(); if (!resultType) { return emitOpError("Result must be a ranked tensor"); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 1ce9a43748..54c938ef29 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -18,6 +18,8 @@ #include #include +#include + using namespace mlir; using namespace mlir::qtensor; From cdc77e8c9cc7abd00624454e90c8c2359ee4cf9d Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 14:48:37 +0100 Subject: [PATCH 045/108] address more coderabbit feedback --- .../mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 2 +- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 2 +- mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 4 ++++ .../QTensor/IR/Operations/ExtractSliceOp.cpp | 14 +++++++------- 4 files changed, 13 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 1a1eb9f52e..3faebef10e 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -232,7 +232,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * auto tensor = builder.allocTensor(3); * ``` * ```mlir - * %tensor = qtensor.alloc : tensor<3x!qco.qubit> + * %tensor = qtensor.alloc(3) : tensor<3x!qco.qubit> * ``` */ Value allocTensor(int64_t size); diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index e321f9ba65..44f5898428 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -78,7 +78,7 @@ class QTensorOpWithOffsetSizesAndStrides { + MemoryEffects<[MemAlloc]>]> { let summary = "Tensor alloc operation"; let description = [{ Allocates a qubit tensor with the given size and returns the allocated tensor. diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 1836646dcc..3c5380bed8 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -209,6 +209,10 @@ Value QCOProgramBuilder::allocTensor(int64_t size) { Value QCOProgramBuilder::fromElements(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!"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 2797457f4d..1dd5c67f6b 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -266,7 +266,7 @@ LogicalResult ExtractSliceOp::verify() { } llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { - return ::getDroppedDims(getType().getShape(), getMixedSizes()); + return mlir::qtensor::getDroppedDims(getType().getShape(), getMixedSizes()); } namespace { @@ -381,24 +381,24 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, ExtractSliceOpCastFolder>(context); } -static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { +static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { auto insertOp = extractOp.getSource().getDefiningOp(); auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; if (insertOp && insertOp.getSource().getType() == extractOp.getType() && insertOp.isSameAs(extractOp, isSame)) { - return insertOp.getSource(); + return insertOp; } - return {}; + return nullptr; } LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { - if (Value slice = foldExtractAfterInsertSlice(*this)) { - results.push_back(slice); - results.push_back(getSource()); + if (auto insertOp = foldExtractAfterInsertSlice(*this)) { + results.push_back(insertOp.getSource()); + results.push_back(insertOp.getDest()); return success(); } From 8473d1dec8915b4732f94bcb7c6a40cd4a044445 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 15:05:58 +0100 Subject: [PATCH 046/108] more coderabbit suggestions --- mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 14 ++++++++------ .../QTensor/IR/Operations/AllocTensorOp.cpp | 3 --- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 3c5380bed8..dba117a547 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -202,8 +202,9 @@ Value QCOProgramBuilder::allocTensor(int64_t size) { } auto allocOp = qtensor::AllocOp::create(*this, size); - validTensors.insert(allocOp); - return allocOp.getResult(); + auto result = allocOp.getResult(); + validTensors.insert(result); + return result; } Value QCOProgramBuilder::fromElements(ValueRange elements) { @@ -222,8 +223,9 @@ Value QCOProgramBuilder::fromElements(ValueRange elements) { } auto fromElementsOp = qtensor::FromElementsOp::create(*this, elements); - validTensors.insert(fromElementsOp); - return fromElementsOp.getResult(); + auto result = fromElementsOp.getResult(); + validTensors.insert(result); + return result; } std::pair @@ -356,10 +358,10 @@ QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { llvm::reportFatalUsageError("Elements must be of QubitType!"); } - qtensor::DeallocOp::create(*this, tensor); - validateTensorValue(tensor); validTensors.erase(tensor); + qtensor::DeallocOp::create(*this, tensor); + return *this; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index 8566112cbe..5778ce327d 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -32,9 +32,6 @@ void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { LogicalResult AllocOp::verify() { auto resultType = getResult().getType(); - if (!resultType) { - return emitOpError("Result must be a ranked tensor"); - } if (!llvm::isa(resultType.getElementType())) { return emitOpError("Result element type must be of qubit type"); From 5109e32cf4abf9c5960854db28a095a7831afc43 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 9 Mar 2026 15:27:08 +0100 Subject: [PATCH 047/108] more coderabbit feedback --- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 24 +++++++++---------- .../QTensor/IR/Operations/AllocTensorOp.cpp | 3 +++ 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index dba117a547..1d38f3ed81 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -360,6 +360,7 @@ QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { validateTensorValue(tensor); validTensors.erase(tensor); + qtensor::DeallocOp::create(*this, tensor); return *this; @@ -980,17 +981,20 @@ 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); } @@ -999,14 +1003,8 @@ OwningOpRef QCOProgramBuilder::finalize() { // Sort tensors for deterministic output llvm::SmallVector sortedTensors(validTensors.begin(), validTensors.end()); - llvm::sort(sortedTensors, [](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); - }); + llvm::sort(sortedTensors, blockOrderComparator); + for (auto tensor : sortedTensors) { qtensor::DeallocOp::create(*this, tensor); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index 5778ce327d..f80bed78c8 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -18,12 +18,15 @@ #include #include +#include #include using namespace mlir; using namespace mlir::qtensor; void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { + assert(size > 0 && "qtensor.alloc size must be positive"); + auto resultType = RankedTensorType::get({size}, qco::QubitType::get(builder.getContext())); build(builder, result, resultType, From b2766ed7af37c4b53ff5e06b85865249dd886f0d Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 10:24:02 +0100 Subject: [PATCH 048/108] address review feedback --- docs/mlir/QTensor.md | 8 ++++++ docs/mlir/index.md | 2 ++ .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 2 -- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 28 +++++++++---------- 4 files changed, 23 insertions(+), 17 deletions(-) create mode 100644 docs/mlir/QTensor.md diff --git a/docs/mlir/QTensor.md b/docs/mlir/QTensor.md new file mode 100644 index 0000000000..ae63cb6eae --- /dev/null +++ b/docs/mlir/QTensor.md @@ -0,0 +1,8 @@ +--- +tocdepth: 3 +--- + +```{include} Dialects/MLIRQTensorDialect.md + +:heading-offset: 1 +``` diff --git a/docs/mlir/index.md b/docs/mlir/index.md index 3b5e91b20c..fc769e1f66 100644 --- a/docs/mlir/index.md +++ b/docs/mlir/index.md @@ -8,6 +8,8 @@ We define multiple dialects, each with its dedicated purpose: - The {doc}`QCO dialect ` uses value semantics and is mainly designed for running optimizations. +- The {doc}`QTensor dialect ` between dialects. diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h index 4f7e832057..6ebfd5f966 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -50,5 +50,3 @@ foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, //===----------------------------------------------------------------------===// // QTensor Dialect Helpers //===----------------------------------------------------------------------===// - -namespace mlir::qtensor {} // namespace mlir::qtensor diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 44f5898428..3b37ad1d04 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -32,11 +32,9 @@ def QTensorDialect : Dialect { let summary = "The QTensor dialect for tensors with linear typing."; let description = [{ - The Qtensor dialect is an adjusted variant of the standard tensor dialect - of MLIR that supports linear typing for tensor types. The extract operations - of this dialect are modified so that they also return the updated input tensor - as return value. In addition, the alloc/dealloc operations are added to the dialect - to support the bulk allocation and deallocation of tensors with linear types. + The Qtensor dialect is an adjusted variant of the standard tensor dialect of MLIR that supports linear typing for tensor types. + 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 tensors with linear types. }]; let hasCanonicalizer = 1; @@ -104,9 +102,9 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ def QTensor_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. + This operation deallocates the tensor and the values it holds, releasing its resources. - Example: + Example: ```mlir qtensor.dealloc %tensor : <3x!qco.qubit> ``` @@ -128,7 +126,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ ]> { let summary = "tensor from elements operation"; let description = [{ - The `qtensor.from_elements`is a wrapper operation for the `tensor.from_elements` operation + The `qtensor.from_elements` is a wrapper operation for the `tensor.from_elements` operation of the tensor dialect. This operation creates a tensor using the operands as its values. During the canonicalization this operation is converted into a `tensor.from_elements` operation. @@ -358,13 +356,13 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let description = [{ The `qtensor.insert` operation is a wrapper operation for the `tensor.insert` operation of the tensor dialect. This operation inserts a scalar into a ranked tensor as specified by the operation`s indices. - During the canonicalization this operation is converted into a `tensor.insert` operation. + During the canonicalization, this operation is converted into a `tensor.insert` operation. Example: - ```mlir - %newTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> - ``` + ```mlir + %newTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> + ``` }]; let arguments = (ins AnyType:$scalar, @@ -411,9 +409,9 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ a dynamic value. Example: - ```mlir - %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] - : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + ```mlir + %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] + : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> ``` }]; From a108608c30bf033cac23a1172f64dc473007a390 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 10:24:29 +0100 Subject: [PATCH 049/108] correct error message --- mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp index f80bed78c8..307915ff00 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp @@ -47,7 +47,7 @@ LogicalResult AllocOp::verify() { } if (resultType.getShape()[0] != size) { - return emitOpError("Tensor length must match size operand (") + return emitOpError("Tensor length must match size attribute (") << size << "), but got " << resultType.getShape()[0]; } From c91e17daec78e5a1eaa5fa90002fa11393e797b8 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 10:54:18 +0100 Subject: [PATCH 050/108] use newer style of operation creation --- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 712750b8e4..472d7587c0 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -177,8 +177,8 @@ struct FoldTensorCastProducerOp for (auto [oldResult, newResult] : llvm::zip(op->getResults(), newOp->getResults())) { if (newResult.getType() != oldResult.getType()) { - replacements.push_back(rewriter.create( - op->getLoc(), oldResult.getType(), newResult)); + replacements.push_back(tensor::CastOp::create( + rewriter, op->getLoc(), oldResult.getType(), newResult)); } else { replacements.push_back(newResult); } From 840a0fb6a35951a3ddd4388cf4d162176bc1dfae Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 10:55:11 +0100 Subject: [PATCH 051/108] add missing word --- docs/mlir/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/mlir/index.md b/docs/mlir/index.md index fc769e1f66..5fb4770436 100644 --- a/docs/mlir/index.md +++ b/docs/mlir/index.md @@ -8,7 +8,7 @@ We define multiple dialects, each with its dedicated purpose: - The {doc}`QCO dialect ` uses value semantics and is mainly designed for running optimizations. -- The {doc}`QTensor dialect Date: Tue, 10 Mar 2026 11:00:42 +0100 Subject: [PATCH 052/108] fix name in CMakeLists --- mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt index 81bd22fbc9..1fedc2350e 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt @@ -7,4 +7,4 @@ # Licensed under the MIT License add_mlir_dialect(QTensorOps qtensor) -add_mlir_doc(QTensorOps QTensorDialect Dialects/ -gen-dialect-doc) +add_mlir_doc(QTensorOps MLIRQTensorDialect Dialects/ -gen-dialect-doc) From c6593d6b2449e12d615fdaa8796fc8b10ef1b8d6 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 11:19:14 +0100 Subject: [PATCH 053/108] more coderabbit fixes --- docs/mlir/index.md | 5 +++-- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 22 +++++++++---------- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/docs/mlir/index.md b/docs/mlir/index.md index 5fb4770436..42ff96def3 100644 --- a/docs/mlir/index.md +++ b/docs/mlir/index.md @@ -8,9 +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. -- The {doc}`QTensor dialect ` adds support for tensors with linear typing and is used in the QCO dialect to represent registers. -Both dialects define various canonicalization and transformations that enable the compilation of quantum programs to native quantum hardware. +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. @@ -20,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/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 3b37ad1d04..f1288e0e9f 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -106,7 +106,7 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { Example: ```mlir - qtensor.dealloc %tensor : <3x!qco.qubit> + qtensor.dealloc %tensor : tensor<3x!qco.qubit> ``` }]; @@ -210,8 +210,7 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", Example: ```mlir - %extractedSlice, %outSourcetensor = qtensor.extract_slice %sourceTensor[%c0][%c2][%c1] : - : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + %extractedSlice, %outSourceTensor = qtensor.extract_slice %sourceTensor[%c0][%c2][%c1] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> ``` }]; @@ -402,16 +401,15 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ * strides: tensor-rank number of strides that specify subsampling in each dimension. - The representation based on offsets, sizes and strides support a - partially-static specification via attributes specified through the - `static_offsets`, `static_sizes` and `static_strides` arguments. A special - sentinel value ShapedType::kDynamic encodes that the corresponding entry has - a dynamic value. + The representation based on offsets, sizes and strides support a + partially-static specification via attributes specified through the + `static_offsets`, `static_sizes` and `static_strides` arguments. A special + sentinel value ShapedType::kDynamic encodes that the corresponding entry has + a dynamic value. - Example: - ```mlir - %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] - : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + Example: + ```mlir + %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> ``` }]; From 76459db774c7be295da67fef3b824dd843751b36 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 11:51:11 +0100 Subject: [PATCH 054/108] more coderabbit fixes --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index f1288e0e9f..08734983ed 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -484,4 +484,4 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ let hasVerifier = 1; } -#endif // TENSOR_OPS +#endif // QTENSOROPS From 1d78b38b3d5a8f5e7989b60432805cd8fc7ecf1c Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 15:30:52 +0100 Subject: [PATCH 055/108] make strides optional and set default values --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 12 ++--- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 24 +++++---- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 52 ++++++++++++++++--- .../QTensor/IR/Operations/InsertSliceOp.cpp | 14 +++-- mlir/unittests/programs/qco_programs.cpp | 14 ++--- 5 files changed, 82 insertions(+), 34 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 3faebef10e..01d285d616 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -274,12 +274,12 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @param tensor Source tensor * @param offset The offset from where the slice is extracted * @param size The size of the extracted slice - * @param strides The strides from where the values are extracted + * @param strides The strides from where the values are extracted (default: 1) * @return Pair of (extractedSlice, outTensor) * * @par Example: * ```c++ - * auto [extractedSlice, outTensor] = builder.extractSlice(tensor, 0, 2, 1); + * auto [extractedSlice, outTensor] = builder.extractSlice(tensor, 0, 2); * ``` * ```mlir * %extractedSlice, %outTensor = qtensor.extract_slice %tensor[%c0][%c2][%c1] @@ -289,7 +289,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { std::pair extractSlice(Value tensor, const std::variant& offset, const std::variant& sizes, - const std::variant& strides); + const std::variant& strides = 1); /** * @brief Insert a qubit into a tensor @@ -315,12 +315,12 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @param destTensor The tensor where the slice is inserted * @param offset The offset into where the slice is inserted * @param size The size of the inserted slice - * @param strides The strides of where the qubits are inserted + * @param strides The strides of where the qubits are inserted (default: 1) * @return The output tensor * * @par Example: * ```c++ - * auto outTensor = builder.insertSlice(slicedTensor, tensor, 0, 2, 1); + * auto outTensor = builder.insertSlice(slicedTensor, tensor, 0, 2); * ``` * ```mlir * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] @@ -330,7 +330,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { Value insertSlice(Value sourceTensor, Value destTensor, const std::variant& offset, const std::variant& sizes, - const std::variant& strides); + const std::variant& strides = 1); /** * @brief Explicitly deallocate a tensor diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 08734983ed..5c5924b195 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -190,7 +190,9 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", 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 ranked tensor and returns the extracted tensor specified by the - offsets, sizes and strides arguments. In addition, it also returns the updated input tensor as result. + offsets, sizes and strides arguments. If the strides arguments are omitted, they are set to default values of 1. + In addition, it also returns the updated input tensor as result. + The extract_slice operation supports the following arguments: @@ -237,29 +239,29 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", // Build an ExtractSliceOp with mixed static and dynamic entries and // inferred result type. OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, - "ArrayRef":$sizes, "ArrayRef":$strides, + "ArrayRef":$sizes, CArg<"ArrayRef", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with mixed static and dynamic entries and custom // result type. If the type passed is nullptr, it is inferred. OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with dynamic entries and custom result type. If // the type passed is nullptr, it is inferred. OpBuilder<(ins "Value":$source, "ValueRange":$offsets, - "ValueRange":$sizes, "ValueRange":$strides, + "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with dynamic entries and inferred result type. OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + "ValueRange":$offsets, "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + "ValueRange":$offsets, "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an ExtractSliceOp with mixed static and dynamic entries packed in // a Range vector. @@ -269,7 +271,7 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", // result type, offsets set to 0 and strides set to 1. OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, "ArrayRef":$sizes, - CArg<"ArrayRef", "{}">:$attrs)>, + CArg<"ArrayRef", "{}">:$attrs)> ]; let extraClassDeclaration = extraBaseClassDeclaration # [{ @@ -388,7 +390,7 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ let description = [{ The `qtensor.insert_slice` operation is a minimally adjusted 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, sizes and strides arguments. + operation`s offset, sizes and strides arguments. If the strides arguments are omitted, they are set to default values of 1. The insert_slice operation supports the following arguments: @@ -437,11 +439,11 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ // result type. OpBuilder<(ins "Value":$source, "Value":$dest, "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build a InsertSliceOp with dynamic entries and inferred result type. OpBuilder<(ins "Value":$source, "Value":$dest, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + "ValueRange":$offsets, "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, CArg<"ArrayRef", "{}">:$attrs)>, // Build an InsertSliceOp with mixed static and dynamic entries packed in // a Range vector and inferred result type. diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 1dd5c67f6b..311bf7525f 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -29,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -145,9 +147,15 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, SmallVector dynamicOffsets; SmallVector dynamicSizes; SmallVector dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + if (strides.empty()) { + staticStrides.resize(sizes.size(), 1); + } else { + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + } + auto sourceRankedTensorType = llvm::cast(source.getType()); // Structuring implementation this way avoids duplication between builders. if (!resultType) { @@ -168,7 +176,11 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, ArrayRef strides, ArrayRef attrs) { build(b, result, resultType, cast(source.getType()), source, - offsets, sizes, strides, attrs); + offsets, sizes, + strides.empty() + ? SmallVector(sizes.size(), b.getIndexAttr(1)) + : strides, + attrs); } /// Build an ExtractSliceOp with mixed static and dynamic entries and inferred @@ -179,7 +191,11 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, ArrayRef strides, ArrayRef attrs) { build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, strides, attrs); + source, offsets, sizes, + strides.empty() + ? SmallVector(sizes.size(), b.getIndexAttr(1)) + : strides, + attrs); } /// Build an ExtractSliceOp with mixed static and dynamic entries packed into @@ -203,25 +219,47 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); SmallVector sizeValues = llvm::to_vector<4>( llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = llvm::to_vector<4>( - llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = + strides.empty() + ? SmallVector(sizes.size(), b.getIndexAttr(1)) + : llvm::to_vector(llvm::map_range( + strides, [](Value v) -> OpFoldResult { return v; })); build(b, result, resultType, outSourceType, source, offsetValues, sizeValues, strideValues, attrs); } + void ExtractSliceOp::build(OpBuilder& b, OperationState& result, RankedTensorType resultType, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef attrs) { + + SmallVector defaultStrides; + + if (strides.empty()) { + auto one = b.create(result.location, 1); + defaultStrides.assign(sizes.size(), one); + } + build(b, result, resultType, cast(source.getType()), source, - offsets, sizes, strides, attrs); + offsets, sizes, strides.empty() ? ValueRange(defaultStrides) : strides, + attrs); } /// Build an ExtractSliceOp with dynamic entries and inferred result type. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef attrs) { + + SmallVector defaultStrides; + + if (strides.empty()) { + auto one = b.create(result.location, 1); + defaultStrides.assign(sizes.size(), one); + } + build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, strides, attrs); + source, offsets, sizes, + strides.empty() ? ValueRange(defaultStrides) : strides, attrs); } /// Build an ExtractSliceOp with mixed static and dynamic sizes, inferred diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 7d57024142..b612c8d5a1 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -56,7 +56,12 @@ void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, SmallVector dynamicStrides; dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + if (strides.empty()) { + staticStrides.resize(sizes.size(), 1); + } else { + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + } + result.addAttributes(attrs); build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), @@ -81,8 +86,11 @@ void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); SmallVector sizeValues = llvm::to_vector<4>( llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = llvm::to_vector<4>( - llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = + strides.empty() + ? SmallVector(sizes.size(), b.getIndexAttr(1)) + : llvm::to_vector(llvm::map_range( + strides, [](Value v) -> OpFoldResult { return v; })); build(b, result, source, dest, offsetValues, sizeValues, strideValues, attrs); } diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index a8b470bd80..de015fb90d 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2087,16 +2087,16 @@ void dynamicInsertTensor(QCOProgramBuilder& b) { void extractSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); - b.extractSlice(qtensor, 0, 2, 1); + b.extractSlice(qtensor, 0, 2); } void insertSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); + auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2); auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); auto q1 = b.h(q0); auto insertOutTensor = b.insert(q1, extractOutTensor, 0); - b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2, 1); + b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } void dynamicInsertSliceTensor(QCOProgramBuilder& b) { @@ -2120,16 +2120,16 @@ void extractInsertTensor(QCOProgramBuilder& b) { void extractSliceInsertSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); - b.insertSlice(slicedTensor, extractSliceOutTensor, 0, 2, 1); + auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2); + b.insertSlice(slicedTensor, extractSliceOutTensor, 0, 2); } void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2, 1); + auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2); auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); auto insertOutTensor = b.insert(q0, extractOutTensor, 0); - b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2, 1); + b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } } // namespace mlir::qco From 8cb17d9849085650d2f5591e27271a05622ca83b Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 17:11:20 +0100 Subject: [PATCH 056/108] remove conversion to tensor operations for insertOp and fromElementsOp --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 8 +++--- .../QTensor/IR/Operations/ExtractOp.cpp | 2 +- .../QTensor/IR/Operations/FromTensorOp.cpp | 25 ------------------- .../QTensor/IR/Operations/InsertOp.cpp | 25 ++++++++++++++----- 4 files changed, 25 insertions(+), 35 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 5c5924b195..3ce4c071ac 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -128,7 +128,6 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ let description = [{ The `qtensor.from_elements` is a wrapper operation for the `tensor.from_elements` operation of the tensor dialect. This operation creates a tensor using the operands as its values. - During the canonicalization this operation is converted into a `tensor.from_elements` operation. Example: @@ -147,7 +146,6 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ ]; let hasVerifier = 1; - let hasCanonicalizer = 1; } def QTensor_ExtractOp : QTensorOp<"extract", [ @@ -346,6 +344,7 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", } def QTensor_InsertOp : QTensorOp<"insert", [ + DestinationStyleOpInterface, Pure, TypesMatchWith<"result type matches type of dest", "dest", "result", @@ -357,7 +356,6 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let description = [{ The `qtensor.insert` operation is a wrapper operation for the `tensor.insert` operation of the tensor dialect. This operation inserts a scalar into a ranked tensor as specified by the operation`s indices. - During the canonicalization, this operation is converted into a `tensor.insert` operation. Example: @@ -374,6 +372,10 @@ def QTensor_InsertOp : QTensorOp<"insert", [ $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) }]; + let extraClassDeclaration = [{ + MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } + }]; + let hasVerifier = 1; let hasCanonicalizer = 1; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index d60d9a5fbc..af65b4a798 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -67,7 +67,7 @@ struct ExtractFromTensorCast : public OpRewritePattern { * return the scalar from the InsertOp directly. */ static Value foldExtractAfterInsert(ExtractOp extractOp) { - auto insertOp = extractOp.getTensor().getDefiningOp(); + auto insertOp = extractOp.getTensor().getDefiningOp(); auto isSame = [](Value a, Value b) { return getAsOpFoldResult(a) == getAsOpFoldResult(b); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 70cbb88ddd..2df24badc5 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -12,12 +12,9 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include -#include #include -#include #include #include -#include #include #include @@ -50,25 +47,3 @@ LogicalResult FromElementsOp::verify() { } return success(); } - -namespace { - -struct ConvertFromElementsOpToTensorOp - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(qtensor::FromElementsOp fromElementsOp, - PatternRewriter& rewriter) const final { - rewriter.replaceOpWithNewOp( - fromElementsOp, fromElementsOp.getType(), fromElementsOp.getElements()); - - return success(); - } -}; - -} // namespace - -void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 54c938ef29..e90d562deb 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -12,7 +12,6 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include -#include #include #include #include @@ -39,14 +38,28 @@ LogicalResult InsertOp::verify() { namespace { -struct ConvertInsertOpToTensorOp : public OpRewritePattern { +/** + * @brief If an InsertOp does not return a tensor with a static shape but the + * destination tensor has one, replace the InsertOp with a new one that has a + * static shape. + */ +struct ConvertInsertOpToStaticShape + : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(qtensor::InsertOp insertOp, PatternRewriter& rewriter) const final { - rewriter.replaceOpWithNewOp( - insertOp, insertOp.getScalar(), insertOp.getDest(), - insertOp.getIndices()); + if (insertOp.getResult().getType().hasStaticShape()) { + return failure(); + } + if (!insertOp.getDest().getType().hasStaticShape()) { + return failure(); + } + + rewriter.replaceOpWithNewOp(insertOp, insertOp.getScalar(), + insertOp.getDest(), + insertOp.getIndices()); + return success(); } }; @@ -82,5 +95,5 @@ struct InsertFromExtractOp : public OpRewritePattern { void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add(context); } From df1f5d844895e0e4c810bd08d5f0ea74a4f65c53 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 17:43:57 +0100 Subject: [PATCH 057/108] address coderabbit suggestions --- .../lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 13 ++++++------- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 3 ++- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index af65b4a798..987de1f236 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -66,7 +66,7 @@ struct ExtractFromTensorCast : public OpRewritePattern { * @brief If an ExtractOp consumes an InsertOp with identical indices, * return the scalar from the InsertOp directly. */ -static Value foldExtractAfterInsert(ExtractOp extractOp) { +static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { auto insertOp = extractOp.getTensor().getDefiningOp(); auto isSame = [](Value a, Value b) { @@ -74,17 +74,16 @@ static Value foldExtractAfterInsert(ExtractOp extractOp) { }; if (insertOp && insertOp.getScalar().getType() == extractOp.getType(0) && llvm::equal(insertOp.getIndices(), extractOp.getIndices(), isSame)) { - return insertOp.getScalar(); + return insertOp; } - return {}; + return nullptr; } LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { - if (Value result = foldExtractAfterInsert(*this)) { - results.push_back(result); - results.push_back(getTensor()); - return success(); + if (auto insertOp = foldExtractAfterInsert(*this)) { + results.push_back(insertOp.getScalar()); + results.push_back(insertOp.getDest()); } return failure(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 311bf7525f..7a3c48a104 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -117,7 +117,8 @@ RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( } } inferredType = - RankedTensorType::get(projectedShape, inferredType.getElementType()); + RankedTensorType::get(projectedShape, inferredType.getElementType(), + inferredType.getEncoding()); } return inferredType; } From 7cc00fa321d091f73b71eb8483edc7416e75e0d2 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Tue, 10 Mar 2026 17:44:24 +0100 Subject: [PATCH 058/108] remove redundant tests --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 6 ------ mlir/unittests/programs/qco_programs.cpp | 21 ------------------- mlir/unittests/programs/qco_programs.h | 6 ------ 3 files changed, 33 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 6196426c33..d3ffa951d1 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1042,16 +1042,10 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(extractTensor)}, QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(insertTensor), MQT_NAMED_BUILDER(insertTensor)}, - QCOTestCase{"InsertTensorDynamicIndex", - MQT_NAMED_BUILDER(dynamicInsertTensor), - MQT_NAMED_BUILDER(dynamicInsertTensor)}, QCOTestCase{"ExtractSliceTensor", MQT_NAMED_BUILDER(extractSliceTensor), MQT_NAMED_BUILDER(extractSliceTensor)}, QCOTestCase{"InsertSliceTensor", MQT_NAMED_BUILDER(insertSliceTensor), MQT_NAMED_BUILDER(insertSliceTensor)}, - QCOTestCase{"InsertSliceTensorDynamicIndex", - MQT_NAMED_BUILDER(dynamicInsertSliceTensor), - MQT_NAMED_BUILDER(dynamicInsertSliceTensor)}, QCOTestCase{"ExtractInsert", MQT_NAMED_BUILDER(extractInsertTensor), MQT_NAMED_BUILDER(allocTensor)}, QCOTestCase{"ExtractSliceInsertSlice", diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index f79d3a4aa6..d99fa1d42a 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2082,14 +2082,6 @@ void insertTensor(QCOProgramBuilder& b) { b.insert(q1, extractOutTensor, 0); } -void dynamicInsertTensor(QCOProgramBuilder& b) { - auto c0 = b.indexConstant(0); - auto qtensor = b.allocTensor(3); - auto [q0, extractOutTensor] = b.extract(qtensor, c0); - auto q1 = b.h(q0); - b.insert(q1, extractOutTensor, c0); -} - void extractSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); b.extractSlice(qtensor, 0, 2); @@ -2104,19 +2096,6 @@ void insertSliceTensor(QCOProgramBuilder& b) { b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } -void dynamicInsertSliceTensor(QCOProgramBuilder& b) { - auto c0 = b.indexConstant(0); - auto c1 = b.indexConstant(1); - auto c2 = b.indexConstant(2); - auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = - b.extractSlice(qtensor, c0, c2, c1); - auto [q0, extractOutTensor] = b.extract(slicedTensor, c0); - auto q1 = b.h(q0); - auto insertOutTensor = b.insert(q1, extractOutTensor, c0); - b.insertSlice(insertOutTensor, extractSliceOutTensor, c0, c2, c1); -} - void extractInsertTensor(QCOProgramBuilder& b) { auto qtensor = b.allocTensor(3); auto [q0, extractOutTensor] = b.extract(qtensor, 0); diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 4b5e882cb9..1d913c2113 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -951,18 +951,12 @@ void extractTensor(QCOProgramBuilder& b); /// Inserts a qubit in a tensor. void insertTensor(QCOProgramBuilder& b); -/// Inserts a qubit in a tensor with dynamic index. -void dynamicInsertTensor(QCOProgramBuilder& b); - /// Extracts a slice from a tensor. void extractSliceTensor(QCOProgramBuilder& b); /// Inserts a slice from a tensor. void insertSliceTensor(QCOProgramBuilder& b); -/// Inserts a slice from a tensor with dynamic indices. -void dynamicInsertSliceTensor(QCOProgramBuilder& b); - /// Extracts a qubit from a tensor and insert it immediately. void extractInsertTensor(QCOProgramBuilder& b); From 38d8d37bed7d4e572e3e11adc6feaa3b00e3b2b6 Mon Sep 17 00:00:00 2001 From: li-mingbao <74404929+li-mingbao@users.noreply.github.com> Date: Thu, 12 Mar 2026 17:36:53 +0100 Subject: [PATCH 059/108] Apply suggestions from code review Co-authored-by: Daniel Haag <121057143+denialhaag@users.noreply.github.com> Signed-off-by: li-mingbao <74404929+li-mingbao@users.noreply.github.com> --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 3ce4c071ac..21e335035c 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -82,7 +82,7 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ Allocates a qubit tensor with the given size and returns the allocated tensor. The qubits are initialized to the |0> state. - Example : + Example: ```mlir %qtensor = qtensor.alloc(3) : tensor<3x!qco.qubit> ``` @@ -126,14 +126,11 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ ]> { let summary = "tensor from elements operation"; let description = [{ - The `qtensor.from_elements` is a wrapper operation for the `tensor.from_elements` operation - of the tensor dialect. This operation creates a tensor using the operands as its values. + The `qtensor.from_elements` operation is a wrapper operation of the `tensor.from_elements` operation. + This operation creates a tensor using the operands as its values. Example: - - ```mlir - %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> - ``` + }]; let arguments = (ins Variadic:$elements); From 7cf19f3caa0410b4595f3c9e4bec206f4347828f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Mar 2026 16:37:19 +0000 Subject: [PATCH 060/108] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 21e335035c..b9fa11014f 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -130,7 +130,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ This operation creates a tensor using the operands as its values. Example: - + }]; let arguments = (ins Variadic:$elements); From b27e5d95ef697b323f8bcd404abc754dd8471127 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 11 Mar 2026 20:12:45 +0100 Subject: [PATCH 061/108] simplify implementation by only allowing dynamic size and offset --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 12 +- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 116 ++--- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 16 +- .../QTensor/IR/Operations/ExtractOp.cpp | 13 +- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 380 ++++------------ .../QTensor/IR/Operations/InsertOp.cpp | 11 +- .../QTensor/IR/Operations/InsertSliceOp.cpp | 412 +++++------------- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 5 - 8 files changed, 232 insertions(+), 733 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 01d285d616..9178022647 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -274,7 +274,6 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @param tensor Source tensor * @param offset The offset from where the slice is extracted * @param size The size of the extracted slice - * @param strides The strides from where the values are extracted (default: 1) * @return Pair of (extractedSlice, outTensor) * * @par Example: @@ -282,14 +281,13 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * auto [extractedSlice, outTensor] = builder.extractSlice(tensor, 0, 2); * ``` * ```mlir - * %extractedSlice, %outTensor = qtensor.extract_slice %tensor[%c0][%c2][%c1] + * %extractedSlice, %outTensor = qtensor.extract_slice %tensor[%c0][%c2] * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> * ``` */ std::pair extractSlice(Value tensor, const std::variant& offset, - const std::variant& sizes, - const std::variant& strides = 1); + const std::variant& sizes); /** * @brief Insert a qubit into a tensor @@ -315,7 +313,6 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @param destTensor The tensor where the slice is inserted * @param offset The offset into where the slice is inserted * @param size The size of the inserted slice - * @param strides The strides of where the qubits are inserted (default: 1) * @return The output tensor * * @par Example: @@ -323,14 +320,13 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * auto outTensor = builder.insertSlice(slicedTensor, tensor, 0, 2); * ``` * ```mlir - * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2][%c1] + * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2] * : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> * ``` */ Value insertSlice(Value sourceTensor, Value destTensor, const std::variant& offset, - const std::variant& sizes, - const std::variant& strides = 1); + const std::variant& sizes); /** * @brief Explicitly deallocate a tensor diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index b9fa11014f..d38b26bf57 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -71,6 +71,25 @@ class QTensorOpWithOffsetSizesAndStrides traits = []> + : QTensorOp { + code extraBaseClassDeclaration = [{ + /// Return the type of the base tensor operand. + ::mlir::RankedTensorType getSourceType() { + return ::llvm::cast(getSource().getType()); + } + + /// Return the type of the result tensor. + ::mlir::RankedTensorType getResultType() { + return ::llvm::cast(getResult().getType()); + } + + /// Return the dynamic sizes for this subview operation if specified. + ::mlir::TypedValue getDynamicSizes() { return getSize(); } + + }]; +} //===----------------------------------------------------------------------===// // Operations //===----------------------------------------------------------------------===// @@ -165,19 +184,17 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ ``` }]; - let arguments = (ins AnyRankedTensor:$tensor, Variadic:$indices); + let arguments = (ins AnyRankedTensor:$tensor, Index:$index); let results = (outs AnyType:$result, AnyRankedTensor:$out_tensor); - let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)"; + let assemblyFormat = "$tensor `[` $index `]` attr-dict `:` type($tensor)"; let hasFolder = 1; let hasVerifier = 1; let hasCanonicalizer = 1; } -def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", [ - AttrSizedOperandSegments, +def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStridess<"extract_slice", [ Pure, - OffsetSizeAndStrideOpInterface, TypesMatchWith<"returned tensor type matches input tensor", "source", "out_source", "$_self"> ]> { @@ -213,59 +230,19 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStrides<"extract_slice", let arguments = (ins AnyRankedTensor:$source, - Variadic:$offsets, - Variadic:$sizes, - Variadic:$strides, - DenseI64ArrayAttr:$static_offsets, - DenseI64ArrayAttr:$static_sizes, - DenseI64ArrayAttr:$static_strides + Index:$offset, + Index:$size ); let results = (outs AnyRankedTensor:$result, AnyRankedTensor:$out_source); let assemblyFormat = [{ - $source `` - custom($offsets, $static_offsets) - custom($sizes, $static_sizes) - custom($strides, $static_strides) + $source `[`$offset `]` `[`$size`]` attr-dict `:` type($source) `to` type($result) }]; let builders = [ - // Build an ExtractSliceOp with mixed static and dynamic entries and - // inferred result type. - OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, - "ArrayRef":$sizes, CArg<"ArrayRef", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with mixed static and dynamic entries and custom - // result type. If the type passed is nullptr, it is inferred. - OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, - "ArrayRef":$offsets, "ArrayRef":$sizes, - CArg<"ArrayRef", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ArrayRef":$offsets, "ArrayRef":$sizes, - CArg<"ArrayRef", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with dynamic entries and custom result type. If - // the type passed is nullptr, it is inferred. - OpBuilder<(ins "Value":$source, "ValueRange":$offsets, - "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with dynamic entries and inferred result type. - OpBuilder<(ins "RankedTensorType":$resultType, "RankedTensorType":$outSourceType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with mixed static and dynamic entries packed in - // a Range vector. - OpBuilder<(ins "Value":$source, "ArrayRef":$ranges, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an ExtractSliceOp with mixed static and dynamic sizes, inferred - // result type, offsets set to 0 and strides set to 1. - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ArrayRef":$sizes, + OpBuilder<(ins "Value":$source, "Value":$offset, + "Value":$size, CArg<"ArrayRef", "{}">:$attrs)> ]; @@ -352,7 +329,7 @@ def QTensor_InsertOp : QTensorOp<"insert", [ 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 scalar into a ranked tensor as specified by the operation`s indices. + This operation inserts a scalar into a ranked tensor as specified by the operation`s index. Example: @@ -363,10 +340,10 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let arguments = (ins AnyType:$scalar, AnyRankedTensor:$dest, - Variadic:$indices); + Index:$index); let results = (outs AnyRankedTensor:$result); let assemblyFormat = [{ - $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) + $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) }]; let extraClassDeclaration = [{ @@ -377,11 +354,9 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let hasCanonicalizer = 1; } -def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ - AttrSizedOperandSegments, +def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStridess<"insert_slice", [ DestinationStyleOpInterface, Pure, - OffsetSizeAndStrideOpInterface, TypesMatchWith<"expected result type to match dest type", "dest", "result", "$_self"> ]> { @@ -417,37 +392,18 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStrides<"insert_slice", [ let arguments = (ins AnyRankedTensor:$source, AnyRankedTensor:$dest, - Variadic:$offsets, - Variadic:$sizes, - Variadic:$strides, - DenseI64ArrayAttr:$static_offsets, - DenseI64ArrayAttr:$static_sizes, - DenseI64ArrayAttr:$static_strides + Index:$offset, + Index:$size ); let results = (outs AnyRankedTensor:$result); let assemblyFormat = [{ - $source `into` $dest `` - custom($offsets, $static_offsets) - custom($sizes, $static_sizes) - custom($strides, $static_strides) + $source `into` $dest `[`$offset `]` `[`$size`]` attr-dict `:` type($source) `into` type($dest) }]; let builders = [ - // Build a InsertSliceOp with mixed static and dynamic entries and inferred - // result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ArrayRef":$offsets, "ArrayRef":$sizes, - CArg<"ArrayRef", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build a InsertSliceOp with dynamic entries and inferred result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ValueRange":$offsets, "ValueRange":$sizes, CArg<"ValueRange", "{}">:$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build an InsertSliceOp with mixed static and dynamic entries packed in - // a Range vector and inferred result type. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ArrayRef":$ranges, + OpBuilder<(ins "Value":$source, + "Value":$offset, "Value":$size, CArg<"ArrayRef", "{}">:$attrs)> ]; diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 1d38f3ed81..5e4a66ae22 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -256,8 +256,7 @@ QCOProgramBuilder::extract(Value tensor, std::pair QCOProgramBuilder::extractSlice(Value tensor, const std::variant& offset, - const std::variant& sizes, - const std::variant& strides) { + const std::variant& sizes) { checkFinalized(); auto tensorType = llvm::dyn_cast(tensor.getType()); @@ -271,9 +270,8 @@ QCOProgramBuilder::extractSlice(Value tensor, auto offsetValue = utils::variantToValue(*this, getLoc(), offset); auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); - auto stridesValue = utils::variantToValue(*this, getLoc(), strides); - auto extractSliceOp = qtensor::ExtractSliceOp::create( - *this, tensor, offsetValue, sizesValue, stridesValue); + auto extractSliceOp = + qtensor::ExtractSliceOp::create(*this, tensor, offsetValue, sizesValue); auto slicedTensor = extractSliceOp.getResult(); auto outTensor = extractSliceOp.getOutSource(); @@ -309,8 +307,7 @@ Value QCOProgramBuilder::insert(Value scalar, Value tensor, Value QCOProgramBuilder::insertSlice( Value source, Value dest, const std::variant& offset, - const std::variant& sizes, - const std::variant& strides) { + const std::variant& sizes) { checkFinalized(); auto sourceTensorType = llvm::dyn_cast(source.getType()); @@ -333,9 +330,8 @@ Value QCOProgramBuilder::insertSlice( auto offsetValue = utils::variantToValue(*this, getLoc(), offset); auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); - auto stridesValue = utils::variantToValue(*this, getLoc(), strides); - auto insertSliceOp = qtensor::InsertSliceOp::create( - *this, source, dest, offsetValue, sizesValue, stridesValue); + auto insertSliceOp = qtensor::InsertSliceOp::create(*this, source, dest, + offsetValue, sizesValue); auto outTensor = insertSliceOp.getResult(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 987de1f236..6763388c1f 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -24,8 +24,6 @@ #include #include -#include - using namespace mlir; using namespace mlir::qtensor; @@ -33,14 +31,10 @@ using namespace mlir::qtensor; // https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp LogicalResult ExtractOp::verify() { - // Verify the # indices match if we have a ranked type. auto tensorType = llvm::cast(getTensor().getType()); if (!llvm::isa(tensorType.getElementType())) { return emitOpError("Elements of tensor must be of qubit type"); } - if (tensorType.getRank() != static_cast(getIndices().size())) { - return emitOpError("incorrect number of indices for extract_element"); - } return success(); } @@ -57,7 +51,7 @@ struct ExtractFromTensorCast : public OpRewritePattern { return failure(); } rewriter.replaceOpWithNewOp(extract, tensorCast.getSource(), - extract.getIndices()); + extract.getIndex()); return success(); } }; @@ -69,11 +63,8 @@ struct ExtractFromTensorCast : public OpRewritePattern { static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { auto insertOp = extractOp.getTensor().getDefiningOp(); - auto isSame = [](Value a, Value b) { - return getAsOpFoldResult(a) == getAsOpFoldResult(b); - }; if (insertOp && insertOp.getScalar().getType() == extractOp.getType(0) && - llvm::equal(insertOp.getIndices(), extractOp.getIndices(), isSame)) { + insertOp.getIndex() == extractOp.getIndex()) { return insertOp; } return nullptr; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 7a3c48a104..74e20528ec 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -37,7 +37,6 @@ #include #include -#include using namespace mlir; using namespace mlir::qtensor; @@ -45,270 +44,68 @@ using namespace mlir::qtensor; // Adjusted from // https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp -/** - * @brief Infers the result type of an extract_slice operation when it is - * not rank-reduced. - * - * @details - * The result type can be inferred from the source tensor type and the - * static representation of offsets, sizes, and strides. Special sentinel - * values are used to encode dynamic entries. - * - * @param sourceTensorType The ranked source tensor type. - * @param staticSizes The static sizes of the slice (sentinel values - * indicate dynamic sizes). - * @return The inferred RankedTensorType for the resulting slice. - */ -RankedTensorType -ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, - ArrayRef staticSizes) { - // An extract_slice op may specify only a leading subset of offset/sizes/ - // strides in which case we complete with offset=0, sizes from memref type - // and strides=1. - assert(static_cast(staticSizes.size()) == - sourceTensorType.getRank() && - "unexpected staticSizes not equal to rank of source"); - return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), - sourceTensorType.getEncoding()); -} - -RankedTensorType -ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType, - ArrayRef sizes) { - SmallVector staticSizes; - std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes); - - assert(static_cast(staticSizes.size()) == - sourceTensorType.getRank() && - "unexpected staticSizes not equal to rank of source"); - return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(), - sourceTensorType.getEncoding()); -} - -/** - * @brief Computes the rank-reduced result type. - * - * @details - * If the desired result rank is smaller than the number of slice sizes, - * rank reduction is performed by dropping dimensions of size 1 until the - * desired rank is reached. - * - * Multiple rank-reduced shapes may be possible. For example, a tensor of - * shape 1x6x1 can be reduced to either 1x6 or 6x1. To ensure deterministic - * behavior, this function always drops the first occurrences of size-1 - * dimensions. - */ -RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( - unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, - ArrayRef sizes) { - // Type inferred in the absence of rank-reducing behavior. - auto inferredType = llvm::cast( - inferResultType(sourceRankedTensorType, sizes)); - int64_t rankDiff = inferredType.getRank() - desiredResultRank; - if (rankDiff > 0) { - auto shape = inferredType.getShape(); - llvm::SmallBitVector dimsToProject = - getPositionsOfShapeOne(rankDiff, shape); - SmallVector projectedShape; - // Best effort rank-reducing: drop 1s in order. - for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) { - if (!dimsToProject.test(pos)) { - projectedShape.push_back(shape[pos]); - } - } - inferredType = - RankedTensorType::get(projectedShape, inferredType.getElementType(), - inferredType.getEncoding()); - } - return inferredType; -} - -RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( - unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, - ArrayRef sizes) { - SmallVector staticSizes; - SmallVector dynamicSizes; - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - return ExtractSliceOp::inferCanonicalRankReducedResultType( - desiredResultRank, sourceRankedTensorType, staticSizes); -} - -/// Build an ExtractSliceOp with mixed static and dynamic entries and custom -/// result type. If the type passed is nullptr, it is inferred. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, - RankedTensorType outSourceType, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - SmallVector staticOffsets; - SmallVector staticSizes; - SmallVector staticStrides; - SmallVector dynamicOffsets; - SmallVector dynamicSizes; - SmallVector dynamicStrides; - - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - if (strides.empty()) { - staticStrides.resize(sizes.size(), 1); - } else { - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); - } - - auto sourceRankedTensorType = llvm::cast(source.getType()); - // Structuring implementation this way avoids duplication between builders. - if (!resultType) { - resultType = llvm::cast( - ExtractSliceOp::inferResultType(sourceRankedTensorType, staticSizes)); - } - result.addAttributes(attrs); - build(b, result, {resultType, outSourceType}, source, dynamicOffsets, - dynamicSizes, dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), - b.getDenseI64ArrayAttr(staticSizes), - b.getDenseI64ArrayAttr(staticStrides)); -} - -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - build(b, result, resultType, cast(source.getType()), source, - offsets, sizes, - strides.empty() - ? SmallVector(sizes.size(), b.getIndexAttr(1)) - : strides, - attrs); -} - -/// Build an ExtractSliceOp with mixed static and dynamic entries and inferred -/// result type. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, - strides.empty() - ? SmallVector(sizes.size(), b.getIndexAttr(1)) - : strides, - attrs); -} - -/// Build an ExtractSliceOp with mixed static and dynamic entries packed into -/// a Range vector. +/// Build an ExtractSliceOp with dynamic entries and inferred result type. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, - ArrayRef ranges, + Value offset, Value size, ArrayRef attrs) { - auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); - build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, strides, attrs); -} + auto optionalVal = getConstantIntValue(size); + auto resultType = RankedTensorType::get( + {optionalVal ? *optionalVal : ShapedType::kDynamic}, + cast(source.getType()).getElementType()); -/// Build an ExtractSliceOp with dynamic entries and custom result type. If -/// the type passed is nullptr, it is inferred. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, - RankedTensorType outSourceType, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { - SmallVector offsetValues = llvm::to_vector<4>( - llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); - SmallVector sizeValues = llvm::to_vector<4>( - llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = - strides.empty() - ? SmallVector(sizes.size(), b.getIndexAttr(1)) - : llvm::to_vector(llvm::map_range( - strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, resultType, outSourceType, source, offsetValues, sizeValues, - strideValues, attrs); + result.addAttributes(attrs); + build(b, result, {resultType, source.getType()}, source, offset, size); } -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { - - SmallVector defaultStrides; +/// Verifier for ExtractSliceOp. +LogicalResult ExtractSliceOp::verify() { + RankedTensorType sourceType = getSourceType(); - if (strides.empty()) { - auto one = b.create(result.location, 1); - defaultStrides.assign(sizes.size(), one); + // Element type check + if (!llvm::isa(sourceType.getElementType())) { + return emitOpError("Elements of source tensor must be of qubit type"); } - build(b, result, resultType, cast(source.getType()), source, - offsets, sizes, strides.empty() ? ValueRange(defaultStrides) : strides, - attrs); -} - -/// Build an ExtractSliceOp with dynamic entries and inferred result type. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { - - SmallVector defaultStrides; + auto srcDim = sourceType.getDimSize(0); - if (strides.empty()) { - auto one = b.create(result.location, 1); - defaultStrides.assign(sizes.size(), one); - } - - build(b, result, RankedTensorType(), cast(source.getType()), - source, offsets, sizes, - strides.empty() ? ValueRange(defaultStrides) : strides, attrs); -} + if (auto constSize = getConstantIntValue(getSize())) { + if (*constSize < 0) { + return emitOpError("Size must be non-negative"); + } -/// Build an ExtractSliceOp with mixed static and dynamic sizes, inferred -/// result type, offsets set to 0 and strides set to 1. -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, - RankedTensorType resultType, Value source, - ArrayRef sizes, - ArrayRef attrs) { - Attribute zeroIdxAttr = b.getIndexAttr(0); - Attribute oneIdxAttr = b.getIndexAttr(1); - SmallVector readStrides(sizes.size(), oneIdxAttr); - SmallVector readOffsets(sizes.size(), zeroIdxAttr); - build(b, result, resultType, cast(source.getType()), source, - readOffsets, sizes, readStrides, attrs); -} + // Check size fits in source + if (!ShapedType::isDynamic(srcDim) && *constSize > srcDim) { + return emitOpError("Size exceeds source dimension"); + } -/// Verifier for ExtractSliceOp. -LogicalResult ExtractSliceOp::verify() { - RankedTensorType sourceType = getSourceType(); + if (auto constOffset = getConstantIntValue(getOffset())) { + if (*constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } - if (!llvm::isa(sourceType.getElementType())) { - return emitOpError("Elements of tensor must be of qubit type"); - } - // Verify result type against inferred type. - RankedTensorType expectedType = - ExtractSliceOp::inferResultType(sourceType, getMixedSizes()); - SliceVerificationResult result = isRankReducedType(expectedType, getType()); - if (result != SliceVerificationResult::Success) { - return produceSliceErrorMsg(result, *this, expectedType); + if (!ShapedType::isDynamic(srcDim) && + *constOffset + *constSize > srcDim) { + return emitOpError("Offset + Size exceeds source dimension"); + } + } + } else if (auto constOffset = getConstantIntValue(getOffset())) { + if (*constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } + if (!ShapedType::isDynamic(srcDim) && *constOffset >= srcDim) { + return emitOpError("Offset out of bounds"); + } } - // Verify that offsets, sizes, strides do not run out-of-bounds with respect - // to the source tensor. - SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( - sourceType.getShape(), getStaticOffsets(), getStaticSizes(), - getStaticStrides(), /*generateErrorMessage=*/true); - if (!boundsResult.isValid) { - return getOperation()->emitError(boundsResult.errorMessage); + // Verify result slice type matches source element type + RankedTensorType resultType = getOutSource().getType(); // or getResult() + if (resultType.getElementType() != sourceType.getElementType()) { + return emitOpError("result element type must match source element type"); } return success(); } -llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { - return mlir::qtensor::getDroppedDims(getType().getShape(), getMixedSizes()); -} - -namespace { /** * @brief Rewrite pattern that pushes tensor.cast past tensor.extract_slice. * @@ -334,19 +131,22 @@ namespace { * This effectively folds the cast into the consumer operation and enables * further canonicalization opportunities. */ +namespace { + class ExtractSliceOpCastFolder final : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(ExtractSliceOp sliceOp, PatternRewriter& rewriter) const override { - // Any constant operand, just return to let the constant folder kick in. - if (llvm::any_of(sliceOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) { + + // Let constant folding handle constant operands + if (matchPattern(sliceOp.getOffset(), matchConstantIndex()) || + matchPattern(sliceOp.getSize(), matchConstantIndex())) { return failure(); } + // Look for tensor.cast producer auto castOp = sliceOp.getSource().getDefiningOp(); if (!castOp) { return failure(); @@ -356,76 +156,57 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern { return failure(); } - // Pattern does not apply if the produced op would not verify. - SliceBoundsVerificationResult sliceResult = verifyInBoundsSlice( - cast(castOp.getSource().getType()).getShape(), - sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), - sliceOp.getStaticStrides()); - if (!sliceResult.isValid) { - return failure(); + // Verify bounds using the original tensor + auto srcType = cast(castOp.getSource().getType()); + + int64_t dim = srcType.getShape()[0]; + + auto offsetVal = getConstantIntValue(sliceOp.getOffset()); + auto sizeVal = getConstantIntValue(sliceOp.getSize()); + + if (offsetVal && sizeVal) { + if (*offsetVal + *sizeVal > dim) { + return failure(); + } } - // Create folded extract. Location loc = sliceOp.getLoc(); - auto newResult = ExtractSliceOp::create( - rewriter, loc, sliceOp.getType(), castOp.getSource(), - sliceOp.getOffsets(), sliceOp.getSizes(), sliceOp.getStrides(), - sliceOp.getStaticOffsets(), sliceOp.getStaticSizes(), - sliceOp.getStaticStrides()); - Value newOutSource = newResult->getResult(1); + + // Create new slice directly on the original tensor + auto newSlice = rewriter.create( + loc, sliceOp.getResult().getType(), sliceOp.getOutSource().getType(), + castOp.getSource(), sliceOp.getOffset(), sliceOp.getSize()); + + Value newResult = newSlice.getResult(); + Value newOutSource = newSlice.getOutSource(); + + // Preserve expected type of out_source if (newOutSource.getType() != sliceOp.getOutSource().getType()) { - newOutSource = tensor::CastOp::create( - rewriter, loc, sliceOp.getOutSource().getType(), newOutSource); + newOutSource = rewriter.create( + loc, sliceOp.getOutSource().getType(), newOutSource); } - rewriter.replaceOp(sliceOp, {newResult->getResult(0), newOutSource}); + + rewriter.replaceOp(sliceOp, {newResult, newOutSource}); + if (castOp->use_empty()) { rewriter.eraseOp(castOp); } + return success(); } }; } // namespace -/// Return the canonical type of the result of an extract_slice op. -struct SliceReturnTypeCanonicalizer { - RankedTensorType operator()(ExtractSliceOp op, - ArrayRef /*mixedOffsets*/, - ArrayRef mixedSizes, - ArrayRef /*mixedStrides*/) { - return ExtractSliceOp::inferCanonicalRankReducedResultType( - op.getType().getRank(), op.getSourceType(), mixedSizes); - } -}; - -/// A canonicalizer wrapper to replace ExtractSliceOps. -struct SliceCanonicalizer { - void operator()(PatternRewriter& rewriter, ExtractSliceOp op, - ExtractSliceOp newOp) { - Value replacement = newOp.getResult(); - Value outSource = newOp.getOutSource(); - if (replacement.getType() != op.getType()) { - replacement = tensor::CastOp::create(rewriter, op.getLoc(), op.getType(), - replacement); - } - rewriter.replaceOp(op, {replacement, outSource}); - } -}; - void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add< - OpWithOffsetSizesAndStridesConstantArgumentFolder< - ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>, - ExtractSliceOpCastFolder>(context); + results.add(context); } static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { auto insertOp = extractOp.getSource().getDefiningOp(); - auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; - if (insertOp && insertOp.getSource().getType() == extractOp.getType() && - insertOp.isSameAs(extractOp, isSame)) { + if (insertOp && insertOp.getSource().getType() == extractOp.getType()) { return insertOp; } @@ -434,7 +215,6 @@ static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { - if (auto insertOp = foldExtractAfterInsertSlice(*this)) { results.push_back(insertOp.getSource()); results.push_back(insertOp.getDest()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index e90d562deb..6b8fd7baf3 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -17,8 +17,6 @@ #include #include -#include - using namespace mlir; using namespace mlir::qtensor; @@ -30,9 +28,6 @@ LogicalResult InsertOp::verify() { if (!llvm::isa(destType.getElementType())) { return emitOpError("Elements of dest tensor must be of qubit type"); } - if (destType.getRank() != static_cast(getIndices().size())) { - return emitOpError("incorrect number of indices for insert"); - } return success(); } @@ -58,7 +53,7 @@ struct ConvertInsertOpToStaticShape rewriter.replaceOpWithNewOp(insertOp, insertOp.getScalar(), insertOp.getDest(), - insertOp.getIndices()); + insertOp.getIndex()); return success(); } @@ -81,7 +76,7 @@ struct InsertFromExtractOp : public OpRewritePattern { if (insertOp.getDest() != extractOp.getOutTensor()) { return failure(); } - if (insertOp.getIndices() != extractOp.getIndices()) { + if (insertOp.getIndex() != extractOp.getIndex()) { return failure(); } @@ -95,5 +90,5 @@ struct InsertFromExtractOp : public OpRewritePattern { void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add(context); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index b612c8d5a1..e7c4a8d970 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -32,9 +32,7 @@ #include #include -#include #include -#include using namespace mlir; using namespace mlir::qtensor; @@ -42,81 +40,8 @@ using namespace mlir::qtensor; // Adjusted from // https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp -// Build a InsertSliceOp with mixed static and dynamic entries. -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - SmallVector staticOffsets; - SmallVector staticSizes; - SmallVector staticStrides; - SmallVector dynamicOffsets; - SmallVector dynamicSizes; - SmallVector dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes); - if (strides.empty()) { - staticStrides.resize(sizes.size(), 1); - } else { - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); - } - - result.addAttributes(attrs); - build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, - dynamicStrides, b.getDenseI64ArrayAttr(staticOffsets), - b.getDenseI64ArrayAttr(staticSizes), - b.getDenseI64ArrayAttr(staticStrides)); -} - -/// Build an InsertSliceOp with mixed static and dynamic entries packed into a -/// Range vector. -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ArrayRef ranges, - ArrayRef attrs) { - auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); - build(b, result, source, dest, offsets, sizes, strides, attrs); -} - -// Build a InsertSliceOp with dynamic entries. -void InsertSliceOp::build(OpBuilder& b, OperationState& result, Value source, - Value dest, ValueRange offsets, ValueRange sizes, - ValueRange strides, ArrayRef attrs) { - SmallVector offsetValues = llvm::to_vector<4>( - llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); - SmallVector sizeValues = llvm::to_vector<4>( - llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = - strides.empty() - ? SmallVector(sizes.size(), b.getIndexAttr(1)) - : llvm::to_vector(llvm::map_range( - strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, source, dest, offsetValues, sizeValues, strideValues, attrs); -} - -/// Rank-reducing type verification for both InsertSliceOp and -/// ParallelInsertSliceOp. -static SliceVerificationResult -verifyInsertSliceOp(RankedTensorType srcType, RankedTensorType dstType, - ArrayRef /*staticOffsets*/, - ArrayRef staticSizes, - ArrayRef /*staticStrides*/, - RankedTensorType* expectedType = nullptr) { - // insert_slice is the inverse of extract_slice, use the same type - // inference. - RankedTensorType expected = - ExtractSliceOp::inferResultType(dstType, staticSizes); - if (expectedType != nullptr) { - *expectedType = expected; - } - return isRankReducedType(expected, srcType); -} - /// Verifier for InsertSliceOp. LogicalResult InsertSliceOp::verify() { - // Verify result type against inferred type. - RankedTensorType expectedType; - if (!llvm::isa(getSourceType().getElementType())) { return emitOpError("Elements of source tensor must be of qubit type"); } @@ -124,20 +49,37 @@ LogicalResult InsertSliceOp::verify() { return emitOpError("Elements of dest tensor must be of qubit type"); } - SliceVerificationResult result = - verifyInsertSliceOp(getSourceType(), getType(), getStaticOffsets(), - getStaticSizes(), getStaticStrides(), &expectedType); - if (result != SliceVerificationResult::Success) { - return produceSliceErrorMsg(result, *this, expectedType); - } + auto dstDim = getDestType().getDimSize(0); + auto srcDim = getSourceType().getDimSize(0); - // Verify that offsets, sizes, strides do not run out-of-bounds with respect - // to the destination tensor. - SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice( - getDestType().getShape(), getStaticOffsets(), getStaticSizes(), - getStaticStrides(), /*generateErrorMessage=*/true); - if (!boundsResult.isValid) { - return getOperation()->emitError(boundsResult.errorMessage); + if (auto constSize = getConstantIntValue(getSize())) { + if (*constSize < 0) { + return emitOpError("Size must be non-negative"); + } + + // Check size fits in source + if (!ShapedType::isDynamic(srcDim) && *constSize > srcDim) { + return emitOpError("Size exceeds source dimension"); + } + + if (auto constOffset = getConstantIntValue(getOffset())) { + if (*constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } + + // Check slice fits in dest + if (!ShapedType::isDynamic(dstDim) && + *constOffset + *constSize > dstDim) { + return emitOpError("Offset + Size exceeds destination dimension"); + } + } + } else if (auto constOffset = getConstantIntValue(getOffset())) { + if (*constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } + if (!ShapedType::isDynamic(dstDim) && *constOffset >= dstDim) { + return emitOpError("Offset out of bounds"); + } } return success(); @@ -165,15 +107,32 @@ LogicalResult InsertSliceOp::verify() { * ``` */ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { + // Check if the destination of current insert is another insert auto prevInsertOp = insertOp.getDest().getDefiningOp(); + if (!prevInsertOp) { + return failure(); + } - auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; - if (!prevInsertOp || - prevInsertOp.getSource().getType() != insertOp.getSource().getType() || - !prevInsertOp.isSameAs(insertOp, isSame)) { + // Check source types + if (prevInsertOp.getSource().getType() != insertOp.getSource().getType()) { return failure(); } + // Check offset and size (only 1D now) + auto prevOffsetOpt = getConstantIntValue(prevInsertOp.getOffset()); + auto prevSizeOpt = getConstantIntValue(prevInsertOp.getSize()); + auto curOffsetOpt = getConstantIntValue(insertOp.getOffset()); + auto curSizeOpt = getConstantIntValue(insertOp.getSize()); + + // Only fold if offsets and sizes are **statically known and identical** + if (!prevOffsetOpt || !prevSizeOpt || !curOffsetOpt || !curSizeOpt) { + return failure(); + } + if (*prevOffsetOpt != *curOffsetOpt || *prevSizeOpt != *curSizeOpt) { + return failure(); + } + + // Fold: bypass previous insert insertOp.getDestMutable().assign(prevInsertOp.getDest()); return success(); } @@ -198,82 +157,58 @@ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { */ static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { auto extractOp = insertOp.getSource().getDefiningOp(); - auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; - if (!extractOp || extractOp.getOutSource() != insertOp.getDest() || - !extractOp.isSameAs(insertOp, isSame)) { + if (!extractOp) { + return nullptr; + } + + // Ensure the insert destination is the original source tensor of extract + if (extractOp.getOutSource() != insertOp.getDest()) { return nullptr; } + + // Optionally check that the offset and size match exactly + if (extractOp.getOffset() != insertOp.getOffset() || + extractOp.getSize() != insertOp.getSize()) { + return nullptr; + } + return extractOp.getSource(); } OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { - if (getSourceType().hasStaticShape() && getType().hasStaticShape() && - getSourceType() == getType() && - succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) { - return this->getSource(); + // Identity fold: full overwrite + if (getSource() == getDest()) { + if (auto constOffset = getConstantIntValue(getOffset())) { + if (*constOffset == 0) { + if (auto constSize = getConstantIntValue(getSize())) { + if (*constSize == getSourceType().getDimSize(0)) { + return getSource(); + } + } + } + } } + // Fold nested insert after insert if (succeeded(foldInsertAfterInsertSlice(*this))) { return getResult(); } + + // Fold after extract_slice if (auto result = foldInsertAfterExtractSlice(*this)) { return result; } - if (llvm::any_of(getMixedSizes(), isZeroInteger)) { - return getDest(); + + // Zero-length insert + if (auto constSize = getConstantIntValue(getSize())) { + if (*constSize == 0) { + return getDest(); + } } + return {}; } - namespace { -template -class InsertSliceOpConstantArgumentFolder final - : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, - PatternRewriter& rewriter) const override { - SmallVector mixedOffsets(insertSliceOp.getMixedOffsets()); - SmallVector mixedSizes(insertSliceOp.getMixedSizes()); - SmallVector mixedStrides(insertSliceOp.getMixedStrides()); - - // No constant operands were folded, just return; - if (failed(foldDynamicOffsetSizeList(mixedOffsets)) && - failed(foldDynamicOffsetSizeList(mixedSizes)) && - failed(foldDynamicStrideList(mixedStrides))) { - return failure(); - } - - // Pattern does not apply if the produced op would not verify. - SliceBoundsVerificationResult sliceResult = - verifyInBoundsSlice(insertSliceOp.getDest().getType().getShape(), - mixedOffsets, mixedSizes, mixedStrides); - if (!sliceResult.isValid) { - return failure(); - } - - // Create the new op in canonical form. - auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType( - insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(), - mixedSizes); - Value toInsert = insertSliceOp.getSource(); - if (sourceType != insertSliceOp.getSourceType()) { - OpBuilder::InsertionGuard g(rewriter); - // The only difference between InsertSliceOp and ParallelInsertSliceOp - // is that the insertion point is just before the InParallelOp in - // the parallel case. - - toInsert = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), - sourceType, toInsert); - } - rewriter.replaceOpWithNewOp( - insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets, - mixedSizes, mixedStrides); - return success(); - } -}; - /** * @brief Folds tensor.cast operations with insert_slice. * @@ -298,182 +233,37 @@ class InsertSliceOpConstantArgumentFolder final * When folding a cast on the destination tensor, the result of the * insert_slice operation is cast to preserve the original result type. */ -template -struct InsertSliceOpCastFolder final : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct InsertSliceOpCastFolder final : public OpRewritePattern { + + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, + LogicalResult matchAndRewrite(InsertSliceOp op, PatternRewriter& rewriter) const override { - if (llvm::any_of(insertSliceOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) { - return failure(); - } - auto getSourceOfCastOp = [](Value v) -> std::optional { - auto castOp = v.getDefiningOp(); - if (!castOp || !canFoldIntoConsumerOp(castOp)) { - return std::nullopt; - } - return castOp.getSource(); - }; - std::optional sourceCastSource = - getSourceOfCastOp(insertSliceOp.getSource()); - std::optional destCastSource = - getSourceOfCastOp(insertSliceOp.getDest()); - if (!sourceCastSource && !destCastSource) { - return failure(); - } + auto srcCast = op.getSource().getDefiningOp(); + auto dstCast = op.getDest().getDefiningOp(); - auto src = - (sourceCastSource ? *sourceCastSource : insertSliceOp.getSource()); - auto dst = (destCastSource ? *destCastSource : insertSliceOp.getDest()); - auto srcType = llvm::dyn_cast(src.getType()); - auto dstType = llvm::dyn_cast(dst.getType()); - if (!srcType || !dstType) { + if (!srcCast && !dstCast) { return failure(); } - // The tensor.cast source could have additional static information not seen - // in the insert slice op static sizes, so we ignore dynamic dims when - // computing the rank reduction mask. - SmallVector staticSizes(insertSliceOp.getStaticSizes()); - auto rankReductionMask = computeRankReductionMask( - staticSizes, srcType.getShape(), /*matchDynamic=*/true); - if (!rankReductionMask.has_value()) { - return failure(); - } - // Replace dimensions in the insert slice op with corresponding static dims - // from the cast source type. If the insert slice sizes have static dims - // that are not static in the tensor.cast source (i.e., when the cast op - // casts a dynamic dim to static), the dim should not be replaced, and the - // pattern will fail later in `verifyInsertSliceOp`. - SmallVector mixedSizes(insertSliceOp.getMixedSizes()); - int64_t rankReducedIdx = 0; - for (auto [idx, size] : enumerate(staticSizes)) { - if (!rankReductionMask.value().contains(idx) && - !srcType.isDynamicDim(rankReducedIdx)) { - mixedSizes[idx] = getAsIndexOpFoldResult( - rewriter.getContext(), srcType.getDimSize(rankReducedIdx)); - size = srcType.getDimSize(rankReducedIdx++); - } - } + Value newSrc = + srcCast ? srcCast.getSource() : static_cast(op.getSource()); + Value newDst = + dstCast ? dstCast.getSource() : static_cast(op.getDest()); - // Pattern does not apply if the produced op would not verify. - if (verifyInsertSliceOp(srcType, dstType, insertSliceOp.getStaticOffsets(), - staticSizes, insertSliceOp.getStaticStrides()) != - SliceVerificationResult::Success) { - return failure(); - } - SliceBoundsVerificationResult sliceResult = - verifyInBoundsSlice(dstType.getShape(), insertSliceOp.getMixedOffsets(), - mixedSizes, insertSliceOp.getMixedStrides()); - if (!sliceResult.isValid) { - return failure(); - } + auto newOp = + rewriter.create(op.getLoc(), op.getType(), newSrc, + newDst, op.getOffset(), op.getSize()); - Operation* replacement = - InsertOpTy::create(rewriter, insertSliceOp.getLoc(), src, dst, - insertSliceOp.getMixedOffsets(), mixedSizes, - insertSliceOp.getMixedStrides()); - - // In the parallel case there is no result and so nothing to cast. - bool isParallelInsert = - std::is_same_v; - if (!isParallelInsert && dst.getType() != insertSliceOp.getDestType()) { - replacement = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), - insertSliceOp.getDestType(), - replacement->getResult(0)); - } - rewriter.replaceOp(insertSliceOp, replacement->getResults()); + rewriter.replaceOp(op, newOp->getResults()); return success(); } }; -/** - * @brief Inserts a tensor.cast before insert_slice when additional static - * type information can be inferred from the slice sizes. - * - * @details - * If the size operands of an insert_slice operation provide additional - * static shape information, an explicit tensor.cast is inserted on the - * source operand to refine its type. This enables further canonicalization - * patterns that match tensor.cast operations, such as - * `ForOpTensorCastFolder` in the SCF dialect. - * - * Example: - * - * ```mlir - * %r = qtensor.insert_slice %0 into %1[...] [3] [1] - * : tensor into ... - * ``` - * - * This folds into: - * - * ```mlir - * %tmp = tensor.cast %0 : tensor to tensor<3x!qco.qubit> - * %r = qtensor.insert_slice %tmp into %1[...] [3] [1] - * : tensor<3x!qco.qubit> into ... - * ``` - */ -template -struct InsertSliceOpSourceCastInserter final - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOpTy insertSliceOp, - PatternRewriter& rewriter) const override { - RankedTensorType srcType = insertSliceOp.getSourceType(); - if (srcType.getRank() != insertSliceOp.getDestType().getRank()) { - return failure(); - } - SmallVector newSrcShape(srcType.getShape()); - for (int64_t i = 0; i < srcType.getRank(); ++i) { - if (std::optional constInt = - getConstantIntValue(insertSliceOp.getMixedSizes()[i])) { - // Bail on invalid IR. - if (*constInt < 0) { - return failure(); - } - newSrcShape[i] = *constInt; - } - } - if (!hasValidSizesOffsets(newSrcShape)) { - return failure(); - } - - RankedTensorType newSrcType = RankedTensorType::get( - newSrcShape, srcType.getElementType(), srcType.getEncoding()); - if (srcType == newSrcType || - !mlir::tensor::preservesStaticInformation(srcType, newSrcType) || - !tensor::CastOp::areCastCompatible(srcType, newSrcType)) { - return failure(); - } - - // newSrcType is: - // 1) Different from srcType. - // 2) "More static" than srcType. - // 3) Cast-compatible with srcType. - // Insert the cast. - OpBuilder::InsertionGuard g(rewriter); - Value cast = tensor::CastOp::create(rewriter, insertSliceOp.getLoc(), - newSrcType, insertSliceOp.getSource()); - rewriter.replaceOpWithNewOp( - insertSliceOp, cast, insertSliceOp.getDest(), - insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(), - insertSliceOp.getMixedStrides()); - return success(); - } -}; } // namespace -llvm::SmallBitVector InsertSliceOp::getDroppedDims() { - return ::getDroppedDims(getSourceType().getShape(), getMixedSizes()); -} - void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add, - InsertSliceOpCastFolder, - InsertSliceOpSourceCastInserter>(context); + results.add(context); } diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 472d7587c0..9f7c2ec4be 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -124,11 +124,6 @@ foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, return failure(); } } - for (OpFoldResult opFold : op.getMixedStrides()) { - if (getConstantIntValue(opFold) != static_cast(1)) { - return failure(); - } - } return success(); } } // namespace mlir::qtensor From a1f86ad4081eff9bcb156ab2cd4fc8a25bff3db5 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 09:50:24 +0100 Subject: [PATCH 062/108] adjust extractSlice and insertSlice operations --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 180 ++---------------- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 5 +- .../QTensor/IR/Operations/InsertSliceOp.cpp | 23 ++- 3 files changed, 31 insertions(+), 177 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index d38b26bf57..5fdd3a1cc5 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -49,47 +49,6 @@ def QTensorDialect : Dialect { class QTensorOp traits = []> : Op; -// Base class for ops with static/dynamic offset, sizes and strides -// attributes/arguments. -class QTensorOpWithOffsetSizesAndStrides traits = []> - : QTensorOp { - code extraBaseClassDeclaration = [{ - /// Return the type of the base tensor operand. - ::mlir::RankedTensorType getSourceType() { - return ::llvm::cast(getSource().getType()); - } - - /// Return the type of the result tensor. - ::mlir::RankedTensorType getResultType() { - return ::llvm::cast(getResult().getType()); - } - - /// Return the dynamic sizes for this subview operation if specified. - ::mlir::Operation::operand_range getDynamicSizes() { return getSizes(); } - - }]; -} - -class QTensorOpWithOffsetSizesAndStridess traits = []> - : QTensorOp { - code extraBaseClassDeclaration = [{ - /// Return the type of the base tensor operand. - ::mlir::RankedTensorType getSourceType() { - return ::llvm::cast(getSource().getType()); - } - - /// Return the type of the result tensor. - ::mlir::RankedTensorType getResultType() { - return ::llvm::cast(getResult().getType()); - } - - /// Return the dynamic sizes for this subview operation if specified. - ::mlir::TypedValue getDynamicSizes() { return getSize(); } - - }]; -} //===----------------------------------------------------------------------===// // Operations //===----------------------------------------------------------------------===// @@ -143,7 +102,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ "::llvm::cast($_self).getNumElements(), " "::llvm::cast($_self).getElementType())"> ]> { - let summary = "tensor from elements operation"; + 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 tensor using the operands as its values. @@ -193,7 +152,7 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ let hasCanonicalizer = 1; } -def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStridess<"extract_slice", [ +def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ Pure, TypesMatchWith<"returned tensor type matches input tensor", "source", "out_source", "$_self"> @@ -202,29 +161,18 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStridess<"extract_slice" 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 ranked tensor and returns the extracted tensor specified by the - offsets, sizes and strides arguments. If the strides arguments are omitted, they are set to default values of 1. - In addition, it also returns the updated input tensor as result. - + offset and size argument. In addition, it also returns the updated input tensor as result. The extract_slice operation supports the following arguments: * source: the "base" tensor from which to extract a slice. - * offsets: tensor-rank number of offsets into the "base" tensor from which - to extract the slice. - * sizes: tensor-rank number of sizes which specify the sizes of the result - tensor type. - * strides: tensor-rank number of strides specifying subsampling in each - dimension. - - The representation based on offsets, sizes and strides support a - partially-static specification via attributes specified through the - `static_offsets`, `static_sizes` and `static_strides` arguments. A special - sentinel value ShapedType::kDynamic encodes that the corresponding entry has - a dynamic value. + * 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 - %extractedSlice, %outSourceTensor = qtensor.extract_slice %sourceTensor[%c0][%c2][%c1] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + %extractedSlice, %outSourceTensor = qtensor.extract_slice %sourceTensor[%c0][%c2] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> ``` }]; @@ -246,71 +194,7 @@ def QTensor_ExtractSliceOp : QTensorOpWithOffsetSizesAndStridess<"extract_slice" CArg<"ArrayRef", "{}">:$attrs)> ]; - let extraClassDeclaration = extraBaseClassDeclaration # [{ - /// The result of an extract_slice is always a tensor. - RankedTensorType getType() { - return getResultType(); - } - - /// Compute the rank-reduction mask that can be applied to map the source - /// tensor type to the result tensor type by dropping unit dims. - std::optional> - computeRankReductionMask() { - return ::mlir::computeRankReductionMask(getSourceType().getShape(), - getType().getShape()); - }; - - /// An extract_slice result type can be inferred, when it is not - /// rank-reduced, from the source type and the static representation of - /// sizes. Special sentinels encode the dynamic case. - static RankedTensorType inferResultType( - RankedTensorType sourceTensorType, - ArrayRef staticSizes); - static RankedTensorType inferResultType( - RankedTensorType sourceTensorType, - ArrayRef staticSizes); - - /// If the rank is reduced (i.e. the desiredResultRank is smaller than the - /// number of sizes), drop as many size 1 as needed to produce an inferred type - /// with the desired rank. - /// - /// Note that there may be multiple ways to compute this rank-reduced type: - /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. - /// - /// To disambiguate, this function always drops the first 1 sizes occurrences. - static RankedTensorType inferCanonicalRankReducedResultType( - unsigned resultRank, - RankedTensorType sourceRankedTensorType, - ArrayRef staticSizes); - static RankedTensorType inferCanonicalRankReducedResultType( - unsigned resultRank, - RankedTensorType sourceRankedTensorType, - ArrayRef staticSizes); - - /// Return the expected rank of each of the`static_offsets`, `static_sizes` - /// and `static_strides` attributes. - std::array getArrayAttrMaxRanks() { - unsigned rank = getSourceType().getRank(); - return {rank, rank, rank}; - } - - /// Return the number of leading operands before the `offsets`, `sizes` and - /// and `strides` operands. - static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } - - /// Return the dimensions of the source that are dropped in the - /// result when the result is rank-reduced. - llvm::SmallBitVector getDroppedDims(); - - /// Given a `value`, asserted to be of RankedTensorType, build an - /// ExtractSliceOp that results in a rank-reducing extract to the desired - /// tensor shape and return the new value created. - /// If the shape of `value` is already the `desiredShape`, just return - /// `value`. - /// If the shape of `value` cannot be rank-reduced to `desiredShape`, fail. - static FailureOr rankReduceIfNeeded( - OpBuilder &b, Location loc, Value value, ArrayRef desiredShape); - }]; + let hasCanonicalizer = 1; let hasFolder = 1; @@ -354,7 +238,7 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let hasCanonicalizer = 1; } -def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStridess<"insert_slice", [ +def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ DestinationStyleOpInterface, Pure, TypesMatchWith<"expected result type to match dest type", @@ -364,28 +248,19 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStridess<"insert_slice", let description = [{ The `qtensor.insert_slice` operation is a minimally adjusted 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, sizes and strides arguments. If the strides arguments are omitted, they are set to default values of 1. + operation`s offset and size argument. The insert_slice operation supports the following arguments: * source: the tensor that is inserted. * dest: the tensor into which the source tensor is inserted. - * offsets: tensor-rank number of offsets into the `dest` tensor into which - the slice is inserted. - * sizes: tensor-rank number of sizes which specify the sizes of the source - tensor type. - * strides: tensor-rank number of strides that specify subsampling in each - dimension. - - The representation based on offsets, sizes and strides support a - partially-static specification via attributes specified through the - `static_offsets`, `static_sizes` and `static_strides` arguments. A special - sentinel value ShapedType::kDynamic encodes that the corresponding entry has - a dynamic value. + * 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][%c1] : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> + %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2] : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> ``` }]; @@ -407,32 +282,7 @@ def QTensor_InsertSliceOp : QTensorOpWithOffsetSizesAndStridess<"insert_slice", CArg<"ArrayRef", "{}">:$attrs)> ]; - let extraClassDeclaration = extraBaseClassDeclaration # [{ - /// The result of a insert_slice is always a tensor. - RankedTensorType getType() { - return getResultType(); - } - - /// The `dest` type is the same as the result type. - RankedTensorType getDestType() { - return getResultType(); - } - - /// Return the expected rank of each of the`static_offsets`, `static_sizes` - /// and `static_strides` attributes. - std::array getArrayAttrMaxRanks() { - unsigned rank = getResultType().getRank(); - return {rank, rank, rank}; - } - - /// Return the dimensions of the dest that are omitted to insert a source - /// when the result is rank-extended. - llvm::SmallBitVector getDroppedDims(); - - /// Return the number of leading operands before the `offsets`, `sizes` and - /// and `strides` operands. - static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; } - + let extraClassDeclaration = [{ MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } }]; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 74e20528ec..f0baaf2cc5 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -59,7 +59,7 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, /// Verifier for ExtractSliceOp. LogicalResult ExtractSliceOp::verify() { - RankedTensorType sourceType = getSourceType(); + RankedTensorType sourceType = getSource().getType(); // Element type check if (!llvm::isa(sourceType.getElementType())) { @@ -206,7 +206,8 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { auto insertOp = extractOp.getSource().getDefiningOp(); - if (insertOp && insertOp.getSource().getType() == extractOp.getType()) { + if (insertOp && + insertOp.getSource().getType() == extractOp.getResult().getType()) { return insertOp; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index e7c4a8d970..8f0d49b3eb 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -42,15 +42,18 @@ using namespace mlir::qtensor; /// Verifier for InsertSliceOp. LogicalResult InsertSliceOp::verify() { - if (!llvm::isa(getSourceType().getElementType())) { + auto sourceType = getSource().getType(); + auto destType = getDest().getType(); + + if (!llvm::isa(sourceType.getElementType())) { return emitOpError("Elements of source tensor must be of qubit type"); } - if (!llvm::isa(getDestType().getElementType())) { + if (!llvm::isa(destType.getElementType())) { return emitOpError("Elements of dest tensor must be of qubit type"); } - auto dstDim = getDestType().getDimSize(0); - auto srcDim = getSourceType().getDimSize(0); + auto dstDim = destType.getDimSize(0); + auto srcDim = sourceType.getDimSize(0); if (auto constSize = getConstantIntValue(getSize())) { if (*constSize < 0) { @@ -96,14 +99,14 @@ LogicalResult InsertSliceOp::verify() { * Example: * * ```mlir - * %0 = qtensor.insert_slice %slice0 into %input[0][2][1] - * %1 = qtensor.insert_slice %slice1 into %0[0][2][1] + * %0 = qtensor.insert_slice %slice0 into %input[%c0][%c2] + * %1 = qtensor.insert_slice %slice1 into %0[%c0][%c2] * ``` * * This folds into: * * ```mlir - * %1 = qtensor.insert_slice %slice1 into %input[0][2][1] + * %1 = qtensor.insert_slice %slice1 into %input[%c0][%c2] * ``` */ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { @@ -118,13 +121,13 @@ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { return failure(); } - // Check offset and size (only 1D now) + // Check offset and size auto prevOffsetOpt = getConstantIntValue(prevInsertOp.getOffset()); auto prevSizeOpt = getConstantIntValue(prevInsertOp.getSize()); auto curOffsetOpt = getConstantIntValue(insertOp.getOffset()); auto curSizeOpt = getConstantIntValue(insertOp.getSize()); - // Only fold if offsets and sizes are **statically known and identical** + // Only fold if offsets and sizes are constant and identical if (!prevOffsetOpt || !prevSizeOpt || !curOffsetOpt || !curSizeOpt) { return failure(); } @@ -181,7 +184,7 @@ OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { if (auto constOffset = getConstantIntValue(getOffset())) { if (*constOffset == 0) { if (auto constSize = getConstantIntValue(getSize())) { - if (*constSize == getSourceType().getDimSize(0)) { + if (*constSize == getSource().getType().getDimSize(0)) { return getSource(); } } From c248490e311481ae127820466fdc6591c72c47e7 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 14:15:53 +0100 Subject: [PATCH 063/108] rename class names --- .../QTensor/IR/Operations/{AllocTensorOp.cpp => AllocOp.cpp} | 0 .../QTensor/IR/Operations/{DeallocTensorOp.cpp => DeallocOp.cpp} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename mlir/lib/Dialect/QTensor/IR/Operations/{AllocTensorOp.cpp => AllocOp.cpp} (100%) rename mlir/lib/Dialect/QTensor/IR/Operations/{DeallocTensorOp.cpp => DeallocOp.cpp} (100%) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp similarity index 100% rename from mlir/lib/Dialect/QTensor/IR/Operations/AllocTensorOp.cpp rename to mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp similarity index 100% rename from mlir/lib/Dialect/QTensor/IR/Operations/DeallocTensorOp.cpp rename to mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp From 81337392987b74a086489716c25faeb4e4104195 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 14:40:57 +0100 Subject: [PATCH 064/108] refactor implementation --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 4 +- .../QTensor/IR/Operations/ExtractOp.cpp | 38 ++++++++++- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 46 +++++++++++-- .../QTensor/IR/Operations/InsertOp.cpp | 65 ++++++++++++------- .../QTensor/IR/Operations/InsertSliceOp.cpp | 50 ++++++++++---- 5 files changed, 156 insertions(+), 47 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 5fdd3a1cc5..9315fd7f03 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -194,8 +194,6 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ CArg<"ArrayRef", "{}">:$attrs)> ]; - - let hasCanonicalizer = 1; let hasFolder = 1; let hasVerifier = 1; @@ -235,7 +233,7 @@ def QTensor_InsertOp : QTensorOp<"insert", [ }]; let hasVerifier = 1; - let hasCanonicalizer = 1; + let hasFolder = 1; } def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 6763388c1f..051f5950b6 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -35,6 +35,16 @@ LogicalResult ExtractOp::verify() { if (!llvm::isa(tensorType.getElementType())) { return emitOpError("Elements of tensor must be of qubit type"); } + auto index = getConstantIntValue(getIndex()); + auto size = getTensor().getType().getDimSize(0); + if (index) { + if (index < 0) { + return emitOpError("Index must be non-negative"); + } + if (index >= size) { + return emitOpError("Index exceeds tensor dimension"); + } + } return success(); } @@ -62,12 +72,34 @@ struct ExtractFromTensorCast : public OpRewritePattern { */ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { auto insertOp = extractOp.getTensor().getDefiningOp(); + if (!insertOp) { + return nullptr; + } + + if (insertOp.getScalar().getType() != extractOp.getType(0)) { + return nullptr; + } - if (insertOp && insertOp.getScalar().getType() == extractOp.getType(0) && - insertOp.getIndex() == extractOp.getIndex()) { + auto insertIndex = insertOp.getIndex(); + auto extractIndex = extractOp.getIndex(); + + // Check if SSA values of the indices are the same + if (insertIndex == extractIndex) { return insertOp; } - return nullptr; + + auto insertIndexValue = getConstantIntValue(insertIndex); + auto extractIndexValue = getConstantIntValue(extractIndex); + + // Check if the indices are constant and equal + if (!insertIndexValue || !extractIndexValue) { + return nullptr; + } + if (*insertIndexValue != *extractIndexValue) { + return nullptr; + } + + return insertOp; } LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index f0baaf2cc5..5770a3729b 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -203,15 +203,49 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, results.add(context); } -static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { - auto insertOp = extractOp.getSource().getDefiningOp(); +static InsertSliceOp +foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { + auto insertSliceOp = + extractSliceOp.getSource().getDefiningOp(); + if (!insertSliceOp) { + return nullptr; + } + + // Source types must match + if (insertSliceOp.getSource().getType() != extractSliceOp.getType(0)) { + return nullptr; + } + + auto insertOffset = insertSliceOp.getOffset(); + auto extractOffset = extractSliceOp.getOffset(); + auto insertSize = insertSliceOp.getSize(); + auto extractSize = extractSliceOp.getSize(); + + // Check if SSA values of the offsets and the sizes are the same + if (insertOffset == extractOffset && insertSize == extractSize) { + return insertSliceOp; + } + + auto insertOffsetValue = getConstantIntValue(insertOffset); + auto extractOffsetValue = getConstantIntValue(extractOffset); + auto insertSizeValue = getConstantIntValue(insertSize); + auto extractSizeValue = getConstantIntValue(extractSize); - if (insertOp && - insertOp.getSource().getType() == extractOp.getResult().getType()) { - return insertOp; + // Check if then offsets and sizes are constant and equal + if (!insertOffsetValue || !extractOffsetValue) { + return nullptr; + } + if (!insertSizeValue || !extractSizeValue) { + return nullptr; + } + if (*insertOffsetValue != *extractOffsetValue) { + return nullptr; + } + if (*insertSizeValue != *extractSizeValue) { + return nullptr; } - return nullptr; + return insertSliceOp; } LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 6b8fd7baf3..53c671fafe 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -28,6 +28,16 @@ LogicalResult InsertOp::verify() { if (!llvm::isa(destType.getElementType())) { return emitOpError("Elements of dest tensor must be of qubit type"); } + auto index = getConstantIntValue(getIndex()); + auto size = destType.getDimSize(0); + if (index) { + if (index < 0) { + return emitOpError("Index must be non-negative"); + } + if (index >= size) { + return emitOpError("Index exceeds tensor dimension"); + } + } return success(); } @@ -58,37 +68,48 @@ struct ConvertInsertOpToStaticShape return success(); } }; +} // namespace /** * @brief If an InsertOp consumes an ExtractOp with identical indices, * return the tensor from the extractOp directly. */ -struct InsertFromExtractOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +static Value foldInsertAfterExtract(InsertOp insertOp) { + auto extractOp = insertOp.getScalar().getDefiningOp(); + if (!extractOp) { + return nullptr; + } + if (insertOp.getDest() != extractOp.getOutTensor()) { + return nullptr; + } - LogicalResult matchAndRewrite(InsertOp insertOp, - PatternRewriter& rewriter) const final { - auto extractOp = insertOp.getScalar().getDefiningOp(); - if (!extractOp) { - return failure(); - } + auto insertIndex = insertOp.getIndex(); + auto extractIndex = extractOp.getIndex(); - if (insertOp.getDest() != extractOp.getOutTensor()) { - return failure(); - } - if (insertOp.getIndex() != extractOp.getIndex()) { - return failure(); - } + // Check if SSA values of the indices are the same + if (insertIndex == extractIndex) { + return extractOp.getTensor(); + } - rewriter.replaceOp(insertOp, extractOp.getTensor()); - rewriter.eraseOp(extractOp); - return success(); + auto insertIndexValue = getConstantIntValue(insertIndex); + auto extractIndexValue = getConstantIntValue(extractIndex); + + // Check if the indices are constant and equal + if (!insertIndexValue || !extractIndexValue) { + return nullptr; + } + if (*insertIndexValue != *extractIndexValue) { + return nullptr; } -}; -} // namespace + return extractOp.getTensor(); +} + +OpFoldResult InsertOp::fold(FoldAdaptor /*adaptor*/) { + // Fold after extract_slice + if (auto result = foldInsertAfterExtract(*this)) { + return result; + } -void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); + return {}; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 8f0d49b3eb..3679b068c1 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -145,37 +145,61 @@ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { * * @details * Detects patterns where a slice is extracted from a tensor and then - * inserted back into the same tensor at the same offsets, sizes, and - * strides. In such cases, the pair of operations forms a no-op and can + * inserted back into the same tensor at the same offset and size. + * In such cases, the pair of operations forms a no-op and can * be folded to the original tensor value. * * Example: * * ```mlir - * %0 = qtensor.extract_slice %val[0][2][1] - * %1 = qtensor.insert_slice %0 into %val[0][2][1] + * %slicedTensor, outTensor = qtensor.extract_slice %tensor[%c0][%c2] + * %newTensor = qtensor.insert_slice %slicedTensor into %outTensor[%c0][%c2] * ``` * - * This can be folded into `%val`. + * This can be folded into `%tensor`. */ -static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { - auto extractOp = insertOp.getSource().getDefiningOp(); - if (!extractOp) { +static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { + auto extractSliceOp = + insertSliceOp.getSource().getDefiningOp(); + if (!extractSliceOp) { return nullptr; } // Ensure the insert destination is the original source tensor of extract - if (extractOp.getOutSource() != insertOp.getDest()) { + if (extractSliceOp.getOutSource() != insertSliceOp.getDest()) { return nullptr; } - // Optionally check that the offset and size match exactly - if (extractOp.getOffset() != insertOp.getOffset() || - extractOp.getSize() != insertOp.getSize()) { + auto insertOffset = insertSliceOp.getOffset(); + auto extractOffset = extractSliceOp.getOffset(); + auto insertSize = insertSliceOp.getSize(); + auto extractSize = extractSliceOp.getSize(); + + // Check if SSA values of the offsets and the sizes are the same + if (insertOffset == extractOffset && insertSize == extractSize) { + return extractSliceOp.getSource(); + } + + auto insertOffsetValue = getConstantIntValue(insertOffset); + auto extractOffsetValue = getConstantIntValue(extractOffset); + auto insertSizeValue = getConstantIntValue(insertSize); + auto extractSizeValue = getConstantIntValue(extractSize); + + // Check if then offsets and sizes are constant and equal + if (!insertOffsetValue || !extractOffsetValue) { + return nullptr; + } + if (!insertSizeValue || !extractSizeValue) { + return nullptr; + } + if (*insertOffsetValue != *extractOffsetValue) { + return nullptr; + } + if (*insertSizeValue != *extractSizeValue) { return nullptr; } - return extractOp.getSource(); + return extractSliceOp.getSource(); } OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { From caa778408b4fc6d228b18294d22cd575e68a1d04 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 15:29:08 +0100 Subject: [PATCH 065/108] rename builder function names --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 24 +++++++++---------- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 23 +++++++++--------- 2 files changed, 23 insertions(+), 24 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 9178022647..7353f14177 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -235,7 +235,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * %tensor = qtensor.alloc(3) : tensor<3x!qco.qubit> * ``` */ - Value allocTensor(int64_t size); + Value qtensorAlloc(int64_t size); /** * @brief Allocate a qubit tensor from a list of qubit values @@ -250,7 +250,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> * ``` */ - Value fromElements(ValueRange elements); + Value qtensorFromElements(ValueRange elements); /** * @brief Extract a qubit from a tensor @@ -266,8 +266,8 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * %q0, %outTensor = qtensor.extract %tensor[%c0]: tensor<3x!qco.qubit> * ``` */ - std::pair extract(Value tensor, - const std::variant& index); + std::pair + qtensorExtract(Value tensor, const std::variant& index); /** * @brief Extract a qubit slice from a tensor @@ -286,8 +286,8 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * ``` */ std::pair - extractSlice(Value tensor, const std::variant& offset, - const std::variant& sizes); + qtensorExtractSlice(Value tensor, const std::variant& offset, + const std::variant& sizes); /** * @brief Insert a qubit into a tensor @@ -304,8 +304,8 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * %outTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> * ``` */ - Value insert(Value scalar, Value tensor, - const std::variant& index); + Value qtensorInsert(Value scalar, Value tensor, + const std::variant& index); /** * @brief Insert a qubit slice into a tensor @@ -324,9 +324,9 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * : tensor<2x!qco.qubit> into tensor<3x!qco.qubit> * ``` */ - Value insertSlice(Value sourceTensor, Value destTensor, - const std::variant& offset, - const std::variant& sizes); + Value qtensorInsertSlice(Value sourceTensor, Value destTensor, + const std::variant& offset, + const std::variant& sizes); /** * @brief Explicitly deallocate a tensor @@ -341,7 +341,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * qtensor.dealloc %tensor : tensor<3x!qco.qubit> * ``` */ - QCOProgramBuilder& deallocTensor(Value tensor); + QCOProgramBuilder& qtensorDealloc(Value tensor); //===--------------------------------------------------------------------===// // Measurement and Reset diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 5e4a66ae22..f42c8d1573 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -194,7 +194,7 @@ void QCOProgramBuilder::updateTensorTracking(Value inputTensor, // QTensor Operations //===----------------------------------------------------------------------===// -Value QCOProgramBuilder::allocTensor(int64_t size) { +Value QCOProgramBuilder::qtensorAlloc(int64_t size) { checkFinalized(); if (size <= 0) { @@ -207,7 +207,7 @@ Value QCOProgramBuilder::allocTensor(int64_t size) { return result; } -Value QCOProgramBuilder::fromElements(ValueRange elements) { +Value QCOProgramBuilder::qtensorFromElements(ValueRange elements) { checkFinalized(); if (elements.empty()) { @@ -229,8 +229,8 @@ Value QCOProgramBuilder::fromElements(ValueRange elements) { } std::pair -QCOProgramBuilder::extract(Value tensor, - const std::variant& index) { +QCOProgramBuilder::qtensorExtract(Value tensor, + const std::variant& index) { checkFinalized(); auto rankedTensorType = llvm::dyn_cast(tensor.getType()); @@ -253,10 +253,9 @@ QCOProgramBuilder::extract(Value tensor, return {qubit, outTensor}; } -std::pair -QCOProgramBuilder::extractSlice(Value tensor, - const std::variant& offset, - const std::variant& sizes) { +std::pair QCOProgramBuilder::qtensorExtractSlice( + Value tensor, const std::variant& offset, + const std::variant& sizes) { checkFinalized(); auto tensorType = llvm::dyn_cast(tensor.getType()); @@ -281,8 +280,8 @@ QCOProgramBuilder::extractSlice(Value tensor, return {slicedTensor, outTensor}; } -Value QCOProgramBuilder::insert(Value scalar, Value tensor, - const std::variant& index) { +Value QCOProgramBuilder::qtensorInsert( + Value scalar, Value tensor, const std::variant& index) { checkFinalized(); auto tensorType = llvm::dyn_cast(tensor.getType()); @@ -305,7 +304,7 @@ Value QCOProgramBuilder::insert(Value scalar, Value tensor, return outTensor; } -Value QCOProgramBuilder::insertSlice( +Value QCOProgramBuilder::qtensorInsertSlice( Value source, Value dest, const std::variant& offset, const std::variant& sizes) { checkFinalized(); @@ -342,7 +341,7 @@ Value QCOProgramBuilder::insertSlice( return outTensor; } -QCOProgramBuilder& QCOProgramBuilder::deallocTensor(Value tensor) { +QCOProgramBuilder& QCOProgramBuilder::qtensorDealloc(Value tensor) { checkFinalized(); auto tensorType = llvm::dyn_cast(tensor.getType()); From 368b69cf44e490464eb20b2983b1d31f057fe08c Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 16:07:21 +0100 Subject: [PATCH 066/108] add sameIndex function to check for identical index values --- .../mlir/Dialect/QTensor/IR/QTensorOps.h | 7 + .../QTensor/IR/Operations/ExtractOp.cpp | 40 +---- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 124 +-------------- .../QTensor/IR/Operations/InsertOp.cpp | 61 +++---- .../QTensor/IR/Operations/InsertSliceOp.cpp | 90 +++-------- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 150 +----------------- 6 files changed, 74 insertions(+), 398 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index ee3bd13a44..0a2017261c 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -12,6 +12,7 @@ // 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" @@ -31,3 +32,9 @@ #define GET_OP_CLASSES #include "mlir/Dialect/QTensor/IR/QTensorOps.h.inc" // IWYU pragma: export + +namespace mlir::qtensor { + +bool isSameIndex(TypedValue index1, TypedValue index2); + +} // namespace mlir::qtensor diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 051f5950b6..281010575e 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -48,24 +48,6 @@ LogicalResult ExtractOp::verify() { return success(); } -struct ExtractFromTensorCast : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ExtractOp extract, - PatternRewriter& rewriter) const final { - auto tensorCast = extract.getTensor().getDefiningOp(); - if (!tensorCast) { - return failure(); - } - if (!llvm::isa(tensorCast.getSource().getType())) { - return failure(); - } - rewriter.replaceOpWithNewOp(extract, tensorCast.getSource(), - extract.getIndex()); - return success(); - } -}; - /** * @brief If an ExtractOp consumes an InsertOp with identical indices, * return the scalar from the InsertOp directly. @@ -83,23 +65,10 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { auto insertIndex = insertOp.getIndex(); auto extractIndex = extractOp.getIndex(); - // Check if SSA values of the indices are the same - if (insertIndex == extractIndex) { + if (isSameIndex(insertIndex, extractIndex)) { return insertOp; } - - auto insertIndexValue = getConstantIntValue(insertIndex); - auto extractIndexValue = getConstantIntValue(extractIndex); - - // Check if the indices are constant and equal - if (!insertIndexValue || !extractIndexValue) { - return nullptr; - } - if (*insertIndexValue != *extractIndexValue) { - return nullptr; - } - - return insertOp; + return nullptr; } LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, @@ -111,8 +80,3 @@ LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, return failure(); } - -void ExtractOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 5770a3729b..3a1201e0c2 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -36,14 +36,10 @@ #include #include -#include using namespace mlir; using namespace mlir::qtensor; -// Adjusted from -// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp - /// Build an ExtractSliceOp with dynamic entries and inferred result type. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value offset, Value size, @@ -106,103 +102,6 @@ LogicalResult ExtractSliceOp::verify() { return success(); } -/** - * @brief Rewrite pattern that pushes tensor.cast past tensor.extract_slice. - * - * @details - * This pattern rewrites a `qtensor.extract_slice` operation whose source - * operand is produced by a `tensor.cast`. When `canFoldIntoConsumerOp` - * evaluates to true, the cast operation is moved after the slice operation. - * - * Conceptually, the slice is applied to the original tensor before the - * cast, avoiding unnecessary intermediate casts. - * - * Example: - * %0 = tensor.cast %V : tensor<3x!qco.qubit> to tensor - * %1, %2 = qtensor.extract_slice %0[0][2][1] - * : tensor to tensor<2x!qco.qubit> - * - * is rewritten into: - * - * %0, %1 = qtensor.extract_slice %V[0][2][1] - * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> - * %2 = tensor.cast %0 : tensor<2x!qco.qubit> to tensor<2x!qco.qubit> - * - * This effectively folds the cast into the consumer operation and enables - * further canonicalization opportunities. - */ -namespace { - -class ExtractSliceOpCastFolder final : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ExtractSliceOp sliceOp, - PatternRewriter& rewriter) const override { - - // Let constant folding handle constant operands - if (matchPattern(sliceOp.getOffset(), matchConstantIndex()) || - matchPattern(sliceOp.getSize(), matchConstantIndex())) { - return failure(); - } - - // Look for tensor.cast producer - auto castOp = sliceOp.getSource().getDefiningOp(); - if (!castOp) { - return failure(); - } - - if (!canFoldIntoConsumerOp(castOp)) { - return failure(); - } - - // Verify bounds using the original tensor - auto srcType = cast(castOp.getSource().getType()); - - int64_t dim = srcType.getShape()[0]; - - auto offsetVal = getConstantIntValue(sliceOp.getOffset()); - auto sizeVal = getConstantIntValue(sliceOp.getSize()); - - if (offsetVal && sizeVal) { - if (*offsetVal + *sizeVal > dim) { - return failure(); - } - } - - Location loc = sliceOp.getLoc(); - - // Create new slice directly on the original tensor - auto newSlice = rewriter.create( - loc, sliceOp.getResult().getType(), sliceOp.getOutSource().getType(), - castOp.getSource(), sliceOp.getOffset(), sliceOp.getSize()); - - Value newResult = newSlice.getResult(); - Value newOutSource = newSlice.getOutSource(); - - // Preserve expected type of out_source - if (newOutSource.getType() != sliceOp.getOutSource().getType()) { - newOutSource = rewriter.create( - loc, sliceOp.getOutSource().getType(), newOutSource); - } - - rewriter.replaceOp(sliceOp, {newResult, newOutSource}); - - if (castOp->use_empty()) { - rewriter.eraseOp(castOp); - } - - return success(); - } -}; - -} // namespace - -void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} - static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { auto insertSliceOp = @@ -221,27 +120,8 @@ foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { auto insertSize = insertSliceOp.getSize(); auto extractSize = extractSliceOp.getSize(); - // Check if SSA values of the offsets and the sizes are the same - if (insertOffset == extractOffset && insertSize == extractSize) { - return insertSliceOp; - } - - auto insertOffsetValue = getConstantIntValue(insertOffset); - auto extractOffsetValue = getConstantIntValue(extractOffset); - auto insertSizeValue = getConstantIntValue(insertSize); - auto extractSizeValue = getConstantIntValue(extractSize); - - // Check if then offsets and sizes are constant and equal - if (!insertOffsetValue || !extractOffsetValue) { - return nullptr; - } - if (!insertSizeValue || !extractSizeValue) { - return nullptr; - } - if (*insertOffsetValue != *extractOffsetValue) { - return nullptr; - } - if (*insertSizeValue != *extractSizeValue) { + if (!isSameIndex(insertOffset, extractOffset) || + !isSameIndex(insertSize, extractSize)) { return nullptr; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 53c671fafe..9bf735014a 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -41,35 +41,6 @@ LogicalResult InsertOp::verify() { return success(); } -namespace { - -/** - * @brief If an InsertOp does not return a tensor with a static shape but the - * destination tensor has one, replace the InsertOp with a new one that has a - * static shape. - */ -struct ConvertInsertOpToStaticShape - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(qtensor::InsertOp insertOp, - PatternRewriter& rewriter) const final { - if (insertOp.getResult().getType().hasStaticShape()) { - return failure(); - } - if (!insertOp.getDest().getType().hasStaticShape()) { - return failure(); - } - - rewriter.replaceOpWithNewOp(insertOp, insertOp.getScalar(), - insertOp.getDest(), - insertOp.getIndex()); - - return success(); - } -}; -} // namespace - /** * @brief If an InsertOp consumes an ExtractOp with identical indices, * return the tensor from the extractOp directly. @@ -106,10 +77,40 @@ static Value foldInsertAfterExtract(InsertOp insertOp) { } OpFoldResult InsertOp::fold(FoldAdaptor /*adaptor*/) { - // Fold after extract_slice if (auto result = foldInsertAfterExtract(*this)) { return result; } return {}; } + +namespace { + +struct InsertAfterInsertOp : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertOp insertOp, + PatternRewriter& rewriter) const final { + auto prevInsertOp = insertOp.getDest().getDefiningOp(); + if (!prevInsertOp) { + return failure(); + } + auto insertIndex = insertOp.getIndex(); + auto prevInsertIndex = prevInsertOp.getIndex(); + + if (!isSameIndex(insertIndex, prevInsertIndex)) { + return failure(); + } + + rewriter.replaceOpWithNewOp(insertOp, insertOp.getScalar(), + prevInsertOp.getDest(), insertIndex); + rewriter.eraseOp(prevInsertOp); + return success(); + } +}; +} // namespace + +void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, + MLIRContext* context) { + results.add(context); +} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 3679b068c1..b2463cd1c9 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -135,8 +135,6 @@ static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { return failure(); } - // Fold: bypass previous insert - insertOp.getDestMutable().assign(prevInsertOp.getDest()); return success(); } @@ -175,27 +173,8 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { auto insertSize = insertSliceOp.getSize(); auto extractSize = extractSliceOp.getSize(); - // Check if SSA values of the offsets and the sizes are the same - if (insertOffset == extractOffset && insertSize == extractSize) { - return extractSliceOp.getSource(); - } - - auto insertOffsetValue = getConstantIntValue(insertOffset); - auto extractOffsetValue = getConstantIntValue(extractOffset); - auto insertSizeValue = getConstantIntValue(insertSize); - auto extractSizeValue = getConstantIntValue(extractSize); - - // Check if then offsets and sizes are constant and equal - if (!insertOffsetValue || !extractOffsetValue) { - return nullptr; - } - if (!insertSizeValue || !extractSizeValue) { - return nullptr; - } - if (*insertOffsetValue != *extractOffsetValue) { - return nullptr; - } - if (*insertSizeValue != *extractSizeValue) { + if (!isSameIndex(insertOffset, extractOffset) || + !isSameIndex(insertSize, extractSize)) { return nullptr; } @@ -236,61 +215,44 @@ OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { } namespace { -/** - * @brief Folds tensor.cast operations with insert_slice. - * - * @details - * If the source or destination tensor of an insert_slice operation is - * produced by a tensor.cast that removes static type information, the - * cast can be folded into the insert_slice operation. - * - * Example: - * - * ```mlir - * %1 = tensor.cast %0 : tensor<3!qco.qubit> to tensor - * %2 = qtensor.insert_slice %1 into ... : tensor into ... - * ``` - * - * This folds into: - * - * ```mlir - * %2 = qtensor.insert_slice %0 into ... : tensor<3!qco.qubit> into ... - * ``` - * - * When folding a cast on the destination tensor, the result of the - * insert_slice operation is cast to preserve the original result type. - */ -struct InsertSliceOpCastFolder final : public OpRewritePattern { +struct InsertSliceAfterInsertSlice final + : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(InsertSliceOp op, + LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp, PatternRewriter& rewriter) const override { - - auto srcCast = op.getSource().getDefiningOp(); - auto dstCast = op.getDest().getDefiningOp(); - - if (!srcCast && !dstCast) { + auto prevInsertOp = insertSliceOp.getDest().getDefiningOp(); + if (!prevInsertOp) { return failure(); } - Value newSrc = - srcCast ? srcCast.getSource() : static_cast(op.getSource()); - Value newDst = - dstCast ? dstCast.getSource() : static_cast(op.getDest()); + // Source types must match + if (prevInsertOp.getSource().getType() != + insertSliceOp.getSource().getType()) { + return failure(); + } - auto newOp = - rewriter.create(op.getLoc(), op.getType(), newSrc, - newDst, op.getOffset(), op.getSize()); + auto prevOffset = prevInsertOp.getOffset(); + auto curOffset = insertSliceOp.getOffset(); + auto prevSize = prevInsertOp.getSize(); + auto curSize = insertSliceOp.getSize(); - rewriter.replaceOp(op, newOp->getResults()); + if (!isSameIndex(prevOffset, curOffset) || + !isSameIndex(prevSize, curSize)) { + return failure(); + } + rewriter.replaceOpWithNewOp( + insertSliceOp, insertSliceOp.getSource(), prevInsertOp.getDest(), + curOffset, curSize); + rewriter.eraseOp(prevInsertOp); return success(); - } + }; }; } // namespace void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add(context); } diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 9f7c2ec4be..94352c7014 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -27,15 +27,8 @@ #include #include #include -#include -#include -#include #include -#include -#include -#include - // The following headers are needed for some template instantiations. // IWYU pragma: begin_keep #include @@ -47,149 +40,18 @@ using namespace mlir::qtensor; namespace mlir::qtensor { -llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, - ArrayRef mixedSizes) { - llvm::SmallBitVector droppedDims(mixedSizes.size()); - int64_t shapePos = static_cast(reducedShape.size()) - 1; - - for (const auto& size : enumerate(llvm::reverse(mixedSizes))) { - size_t idx = mixedSizes.size() - size.index() - 1; - // Rank-reduced dims must have a static unit dimension. - bool isStaticUnitSize = - isa(size.value()) && - llvm::cast(cast(size.value())).getInt() == 1; - - if (shapePos < 0) { - // There are no more dims in the reduced shape. All remaining sizes must - // be rank-reduced dims. - assert(isStaticUnitSize && "expected unit dim"); - droppedDims.set(idx); - continue; - } - - // Dim is preserved if the size is not a static 1. - if (!isStaticUnitSize) { - --shapePos; - continue; - } - - // Dim is preserved if the reduced shape dim is also 1. - if (reducedShape[shapePos] == 1) { - --shapePos; - continue; - } - - // Otherwise: Dim is dropped. - droppedDims.set(idx); +bool isSameIndex(TypedValue index1, TypedValue index2) { + if (index1 == index2) { + return true; } - assert(shapePos < 0 && "dimension mismatch"); - return droppedDims; -} + auto val1 = getConstantIntValue(index1); + auto val2 = getConstantIntValue(index2); -LogicalResult produceSliceErrorMsg(SliceVerificationResult result, - Operation* op, - RankedTensorType expectedType) { - switch (result) { - case SliceVerificationResult::Success: - return success(); - case SliceVerificationResult::RankTooLarge: - return op->emitError("expected rank to be smaller or equal to ") - << "the other rank. "; - case SliceVerificationResult::SizeMismatch: - return op->emitError("expected type to be ") - << expectedType << " or a rank-reduced version. (size mismatch) "; - case SliceVerificationResult::ElemTypeMismatch: - return op->emitError("expected element type to be ") - << expectedType.getElementType(); - default: - llvm_unreachable("unexpected extract_slice op verification result"); - } -} - -LogicalResult -foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, - ShapedType shapedType) { - OpBuilder b(op.getContext()); - for (OpFoldResult opFold : op.getMixedOffsets()) { - if (getConstantIntValue(opFold) != static_cast(0)) { - return failure(); - } - } - // Rank-reducing noops only need to inspect the leading dimensions: - // llvm::zip is appropriate. - auto shape = shapedType.getShape(); - for (auto it : llvm::zip(op.getMixedSizes(), shape)) { - if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) { - return failure(); - } - } - return success(); + return val1 && val2 && *val1 == *val2; } } // namespace mlir::qtensor -//===----------------------------------------------------------------------===// -// Common Canonicalizers and Folders. -//===----------------------------------------------------------------------===// - -static bool foldTensorCastPrecondition(DestinationStyleOpInterface op) { - // 1. InsertSliceOp has its own logic about folding tensor.cast ops. - // 2. Exclude DPS ops that are also LoopLike from this interface as they - // might need special handling of attached regions. - if (isa(op.getOperation()) || - isa(op.getOperation())) { - return false; - } - - return mlir::tensor::hasFoldableTensorCastOperand(op); -} - -namespace { -struct FoldTensorCastProducerOp - : public OpInterfaceRewritePattern { - using OpInterfaceRewritePattern< - DestinationStyleOpInterface>::OpInterfaceRewritePattern; - - LogicalResult matchAndRewrite(DestinationStyleOpInterface op, - PatternRewriter& rewriter) const override { - - // Reject PackOp/UnpackOp (i.e. RelayoutOps) - there are dedicated patterns - // for that instead. - if (!foldTensorCastPrecondition(op) || - isa(*op)) { - return failure(); - } - - SmallVector newResultTypes(op->getResultTypes()); - SmallVector newOperands = - mlir::tensor::getUpdatedOperandsAfterCastOpFolding(op, newResultTypes); - - // Clone op - auto newOp = clone(rewriter, op, newResultTypes, newOperands); - - SmallVector replacements; - replacements.reserve(newOp->getNumResults()); - for (auto [oldResult, newResult] : - llvm::zip(op->getResults(), newOp->getResults())) { - if (newResult.getType() != oldResult.getType()) { - replacements.push_back(tensor::CastOp::create( - rewriter, op->getLoc(), oldResult.getType(), newResult)); - } else { - replacements.push_back(newResult); - } - } - rewriter.replaceOp(op, replacements); - - return success(); - } -}; -} // namespace - -void QTensorDialect::getCanonicalizationPatterns( - RewritePatternSet& results) const { - results.add(getContext()); -} - //===----------------------------------------------------------------------===// // Dialect //===----------------------------------------------------------------------===// From 8e869d391c04e2a2be3dfb673ab8f706680c2154 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 16:07:36 +0100 Subject: [PATCH 067/108] remove more redundant code --- .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 19 ------------------- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 15 +-------------- 2 files changed, 1 insertion(+), 33 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h index 6ebfd5f966..f16d153d1d 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -10,30 +10,11 @@ #pragma once -#include #include #include -#include -#include #define DIALECT_NAME_QTensor "qtensor" -namespace mlir::qtensor { - -/// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or -/// rank-extending tensor.insert_slice op. -llvm::SmallBitVector getDroppedDims(ArrayRef reducedShape, - ArrayRef mixedSizes); - -LogicalResult produceSliceErrorMsg(SliceVerificationResult result, - Operation* op, - RankedTensorType expectedType); - -LogicalResult -foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, - ShapedType shapedType); -} // namespace mlir::qtensor - //===----------------------------------------------------------------------===// // QTensor Dialect //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 9315fd7f03..dfb2d4b157 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -37,8 +37,6 @@ def QTensorDialect : Dialect { In addition, alloc/dealloc operations are added to the dialect to support the bulk allocation and deallocation of tensors with linear types. }]; - let hasCanonicalizer = 1; - let cppNamespace = "::mlir::qtensor"; } @@ -149,7 +147,6 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ let hasFolder = 1; let hasVerifier = 1; - let hasCanonicalizer = 1; } def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ @@ -194,13 +191,11 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ CArg<"ArrayRef", "{}">:$attrs)> ]; - let hasCanonicalizer = 1; let hasFolder = 1; let hasVerifier = 1; } def QTensor_InsertOp : QTensorOp<"insert", [ - DestinationStyleOpInterface, Pure, TypesMatchWith<"result type matches type of dest", "dest", "result", @@ -228,16 +223,12 @@ def QTensor_InsertOp : QTensorOp<"insert", [ $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) }]; - let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } - }]; - let hasVerifier = 1; let hasFolder = 1; + let hasCanonicalizer = 1; } def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ - DestinationStyleOpInterface, Pure, TypesMatchWith<"expected result type to match dest type", "dest", "result", "$_self"> @@ -280,10 +271,6 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ CArg<"ArrayRef", "{}">:$attrs)> ]; - let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } - }]; - let hasCanonicalizer = 1; let hasFolder = 1; let hasVerifier = 1; From 910a076cbc209b5a630ccff53cae809534d68ce8 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 16:26:33 +0100 Subject: [PATCH 068/108] minor fixes --- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 5 +- .../QTensor/IR/Operations/InsertOp.cpp | 14 +-- .../QTensor/IR/Operations/InsertSliceOp.cpp | 111 ++++-------------- 3 files changed, 27 insertions(+), 103 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 3a1201e0c2..10590f9f97 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -40,13 +40,12 @@ using namespace mlir; using namespace mlir::qtensor; -/// Build an ExtractSliceOp with dynamic entries and inferred result type. void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value offset, Value size, ArrayRef attrs) { - auto optionalVal = getConstantIntValue(size); + auto sizeValue = getConstantIntValue(size); auto resultType = RankedTensorType::get( - {optionalVal ? *optionalVal : ShapedType::kDynamic}, + {sizeValue ? *sizeValue : ShapedType::kDynamic}, cast(source.getType()).getElementType()); result.addAttributes(attrs); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 9bf735014a..3e68cd4741 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -57,19 +57,7 @@ static Value foldInsertAfterExtract(InsertOp insertOp) { auto insertIndex = insertOp.getIndex(); auto extractIndex = extractOp.getIndex(); - // Check if SSA values of the indices are the same - if (insertIndex == extractIndex) { - return extractOp.getTensor(); - } - - auto insertIndexValue = getConstantIntValue(insertIndex); - auto extractIndexValue = getConstantIntValue(extractIndex); - - // Check if the indices are constant and equal - if (!insertIndexValue || !extractIndexValue) { - return nullptr; - } - if (*insertIndexValue != *extractIndexValue) { + if (!isSameIndex(insertIndex, extractIndex)) { return nullptr; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index b2463cd1c9..9c0ea320b8 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -88,74 +88,6 @@ LogicalResult InsertSliceOp::verify() { return success(); } -/** - * @brief Folds consecutive InsertSliceOp operations writing to the same slice. - * - * @details - * If two consecutive InsertSliceOp operations write to the same slice, - * the destination of the second InsertSliceOp can be updated to the - * destination of the first one, eliminating the intermediate operation. - * - * Example: - * - * ```mlir - * %0 = qtensor.insert_slice %slice0 into %input[%c0][%c2] - * %1 = qtensor.insert_slice %slice1 into %0[%c0][%c2] - * ``` - * - * This folds into: - * - * ```mlir - * %1 = qtensor.insert_slice %slice1 into %input[%c0][%c2] - * ``` - */ -static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { - // Check if the destination of current insert is another insert - auto prevInsertOp = insertOp.getDest().getDefiningOp(); - if (!prevInsertOp) { - return failure(); - } - - // Check source types - if (prevInsertOp.getSource().getType() != insertOp.getSource().getType()) { - return failure(); - } - - // Check offset and size - auto prevOffsetOpt = getConstantIntValue(prevInsertOp.getOffset()); - auto prevSizeOpt = getConstantIntValue(prevInsertOp.getSize()); - auto curOffsetOpt = getConstantIntValue(insertOp.getOffset()); - auto curSizeOpt = getConstantIntValue(insertOp.getSize()); - - // Only fold if offsets and sizes are constant and identical - if (!prevOffsetOpt || !prevSizeOpt || !curOffsetOpt || !curSizeOpt) { - return failure(); - } - if (*prevOffsetOpt != *curOffsetOpt || *prevSizeOpt != *curSizeOpt) { - return failure(); - } - - return success(); -} - -/** - * @brief Folds round-trip extract/insert slice operation pairs. - * - * @details - * Detects patterns where a slice is extracted from a tensor and then - * inserted back into the same tensor at the same offset and size. - * In such cases, the pair of operations forms a no-op and can - * be folded to the original tensor value. - * - * Example: - * - * ```mlir - * %slicedTensor, outTensor = qtensor.extract_slice %tensor[%c0][%c2] - * %newTensor = qtensor.insert_slice %slicedTensor into %outTensor[%c0][%c2] - * ``` - * - * This can be folded into `%tensor`. - */ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { auto extractSliceOp = insertSliceOp.getSource().getDefiningOp(); @@ -163,7 +95,6 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { return nullptr; } - // Ensure the insert destination is the original source tensor of extract if (extractSliceOp.getOutSource() != insertSliceOp.getDest()) { return nullptr; } @@ -181,30 +112,36 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { return extractSliceOp.getSource(); } -OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { - // Identity fold: full overwrite - if (getSource() == getDest()) { - if (auto constOffset = getConstantIntValue(getOffset())) { - if (*constOffset == 0) { - if (auto constSize = getConstantIntValue(getSize())) { - if (*constSize == getSource().getType().getDimSize(0)) { - return getSource(); - } - } - } - } +static Value foldIdentity(InsertSliceOp insertSliceOp) { + auto offsetValue = getConstantIntValue(insertSliceOp.getOffset()); + auto sizeValue = getConstantIntValue(insertSliceOp.getSize()); + auto source = insertSliceOp.getSource(); + auto dest = insertSliceOp.getDest(); + + if (source != dest) { + return nullptr; } - // Fold nested insert after insert - if (succeeded(foldInsertAfterInsertSlice(*this))) { - return getResult(); + if (!offsetValue || !sizeValue) { + return nullptr; + } + if (*offsetValue != 0) { + return nullptr; + } + if (*sizeValue != source.getType().getDimSize(0)) { + return nullptr; + } + return source; +} + +OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { + if (auto result = foldIdentity(*this)) { + return result; } - // Fold after extract_slice if (auto result = foldInsertAfterExtractSlice(*this)) { return result; } - // Zero-length insert if (auto constSize = getConstantIntValue(getSize())) { if (*constSize == 0) { return getDest(); @@ -247,7 +184,7 @@ struct InsertSliceAfterInsertSlice final curOffset, curSize); rewriter.eraseOp(prevInsertOp); return success(); - }; + } }; } // namespace From 9a727a84ab75badf9708525cfc999892e233f6a7 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 17:02:27 +0100 Subject: [PATCH 069/108] rearrange verifiers and add docstrings --- .../mlir/Dialect/QTensor/IR/QTensorOps.h | 7 ++ .../QTensor/IR/Operations/ExtractOp.cpp | 23 +++-- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 62 ++++++------ .../QTensor/IR/Operations/InsertOp.cpp | 20 ++-- .../QTensor/IR/Operations/InsertSliceOp.cpp | 95 +++++++------------ 5 files changed, 99 insertions(+), 108 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index 0a2017261c..bdf9af7adb 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -35,6 +35,13 @@ namespace mlir::qtensor { +/** + * @brief Check if two values of IndexType identical. + * + * @param index1 The first IndexType value. + * @param index2 The second IndexType value. + * @return True if both values are equal. + */ bool isSameIndex(TypedValue index1, TypedValue index2); } // namespace mlir::qtensor diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 281010575e..e813875833 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -27,21 +27,20 @@ using namespace mlir; using namespace mlir::qtensor; -// Adjusted from -// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp - LogicalResult ExtractOp::verify() { - auto tensorType = llvm::cast(getTensor().getType()); + auto tensorType = getTensor().getType(); + auto tensorDim = getTensor().getType().getDimSize(0); + auto index = getConstantIntValue(getIndex()); + if (!llvm::isa(tensorType.getElementType())) { return emitOpError("Elements of tensor must be of qubit type"); } - auto index = getConstantIntValue(getIndex()); - auto size = getTensor().getType().getDimSize(0); + if (index) { if (index < 0) { return emitOpError("Index must be non-negative"); } - if (index >= size) { + if (index >= tensorDim) { return emitOpError("Index exceeds tensor dimension"); } } @@ -49,8 +48,8 @@ LogicalResult ExtractOp::verify() { } /** - * @brief If an ExtractOp consumes an InsertOp with identical indices, - * return the scalar from the InsertOp directly. + * @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(); @@ -65,10 +64,10 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { auto insertIndex = insertOp.getIndex(); auto extractIndex = extractOp.getIndex(); - if (isSameIndex(insertIndex, extractIndex)) { - return insertOp; + if (!isSameIndex(insertIndex, extractIndex)) { + return nullptr; } - return nullptr; + return insertOp; } LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 10590f9f97..95f8d08298 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -52,55 +52,55 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, build(b, result, {resultType, source.getType()}, source, offset, size); } -/// Verifier for ExtractSliceOp. LogicalResult ExtractSliceOp::verify() { - RankedTensorType sourceType = getSource().getType(); + auto sourceType = getSource().getType(); + auto resultType = getResult().getType(); + auto outSourceType = getOutSource().getType(); + auto srcDim = sourceType.getDimSize(0); + auto constOffset = getConstantIntValue(getOffset()); + auto constSize = getConstantIntValue(getSize()); - // Element type check if (!llvm::isa(sourceType.getElementType())) { return emitOpError("Elements of source tensor must be of qubit type"); } - auto srcDim = sourceType.getDimSize(0); + if (constOffset && *constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } - if (auto constSize = getConstantIntValue(getSize())) { - if (*constSize < 0) { - return emitOpError("Size must be non-negative"); - } + if (constSize && *constSize < 0) { + return emitOpError("Size must be non-negative"); + } - // Check size fits in source - if (!ShapedType::isDynamic(srcDim) && *constSize > srcDim) { - return emitOpError("Size exceeds source dimension"); + if (constOffset && constSize && !ShapedType::isDynamic(srcDim)) { + if (*constOffset + *constSize > srcDim) { + return emitOpError("Offset + Size exceeds source dimension"); } + } - if (auto constOffset = getConstantIntValue(getOffset())) { - if (*constOffset < 0) { - return emitOpError("Offset must be non-negative"); - } + if (resultType.getElementType() != sourceType.getElementType()) { + return emitOpError("Result element type must match source element type"); + } - if (!ShapedType::isDynamic(srcDim) && - *constOffset + *constSize > srcDim) { - return emitOpError("Offset + Size exceeds source dimension"); - } - } - } else if (auto constOffset = getConstantIntValue(getOffset())) { - if (*constOffset < 0) { - return emitOpError("Offset must be non-negative"); - } - if (!ShapedType::isDynamic(srcDim) && *constOffset >= srcDim) { - return emitOpError("Offset out of bounds"); - } + if (outSourceType.getElementType() != sourceType.getElementType()) { + return emitOpError( + "OutSourceTensor element type must match source element type"); } - // Verify result slice type matches source element type - RankedTensorType resultType = getOutSource().getType(); // or getResult() - if (resultType.getElementType() != sourceType.getElementType()) { - return emitOpError("result element type must match source element type"); + if (constSize && !ShapedType::isDynamic(resultType.getDimSize(0))) { + if (resultType.getDimSize(0) != *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 = diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 3e68cd4741..1d2b83ab44 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -22,27 +22,31 @@ using namespace mlir::qtensor; LogicalResult InsertOp::verify() { auto destType = getDest().getType(); + auto dstDim = destType.getDimSize(0); + auto index = getConstantIntValue(getIndex()); + if (!llvm::isa(getScalar().getType())) { return emitOpError("Scalar must be of qubit type"); } + if (!llvm::isa(destType.getElementType())) { return emitOpError("Elements of dest tensor must be of qubit type"); } - auto index = getConstantIntValue(getIndex()); - auto size = destType.getDimSize(0); + if (index) { if (index < 0) { return emitOpError("Index must be non-negative"); } - if (index >= size) { + if (index >= dstDim) { return emitOpError("Index exceeds tensor dimension"); } } + return success(); } /** - * @brief If an InsertOp consumes an ExtractOp with identical indices, + * @brief If an InsertOp consumes an ExtractOp with the same index, * return the tensor from the extractOp directly. */ static Value foldInsertAfterExtract(InsertOp insertOp) { @@ -74,7 +78,10 @@ OpFoldResult InsertOp::fold(FoldAdaptor /*adaptor*/) { namespace { -struct InsertAfterInsertOp : public OpRewritePattern { +/** + * @brief Combine subsequent insert operations with the same index. + */ +struct CombineSubsequentInsertOp : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(InsertOp insertOp, @@ -96,9 +103,10 @@ struct InsertAfterInsertOp : public OpRewritePattern { return success(); } }; + } // namespace void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add(context); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 9c0ea320b8..bd76ed8498 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -37,57 +37,53 @@ using namespace mlir; using namespace mlir::qtensor; -// Adjusted from -// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp - -/// Verifier for InsertSliceOp. LogicalResult InsertSliceOp::verify() { auto sourceType = getSource().getType(); auto destType = getDest().getType(); + auto srcDim = sourceType.getDimSize(0); + auto dstDim = destType.getDimSize(0); + auto constOffset = getConstantIntValue(getOffset()); + auto constSize = getConstantIntValue(getSize()); if (!llvm::isa(sourceType.getElementType())) { - return emitOpError("Elements of source tensor must be of qubit type"); + return emitOpError("Elements of sourceTensor must be of qubit type"); } + if (!llvm::isa(destType.getElementType())) { - return emitOpError("Elements of dest tensor must be of qubit type"); + return emitOpError("Elements of destTensor must be of qubit type"); } - auto dstDim = destType.getDimSize(0); - auto srcDim = sourceType.getDimSize(0); + if (constOffset && *constOffset < 0) { + return emitOpError("Offset must be non-negative"); + } - if (auto constSize = getConstantIntValue(getSize())) { - if (*constSize < 0) { - return emitOpError("Size must be non-negative"); - } + if (constSize && *constSize < 0) { + return emitOpError("Size must be non-negative"); + } - // Check size fits in source - if (!ShapedType::isDynamic(srcDim) && *constSize > srcDim) { - return emitOpError("Size exceeds source dimension"); + if (constSize && !ShapedType::isDynamic(srcDim)) { + if (*constSize != srcDim) { + return emitOpError("Size must match source dimension"); } + } - if (auto constOffset = getConstantIntValue(getOffset())) { - if (*constOffset < 0) { - return emitOpError("Offset must be non-negative"); - } - - // Check slice fits in dest - if (!ShapedType::isDynamic(dstDim) && - *constOffset + *constSize > dstDim) { - return emitOpError("Offset + Size exceeds destination dimension"); - } - } - } else if (auto constOffset = getConstantIntValue(getOffset())) { - if (*constOffset < 0) { - return emitOpError("Offset must be non-negative"); - } - if (!ShapedType::isDynamic(dstDim) && *constOffset >= dstDim) { - return emitOpError("Offset out of bounds"); + if (constOffset && constSize && !ShapedType::isDynamic(dstDim)) { + if (*constOffset + *constSize > dstDim) { + return emitOpError("Offset + Size exceeds destination dimension"); } } + if (getResult().getType() != destType) { + return emitOpError("Result type must match dest type"); + } + 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(); @@ -112,36 +108,12 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { return extractSliceOp.getSource(); } -static Value foldIdentity(InsertSliceOp insertSliceOp) { - auto offsetValue = getConstantIntValue(insertSliceOp.getOffset()); - auto sizeValue = getConstantIntValue(insertSliceOp.getSize()); - auto source = insertSliceOp.getSource(); - auto dest = insertSliceOp.getDest(); - - if (source != dest) { - return nullptr; - } - if (!offsetValue || !sizeValue) { - return nullptr; - } - if (*offsetValue != 0) { - return nullptr; - } - if (*sizeValue != source.getType().getDimSize(0)) { - return nullptr; - } - return source; -} - OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { - if (auto result = foldIdentity(*this)) { - return result; - } - if (auto result = foldInsertAfterExtractSlice(*this)) { return result; } + // Fold InsertSliceOp if size is 0 if (auto constSize = getConstantIntValue(getSize())) { if (*constSize == 0) { return getDest(); @@ -150,9 +122,14 @@ OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { return {}; } + namespace { -struct InsertSliceAfterInsertSlice final +/** + * @brief Combine subsequent insertSlice operations with the same offset and + * size. + */ +struct CombineSubsequentInsertSliceOp final : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -191,5 +168,5 @@ struct InsertSliceAfterInsertSlice final void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results.add(context); } From 754600966ce04fd2795f8f2f4a5ab464081fe1d5 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 17:39:10 +0100 Subject: [PATCH 070/108] clean up headers --- .../Dialect/QTensor/IR/Operations/ExtractOp.cpp | 6 ------ .../QTensor/IR/Operations/ExtractSliceOp.cpp | 14 +------------- .../QTensor/IR/Operations/FromTensorOp.cpp | 4 ++-- .../Dialect/QTensor/IR/Operations/InsertOp.cpp | 2 ++ .../QTensor/IR/Operations/InsertSliceOp.cpp | 15 ++------------- 5 files changed, 7 insertions(+), 34 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index e813875833..3032d1b406 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -11,16 +11,10 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include -#include #include -#include #include -#include #include -#include -#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 95f8d08298..b0769b9cbe 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -12,26 +12,14 @@ #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include -#include -#include #include -#include -#include -#include #include #include #include #include #include -#include -#include #include #include -#include -#include -#include -#include #include #include @@ -84,7 +72,7 @@ LogicalResult ExtractSliceOp::verify() { if (outSourceType.getElementType() != sourceType.getElementType()) { return emitOpError( - "OutSourceTensor element type must match source element type"); + "OutSource tensor element type must match source element type"); } if (constSize && !ShapedType::isDynamic(resultType.getDimSize(0))) { diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 2df24badc5..120cddf693 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -29,7 +29,7 @@ using namespace mlir::qtensor; void FromElementsOp::build(OpBuilder& builder, OperationState& result, ValueRange elements) { - assert(!elements.empty() && "expected at least one element"); + assert(!elements.empty() && "Expected at least one element"); Type resultType = RankedTensorType::get( {static_cast(elements.size())}, elements.front().getType()); build(builder, result, resultType, elements); @@ -37,7 +37,7 @@ void FromElementsOp::build(OpBuilder& builder, OperationState& result, LogicalResult FromElementsOp::verify() { if (!llvm::isa(getResult().getType().getElementType())) { - return emitOpError("result tensor must have qubit element type"); + return emitOpError("Result tensor must have qubit element type"); } for (auto type : getElements().getTypes()) { diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 1d2b83ab44..01856cd2ec 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -12,7 +12,9 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include +#include #include +#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index bd76ed8498..fd1e80f05e 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -12,23 +12,12 @@ #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include -#include -#include #include -#include -#include #include -#include -#include #include -#include -#include #include #include #include -#include -#include #include #include @@ -46,11 +35,11 @@ LogicalResult InsertSliceOp::verify() { auto constSize = getConstantIntValue(getSize()); if (!llvm::isa(sourceType.getElementType())) { - return emitOpError("Elements of sourceTensor must be of qubit type"); + return emitOpError("Elements of source tensor must be of qubit type"); } if (!llvm::isa(destType.getElementType())) { - return emitOpError("Elements of destTensor must be of qubit type"); + return emitOpError("Elements of dest tensor must be of qubit type"); } if (constOffset && *constOffset < 0) { From 4656a8656e4a5aeddb88380f793361b4f84feebd Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 17:40:21 +0100 Subject: [PATCH 071/108] apply changes to tests --- mlir/unittests/programs/qco_programs.cpp | 57 +++++++++++++----------- 1 file changed, 30 insertions(+), 27 deletions(-) diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index d99fa1d42a..071ffea869 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2056,64 +2056,67 @@ void nestedFalseIf(QCOProgramBuilder& b) { }); } -void allocTensor(QCOProgramBuilder& b) { b.allocTensor(3); } +void allocTensor(QCOProgramBuilder& b) { b.qtensorAlloc(3); } void allocDeallocTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - b.deallocTensor(qtensor); + auto qtensor = b.qtensorAlloc(3); + b.qtensorDealloc(qtensor); } void fromElements(QCOProgramBuilder& b) { auto q0 = b.allocQubit(); auto q1 = b.allocQubit(); auto q2 = b.allocQubit(); - b.fromElements({q0, q1, q2}); + b.qtensorFromElements({q0, q1, q2}); } void extractTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - b.extract(qtensor, 0); + auto qtensor = b.qtensorAlloc(3); + b.qtensorExtract(qtensor, 0); } void insertTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - auto [q0, extractOutTensor] = b.extract(qtensor, 0); + auto qtensor = b.qtensorAlloc(3); + auto [q0, extractOutTensor] = b.qtensorExtract(qtensor, 0); auto q1 = b.h(q0); - b.insert(q1, extractOutTensor, 0); + b.qtensorInsert(q1, extractOutTensor, 0); } void extractSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - b.extractSlice(qtensor, 0, 2); + auto qtensor = b.qtensorAlloc(3); + b.qtensorExtractSlice(qtensor, 0, 2); } void insertSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2); - auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); + auto qtensor = b.qtensorAlloc(3); + auto [slicedTensor, extractSliceOutTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + auto [q0, extractOutTensor] = b.qtensorExtract(slicedTensor, 0); auto q1 = b.h(q0); - auto insertOutTensor = b.insert(q1, extractOutTensor, 0); - b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); + auto insertOutTensor = b.qtensorInsert(q1, extractOutTensor, 0); + b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } void extractInsertTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - auto [q0, extractOutTensor] = b.extract(qtensor, 0); - b.insert(q0, extractOutTensor, 0); + auto qtensor = b.qtensorAlloc(3); + auto [q0, extractOutTensor] = b.qtensorExtract(qtensor, 0); + b.qtensorInsert(q0, extractOutTensor, 0); } void extractSliceInsertSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2); - b.insertSlice(slicedTensor, extractSliceOutTensor, 0, 2); + auto qtensor = b.qtensorAlloc(3); + auto [slicedTensor, extractSliceOutTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 0, 2); } void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { - auto qtensor = b.allocTensor(3); - auto [slicedTensor, extractSliceOutTensor] = b.extractSlice(qtensor, 0, 2); - auto [q0, extractOutTensor] = b.extract(slicedTensor, 0); - auto insertOutTensor = b.insert(q0, extractOutTensor, 0); - b.insertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); + auto qtensor = b.qtensorAlloc(3); + auto [slicedTensor, extractSliceOutTensor] = + b.qtensorExtractSlice(qtensor, 0, 2); + auto [q0, extractOutTensor] = b.qtensorExtract(slicedTensor, 0); + auto insertOutTensor = b.qtensorInsert(q0, extractOutTensor, 0); + b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } } // namespace mlir::qco From 61eb09848d986dce82aa09b230548f3b8cf96dfd Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 18:05:59 +0100 Subject: [PATCH 072/108] rearrange tablegen file --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index dfb2d4b157..6bd3920d4e 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -89,8 +89,8 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { let arguments = (ins AnyRankedTensor:$tensor); let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; - let hasVerifier = 1; let hasCanonicalizer = 1; + let hasVerifier = 1; } def QTensor_FromElementsOp : QTensorOp<"from_elements", [ @@ -106,7 +106,9 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ This operation creates a tensor using the operands as its values. Example: - + ```mlir + %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> + ``` }]; let arguments = (ins Variadic:$elements); @@ -132,8 +134,8 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ 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 element from the input tensor - at the given indices. In addition, it also returns the updated input tensor as a result. - The indices must be of `index` type. + at the given index. In addition, it also returns the updated input tensor as a result. + The index must be of `index` type. Example: ```mlir @@ -209,7 +211,6 @@ def QTensor_InsertOp : QTensorOp<"insert", [ This operation inserts a scalar into a ranked tensor as specified by the operation`s index. Example: - ```mlir %newTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> ``` @@ -223,9 +224,9 @@ def QTensor_InsertOp : QTensorOp<"insert", [ $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) }]; - let hasVerifier = 1; - let hasFolder = 1; let hasCanonicalizer = 1; + let hasFolder = 1; + let hasVerifier = 1; } def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ @@ -235,7 +236,7 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ ]> { let summary = "Insert slice into tensor"; let description = [{ - The `qtensor.insert_slice` operation is a minimally adjusted version of the `tensor.insert_slice` operation + 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 argument. From 3c4f5dbf57bcc0f238092c4d22c690810133bcc0 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 18:25:25 +0100 Subject: [PATCH 073/108] removed unused interface --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h | 1 - mlir/lib/Conversion/QCOToQC/CMakeLists.txt | 1 - mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 1 - 3 files changed, 3 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index bdf9af7adb..e7adae0de5 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -18,7 +18,6 @@ #pragma clang diagnostic ignored "-Wambiguous-reversed-operator" #endif #include -#include #ifdef __clang__ #pragma clang diagnostic pop #endif diff --git a/mlir/lib/Conversion/QCOToQC/CMakeLists.txt b/mlir/lib/Conversion/QCOToQC/CMakeLists.txt index 8b85d395bf..da409ad0b2 100644 --- a/mlir/lib/Conversion/QCOToQC/CMakeLists.txt +++ b/mlir/lib/Conversion/QCOToQC/CMakeLists.txt @@ -20,7 +20,6 @@ add_mlir_library( MLIRArithDialect MLIRFuncDialect MLIRTransforms - MLIRViewLikeInterface MLIRFuncTransforms DISABLE_INSTALL) diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt index 203af124cf..17c3149d32 100644 --- a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -20,7 +20,6 @@ add_mlir_dialect_library( PUBLIC MLIRIR MLIRArithDialect - MLIRViewLikeInterface MLIRInferTypeOpInterface MLIRSideEffectInterfaces DISABLE_INSTALL) From 97b8b1b42d30d645abbd8cde8f4caf72bb9a1801 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 18:26:55 +0100 Subject: [PATCH 074/108] reorder results of extract and extractSliceOp --- .../mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 12 ++++++------ mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 8 ++++---- mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 4 ++-- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 2 +- mlir/unittests/programs/qco_programs.cpp | 14 +++++++------- 5 files changed, 20 insertions(+), 20 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 7353f14177..6c568a1ab0 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -256,14 +256,14 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @brief Extract a qubit from a tensor * @param tensor Source tensor * @param index The index from where the qubit is extracted - * @return Pair of (extractedQubit, outTensor) + * @return Pair of (outTensor, extractedQubit) * * @par Example: * ```c++ - * auto [q0, outTensor] = builder.extract(tensor, 0); + * auto [outTensor, q0] = builder.extract(tensor, 0); * ``` * ```mlir - * %q0, %outTensor = qtensor.extract %tensor[%c0]: tensor<3x!qco.qubit> + * %outTensor, %q0 = qtensor.extract %tensor[%c0]: tensor<3x!qco.qubit> * ``` */ std::pair @@ -274,14 +274,14 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @param tensor Source tensor * @param offset The offset from where the slice is extracted * @param size The size of the extracted slice - * @return Pair of (extractedSlice, outTensor) + * @return Pair of (outTensor, extractedSlice) * * @par Example: * ```c++ - * auto [extractedSlice, outTensor] = builder.extractSlice(tensor, 0, 2); + * auto [outTensor, extractedSlice] = builder.extractSlice(tensor, 0, 2); * ``` * ```mlir - * %extractedSlice, %outTensor = qtensor.extract_slice %tensor[%c0][%c2] + * %outTensor, %extractedSlice = qtensor.extract_slice %tensor[%c0][%c2] * : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> * ``` */ diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 6bd3920d4e..43df95337a 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -139,12 +139,12 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ Example: ```mlir - %q0, %outTensor = qtensor.extract %tensor[%c0] : tensor<4x!qco.qubit> + %outTensor, %q0 = qtensor.extract %tensor[%c0] : tensor<4x!qco.qubit> ``` }]; let arguments = (ins AnyRankedTensor:$tensor, Index:$index); - let results = (outs AnyType:$result, AnyRankedTensor:$out_tensor); + let results = (outs AnyRankedTensor:$out_tensor, AnyType:$result); let assemblyFormat = "$tensor `[` $index `]` attr-dict `:` type($tensor)"; let hasFolder = 1; @@ -171,7 +171,7 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ Example: ```mlir - %extractedSlice, %outSourceTensor = qtensor.extract_slice %sourceTensor[%c0][%c2] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + %outSourceTensor, %extractedSlice = qtensor.extract_slice %sourceTensor[%c0][%c2] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> ``` }]; @@ -180,7 +180,7 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ Index:$offset, Index:$size ); - let results = (outs AnyRankedTensor:$result, AnyRankedTensor:$out_source); + let results = (outs AnyRankedTensor:$out_source, AnyRankedTensor:$result); let assemblyFormat = [{ $source `[`$offset `]` `[`$size`]` diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index f42c8d1573..d040da67df 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -250,7 +250,7 @@ QCOProgramBuilder::qtensorExtract(Value tensor, validQubits.insert(qubit); updateTensorTracking(tensor, outTensor); - return {qubit, outTensor}; + return {outTensor, qubit}; } std::pair QCOProgramBuilder::qtensorExtractSlice( @@ -277,7 +277,7 @@ std::pair QCOProgramBuilder::qtensorExtractSlice( validTensors.insert(slicedTensor); updateTensorTracking(tensor, outTensor); - return {slicedTensor, outTensor}; + return {outTensor, slicedTensor}; } Value QCOProgramBuilder::qtensorInsert( diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index b0769b9cbe..c9aa56763a 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -37,7 +37,7 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, cast(source.getType()).getElementType()); result.addAttributes(attrs); - build(b, result, {resultType, source.getType()}, source, offset, size); + build(b, result, {source.getType(), resultType}, source, offset, size); } LogicalResult ExtractSliceOp::verify() { diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 071ffea869..96e7c6e040 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2077,7 +2077,7 @@ void extractTensor(QCOProgramBuilder& b) { void insertTensor(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); - auto [q0, extractOutTensor] = b.qtensorExtract(qtensor, 0); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); auto q1 = b.h(q0); b.qtensorInsert(q1, extractOutTensor, 0); } @@ -2089,9 +2089,9 @@ void extractSliceTensor(QCOProgramBuilder& b) { void insertSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); - auto [slicedTensor, extractSliceOutTensor] = + auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); - auto [q0, extractOutTensor] = b.qtensorExtract(slicedTensor, 0); + 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); @@ -2099,22 +2099,22 @@ void insertSliceTensor(QCOProgramBuilder& b) { void extractInsertTensor(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); - auto [q0, extractOutTensor] = b.qtensorExtract(qtensor, 0); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); b.qtensorInsert(q0, extractOutTensor, 0); } void extractSliceInsertSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); - auto [slicedTensor, extractSliceOutTensor] = + auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 0, 2); } void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); - auto [slicedTensor, extractSliceOutTensor] = + auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); - auto [q0, extractOutTensor] = b.qtensorExtract(slicedTensor, 0); + auto [extractOutTensor, q0] = b.qtensorExtract(slicedTensor, 0); auto insertOutTensor = b.qtensorInsert(q0, extractOutTensor, 0); b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } From 9135e7a2aa6f9a43d36abbb6fe8f2af1dd725780 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 19:13:54 +0100 Subject: [PATCH 075/108] remove redundant builder --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 43df95337a..06005ef92a 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -10,17 +10,11 @@ #define QTENSOROPS include "mlir/IR/OpBase.td" -include "mlir/Interfaces/CastInterfaces.td" -include "mlir/Interfaces/ControlFlowInterfaces.td" -include "mlir/Interfaces/DestinationStyleOpInterface.td" include "mlir/Interfaces/InferIntRangeInterface.td" include "mlir/Interfaces/InferTypeOpInterface.td" -include "mlir/Interfaces/ParallelCombiningOpInterface.td" include "mlir/Interfaces/ShapedOpInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/Interfaces/TilingInterface.td" -include "mlir/Interfaces/ViewLikeInterface.td" -include "mlir/IR/OpAsmInterface.td" //===----------------------------------------------------------------------===// // Dialect @@ -266,12 +260,6 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ attr-dict `:` type($source) `into` type($dest) }]; - let builders = [ - OpBuilder<(ins "Value":$source, - "Value":$offset, "Value":$size, - CArg<"ArrayRef", "{}">:$attrs)> - ]; - let hasCanonicalizer = 1; let hasFolder = 1; let hasVerifier = 1; From 21c88602e7777d0277a488810e5cbe9aebdd825d Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 19:37:26 +0100 Subject: [PATCH 076/108] fix linter issues --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 1 - mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 1 - mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 1 - mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 7 ------- 4 files changed, 10 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 3032d1b406..e927650b32 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -13,7 +13,6 @@ #include #include -#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index c9aa56763a..85b642d01e 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -9,7 +9,6 @@ */ #include "mlir/Dialect/QCO/IR/QCODialect.h" -#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index fd1e80f05e..95a724373b 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -9,7 +9,6 @@ */ #include "mlir/Dialect/QCO/IR/QCODialect.h" -#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index 94352c7014..a0dd0ff5d5 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -12,14 +12,7 @@ #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated -#include -#include -#include -#include -#include -#include #include -#include #include #include #include From ba96c3bda86d134251f164918df474cb69788751 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 20:05:07 +0100 Subject: [PATCH 077/108] fix extractInsert fold --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 5 +++-- mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index e927650b32..2c4d1d2c5c 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -50,7 +50,7 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { return nullptr; } - if (insertOp.getScalar().getType() != extractOp.getType(0)) { + if (insertOp.getScalar().getType() != extractOp.getType(1)) { return nullptr; } @@ -66,8 +66,9 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { if (auto insertOp = foldExtractAfterInsert(*this)) { - results.push_back(insertOp.getScalar()); results.push_back(insertOp.getDest()); + results.push_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 index 85b642d01e..9b0d8eddce 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -97,7 +97,7 @@ foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { } // Source types must match - if (insertSliceOp.getSource().getType() != extractSliceOp.getType(0)) { + if (insertSliceOp.getSource().getType() != extractSliceOp.getType(1)) { return nullptr; } @@ -117,8 +117,8 @@ foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { if (auto insertOp = foldExtractAfterInsertSlice(*this)) { - results.push_back(insertOp.getSource()); results.push_back(insertOp.getDest()); + results.push_back(insertOp.getSource()); return success(); } From fdadc4295f1a4de091f07b9bfe4dd96066591db8 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 20:30:27 +0100 Subject: [PATCH 078/108] add additional tests and rename tests --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 53 +++++++++++-------- mlir/unittests/programs/qco_programs.cpp | 43 +++++++++++---- mlir/unittests/programs/qco_programs.h | 30 ++++++----- 3 files changed, 82 insertions(+), 44 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 21e149a66e..056bdfe91e 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1064,27 +1064,36 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P( QTensorTest, QCOTest, testing::Values( - QCOTestCase{"AllocTensor", MQT_NAMED_BUILDER(allocTensor), - MQT_NAMED_BUILDER(allocTensor)}, - QCOTestCase{"AllocDeallocTensor", MQT_NAMED_BUILDER(allocDeallocTensor), - MQT_NAMED_BUILDER(allocTensor)}, - QCOTestCase{"FromElements", MQT_NAMED_BUILDER(fromElements), - MQT_NAMED_BUILDER(fromElements)}, - QCOTestCase{"ExtractTensor", MQT_NAMED_BUILDER(extractTensor), - MQT_NAMED_BUILDER(extractTensor)}, - QCOTestCase{"InsertTensor", MQT_NAMED_BUILDER(insertTensor), - MQT_NAMED_BUILDER(insertTensor)}, - QCOTestCase{"ExtractSliceTensor", MQT_NAMED_BUILDER(extractSliceTensor), - MQT_NAMED_BUILDER(extractSliceTensor)}, - QCOTestCase{"InsertSliceTensor", MQT_NAMED_BUILDER(insertSliceTensor), - MQT_NAMED_BUILDER(insertSliceTensor)}, - QCOTestCase{"ExtractInsert", MQT_NAMED_BUILDER(extractInsertTensor), - MQT_NAMED_BUILDER(allocTensor)}, - QCOTestCase{"ExtractSliceInsertSlice", - MQT_NAMED_BUILDER(extractSliceInsertSliceTensor), - MQT_NAMED_BUILDER(allocTensor)}, + 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{"QTensorExtractInsert", + MQT_NAMED_BUILDER(qtensorExtractInsert), + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{"QTensorInsertExtract", + MQT_NAMED_BUILDER(qtensorInsertExtract), + MQT_NAMED_BUILDER(qtensorInsert)}, + QCOTestCase{"QTensorExtractSliceInsertSlice", + MQT_NAMED_BUILDER(qtensorExtractSliceInsertSlice), + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{"IQTensornsertSliceExtractSlice", + MQT_NAMED_BUILDER(qtensorInsertSliceExtractSlice), + MQT_NAMED_BUILDER(qtensorInsertSlice)}, QCOTestCase{ - "ExtractSliceExtractInsertInsertSliceTensor", - MQT_NAMED_BUILDER(extractSliceExtractInsertInsertSliceTensor), - MQT_NAMED_BUILDER(allocTensor)})); + "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 53747b14d3..f8f84345cb 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2134,38 +2134,38 @@ void nestedFalseIf(QCOProgramBuilder& b) { }); } -void allocTensor(QCOProgramBuilder& b) { b.qtensorAlloc(3); } +void qtensorAlloc(QCOProgramBuilder& b) { b.qtensorAlloc(3); } -void allocDeallocTensor(QCOProgramBuilder& b) { +void qtensorDealloc(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); b.qtensorDealloc(qtensor); } -void fromElements(QCOProgramBuilder& b) { +void qtensorFromElements(QCOProgramBuilder& b) { auto q0 = b.allocQubit(); auto q1 = b.allocQubit(); auto q2 = b.allocQubit(); b.qtensorFromElements({q0, q1, q2}); } -void extractTensor(QCOProgramBuilder& b) { +void qtensorExtract(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); b.qtensorExtract(qtensor, 0); } -void insertTensor(QCOProgramBuilder& b) { +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 extractSliceTensor(QCOProgramBuilder& b) { +void qtensorExtractSlice(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); b.qtensorExtractSlice(qtensor, 0, 2); } -void insertSliceTensor(QCOProgramBuilder& b) { +void qtensorInsertSlice(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); @@ -2175,20 +2175,20 @@ void insertSliceTensor(QCOProgramBuilder& b) { b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } -void extractInsertTensor(QCOProgramBuilder& b) { +void qtensorExtractInsert(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); b.qtensorInsert(q0, extractOutTensor, 0); } -void extractSliceInsertSliceTensor(QCOProgramBuilder& b) { +void qtensorExtractSliceInsertSlice(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 0, 2); } -void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { +void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); @@ -2197,4 +2197,27 @@ void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b) { b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } +void qtensorInsertExtract(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 qtensorInsertSliceExtractSlice(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); +} + } // namespace mlir::qco diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 7b93956439..901bc1f5c7 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -976,34 +976,40 @@ void nestedTrueIf(QCOProgramBuilder& b); void nestedFalseIf(QCOProgramBuilder& b); /// Allocates a tensor of size `3`. -void allocTensor(QCOProgramBuilder& b); +void qtensorAlloc(QCOProgramBuilder& b); /// Allocates and explicitly deallocates a tensor. -void allocDeallocTensor(QCOProgramBuilder& b); +void qtensorDealloc(QCOProgramBuilder& b); /// Constructs a tensor with from_elements. -void fromElements(QCOProgramBuilder& b); +void qtensorFromElements(QCOProgramBuilder& b); /// Extracts a qubit from a tensor. -void extractTensor(QCOProgramBuilder& b); +void qtensorExtract(QCOProgramBuilder& b); -/// Inserts a qubit in a tensor. -void insertTensor(QCOProgramBuilder& b); +/// Inserts a qubit into a tensor. +void qtensorInsert(QCOProgramBuilder& b); /// Extracts a slice from a tensor. -void extractSliceTensor(QCOProgramBuilder& b); +void qtensorExtractSlice(QCOProgramBuilder& b); -/// Inserts a slice from a tensor. -void insertSliceTensor(QCOProgramBuilder& b); +/// Inserts a slice into a tensor. +void qtensorInsertSlice(QCOProgramBuilder& b); /// Extracts a qubit from a tensor and insert it immediately. -void extractInsertTensor(QCOProgramBuilder& b); +void qtensorExtractInsert(QCOProgramBuilder& b); + +/// Inserts a qubit into a tensor and extract it immediately. +void qtensorInsertExtract(QCOProgramBuilder& b); /// Extracts a slice of qubits from a tensor and insert it immediately. -void extractSliceInsertSliceTensor(QCOProgramBuilder& b); +void qtensorExtractSliceInsertSlice(QCOProgramBuilder& b); + +/// Inserts a slice of qubits into a tensor and extract it immediately. +void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b); /// Extracts a slice of qubits, a qubit from the slice, insert the qubit back to /// the slice and the slice back to the tensor immediately. -void extractSliceExtractInsertInsertSliceTensor(QCOProgramBuilder& b); +void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b); } // namespace mlir::qco From aa532e8e5c6e5f31936781cdfde01472e43ce6c0 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Thu, 12 Mar 2026 20:39:34 +0100 Subject: [PATCH 079/108] fix linter issues --- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index a0dd0ff5d5..b1bee05cc7 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -14,13 +14,10 @@ #include #include -#include #include #include -#include -#include #include -#include +#include // The following headers are needed for some template instantiations. // IWYU pragma: begin_keep From db886ad51b7affe6ec3f18636167314a3e3a6796 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Fri, 13 Mar 2026 11:29:57 +0100 Subject: [PATCH 080/108] clean up code a bit --- .../Dialect/QTensor/IR/Operations/DeallocOp.cpp | 1 - .../Dialect/QTensor/IR/Operations/ExtractOp.cpp | 3 ++- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 6 ++---- .../Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 4 ---- .../lib/Dialect/QTensor/IR/Operations/InsertOp.cpp | 10 ++++++++-- .../QTensor/IR/Operations/InsertSliceOp.cpp | 14 ++------------ 6 files changed, 14 insertions(+), 24 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp index 17c0775b4b..dab02ba466 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp @@ -13,7 +13,6 @@ #include #include -#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 2c4d1d2c5c..042b020ae7 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -33,7 +34,7 @@ LogicalResult ExtractOp::verify() { if (index < 0) { return emitOpError("Index must be non-negative"); } - if (index >= tensorDim) { + if (!ShapedType::isDynamic(tensorDim) && index >= tensorDim) { return emitOpError("Index exceeds tensor dimension"); } } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 9b0d8eddce..497a6a9b56 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -22,8 +22,6 @@ #include #include -#include - using namespace mlir; using namespace mlir::qtensor; @@ -55,8 +53,8 @@ LogicalResult ExtractSliceOp::verify() { return emitOpError("Offset must be non-negative"); } - if (constSize && *constSize < 0) { - return emitOpError("Size must be non-negative"); + if (constSize && *constSize <= 0) { + return emitOpError("Size must be positive"); } if (constOffset && constSize && !ShapedType::isDynamic(srcDim)) { diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 120cddf693..4f8efd0256 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -14,7 +14,6 @@ #include #include #include -#include #include #include @@ -24,9 +23,6 @@ using namespace mlir; using namespace mlir::qtensor; -// Adjusted from -// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.0/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp - void FromElementsOp::build(OpBuilder& builder, OperationState& result, ValueRange elements) { assert(!elements.empty() && "Expected at least one element"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 01856cd2ec..5b7ba06c95 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -36,10 +36,10 @@ LogicalResult InsertOp::verify() { } if (index) { - if (index < 0) { + if (*index < 0) { return emitOpError("Index must be non-negative"); } - if (index >= dstDim) { + if (!ShapedType::isDynamic(dstDim) && *index >= dstDim) { return emitOpError("Index exceeds tensor dimension"); } } @@ -53,12 +53,16 @@ LogicalResult InsertOp::verify() { */ static Value foldInsertAfterExtract(InsertOp insertOp) { auto extractOp = insertOp.getScalar().getDefiningOp(); + if (!extractOp) { return nullptr; } if (insertOp.getDest() != extractOp.getOutTensor()) { return nullptr; } + if (extractOp.getTensor().getType() != insertOp.getDest().getType()) { + return nullptr; + } auto insertIndex = insertOp.getIndex(); auto extractIndex = extractOp.getIndex(); @@ -89,9 +93,11 @@ struct CombineSubsequentInsertOp : public OpRewritePattern { LogicalResult matchAndRewrite(InsertOp insertOp, PatternRewriter& rewriter) const final { auto prevInsertOp = insertOp.getDest().getDefiningOp(); + if (!prevInsertOp) { return failure(); } + auto insertIndex = insertOp.getIndex(); auto prevInsertIndex = prevInsertOp.getIndex(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 95a724373b..da302abe30 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -20,8 +20,6 @@ #include #include -#include - using namespace mlir; using namespace mlir::qtensor; @@ -45,8 +43,8 @@ LogicalResult InsertSliceOp::verify() { return emitOpError("Offset must be non-negative"); } - if (constSize && *constSize < 0) { - return emitOpError("Size must be non-negative"); + if (constSize && *constSize <= 0) { + return emitOpError("Size must be positive"); } if (constSize && !ShapedType::isDynamic(srcDim)) { @@ -101,13 +99,6 @@ OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { return result; } - // Fold InsertSliceOp if size is 0 - if (auto constSize = getConstantIntValue(getSize())) { - if (*constSize == 0) { - return getDest(); - } - } - return {}; } @@ -129,7 +120,6 @@ struct CombineSubsequentInsertSliceOp final return failure(); } - // Source types must match if (prevInsertOp.getSource().getType() != insertSliceOp.getSource().getType()) { return failure(); From 3d8d1c7d67caa1c3df10224289957519d014b753 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Fri, 13 Mar 2026 11:35:57 +0100 Subject: [PATCH 081/108] fix linter issues --- mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 4f8efd0256..ccfd036ce5 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include From fc615a4591ff3cce5ea33811d0bf50b109061145 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Fri, 13 Mar 2026 13:10:01 +0100 Subject: [PATCH 082/108] address coderabbit comments --- .../mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 15 ++++++++------- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h | 4 +--- .../lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 3 ++- .../Dialect/QTensor/IR/Operations/ExtractOp.cpp | 4 ++++ .../QTensor/IR/Operations/ExtractSliceOp.cpp | 8 +++++++- .../Dialect/QTensor/IR/Operations/InsertOp.cpp | 4 ++++ .../QTensor/IR/Operations/InsertSliceOp.cpp | 4 ++++ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 2 +- 8 files changed, 31 insertions(+), 13 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 6c568a1ab0..3f90fd3ee3 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -229,7 +229,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto tensor = builder.allocTensor(3); + * auto tensor = builder.qtensorAlloc(3); * ``` * ```mlir * %tensor = qtensor.alloc(3) : tensor<3x!qco.qubit> @@ -244,7 +244,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto tensor = builder.fromElements({q0, q1, q2}); + * auto tensor = builder.qtensorFromElements({q0, q1, q2}); * ``` * ```mlir * %tensor = qtensor.from_elements %q0, %q1, %q2 : tensor<3x!qco.qubit> @@ -260,7 +260,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto [outTensor, q0] = builder.extract(tensor, 0); + * auto [outTensor, q0] = builder.qtensorExtract(tensor, 0); * ``` * ```mlir * %outTensor, %q0 = qtensor.extract %tensor[%c0]: tensor<3x!qco.qubit> @@ -278,7 +278,8 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto [outTensor, extractedSlice] = builder.extractSlice(tensor, 0, 2); + * auto [outTensor, extractedSlice] = builder.qtensorExtractSlice(tensor, 0, + * 2); * ``` * ```mlir * %outTensor, %extractedSlice = qtensor.extract_slice %tensor[%c0][%c2] @@ -298,7 +299,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto outTensor = builder.insert(q0, tensor, 0); + * auto outTensor = builder.qtensorInsert(q0, tensor, 0); * ``` * ```mlir * %outTensor = qtensor.insert %q0 into %tensor[%c0] : tensor<3x!qco.qubit> @@ -317,7 +318,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * auto outTensor = builder.insertSlice(slicedTensor, tensor, 0, 2); + * auto outTensor = builder.qtensorInsertSlice(slicedTensor, tensor, 0, 2); * ``` * ```mlir * %outTensor = qtensor.insert_slice %slicedTensor into %tensor[%c0][%c2] @@ -335,7 +336,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * builder.deallocTensor(tensor); + * builder.qtensorDealloc(tensor); * ``` * ```mlir * qtensor.dealloc %tensor : tensor<3x!qco.qubit> diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index e7adae0de5..7d4c6dff75 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -27,15 +27,13 @@ #include #include -#include - #define GET_OP_CLASSES #include "mlir/Dialect/QTensor/IR/QTensorOps.h.inc" // IWYU pragma: export namespace mlir::qtensor { /** - * @brief Check if two values of IndexType identical. + * @brief Check if two values of IndexType are identical. * * @param index1 The first IndexType value. * @param index2 The second IndexType value. diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index d040da67df..a3713ac232 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -12,6 +12,7 @@ #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" @@ -45,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() { diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 042b020ae7..f3574b3a92 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -23,6 +23,10 @@ using namespace mlir::qtensor; LogicalResult ExtractOp::verify() { auto tensorType = getTensor().getType(); + if (tensorType.getRank() != 1) { + return emitOpError("Tensor must be 1-D"); + } + auto tensorDim = getTensor().getType().getDimSize(0); auto index = getConstantIntValue(getIndex()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 497a6a9b56..1a5d9e162c 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -28,10 +28,11 @@ using namespace mlir::qtensor; void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, Value offset, Value size, ArrayRef attrs) { + auto sourceType = cast(source.getType()); auto sizeValue = getConstantIntValue(size); auto resultType = RankedTensorType::get( {sizeValue ? *sizeValue : ShapedType::kDynamic}, - cast(source.getType()).getElementType()); + sourceType.getElementType(), sourceType.getEncoding()); result.addAttributes(attrs); build(b, result, {source.getType(), resultType}, source, offset, size); @@ -41,6 +42,11 @@ LogicalResult ExtractSliceOp::verify() { auto sourceType = getSource().getType(); auto resultType = getResult().getType(); auto outSourceType = getOutSource().getType(); + if (sourceType.getRank() != 1 || resultType.getRank() != 1 || + outSourceType.getRank() != 1) { + return emitOpError("Tensors must be 1-D"); + } + auto srcDim = sourceType.getDimSize(0); auto constOffset = getConstantIntValue(getOffset()); auto constSize = getConstantIntValue(getSize()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 5b7ba06c95..f0ed6fb996 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -24,6 +24,10 @@ using namespace mlir::qtensor; LogicalResult InsertOp::verify() { auto destType = getDest().getType(); + if (destType.getRank() != 1) { + return emitOpError("Dest tensor must be 1-D"); + } + auto dstDim = destType.getDimSize(0); auto index = getConstantIntValue(getIndex()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index da302abe30..4d316237ae 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -26,6 +26,10 @@ using namespace mlir::qtensor; LogicalResult InsertSliceOp::verify() { auto sourceType = getSource().getType(); auto destType = getDest().getType(); + if (sourceType.getRank() != 1 || destType.getRank() != 1) { + return emitOpError("Tensors must be 1-D"); + } + auto srcDim = sourceType.getDimSize(0); auto dstDim = destType.getDimSize(0); auto constOffset = getConstantIntValue(getOffset()); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 056bdfe91e..99229be742 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1089,7 +1089,7 @@ INSTANTIATE_TEST_SUITE_P( QCOTestCase{"QTensorExtractSliceInsertSlice", MQT_NAMED_BUILDER(qtensorExtractSliceInsertSlice), MQT_NAMED_BUILDER(qtensorAlloc)}, - QCOTestCase{"IQTensornsertSliceExtractSlice", + QCOTestCase{"QTensorInsertSliceExtractSlice", MQT_NAMED_BUILDER(qtensorInsertSliceExtractSlice), MQT_NAMED_BUILDER(qtensorInsertSlice)}, QCOTestCase{ From 6195f23b378d84628df58d46c684f087e1343c91 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Fri, 13 Mar 2026 13:44:32 +0100 Subject: [PATCH 083/108] address more coderabbit comments --- .../mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 15 --------------- .../lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 5 ----- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 5 ++--- 3 files changed, 2 insertions(+), 23 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 3f90fd3ee3..d0b92b188f 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -98,21 +98,6 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { */ Value intConstant(int64_t value); - /** - * @brief Create a constant index value - * @param value The value to store in the constant - * @return The value produced by the constant operation - * - * @par Example: - * ```c++ - * auto c = builder.indexConstant(1); - * ``` - * ```mlir - * %c = arith.constant 1 : index - * ``` - */ - Value indexConstant(int64_t value); - //===--------------------------------------------------------------------===// // Memory Management //===--------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index a3713ac232..226476dc20 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -71,11 +71,6 @@ Value QCOProgramBuilder::intConstant(const int64_t value) { return arith::ConstantOp::create(*this, getI64IntegerAttr(value)).getResult(); } -Value QCOProgramBuilder::indexConstant(const int64_t value) { - checkFinalized(); - return arith::ConstantOp::create(*this, getIndexAttr(value)).getResult(); -} - Value QCOProgramBuilder::allocQubit() { checkFinalized(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 1a5d9e162c..0e5773427b 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -73,9 +73,8 @@ LogicalResult ExtractSliceOp::verify() { return emitOpError("Result element type must match source element type"); } - if (outSourceType.getElementType() != sourceType.getElementType()) { - return emitOpError( - "OutSource tensor element type must match source element type"); + if (outSourceType != sourceType) { + return emitOpError("Outsource tensor type must match source tensor type"); } if (constSize && !ShapedType::isDynamic(resultType.getDimSize(0))) { From a343e0fbf52503f74877cdf9b94831ceee96bf0e Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Fri, 13 Mar 2026 14:02:53 +0100 Subject: [PATCH 084/108] fix typo --- mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 4 ++-- mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index d0b92b188f..3a43d38eab 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -273,7 +273,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { */ std::pair qtensorExtractSlice(Value tensor, const std::variant& offset, - const std::variant& sizes); + const std::variant& size); /** * @brief Insert a qubit into a tensor @@ -312,7 +312,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { */ Value qtensorInsertSlice(Value sourceTensor, Value destTensor, const std::variant& offset, - const std::variant& sizes); + const std::variant& size); /** * @brief Explicitly deallocate a tensor diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 226476dc20..62fdd0f17a 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -251,7 +251,7 @@ QCOProgramBuilder::qtensorExtract(Value tensor, std::pair QCOProgramBuilder::qtensorExtractSlice( Value tensor, const std::variant& offset, - const std::variant& sizes) { + const std::variant& size) { checkFinalized(); auto tensorType = llvm::dyn_cast(tensor.getType()); @@ -264,7 +264,7 @@ std::pair QCOProgramBuilder::qtensorExtractSlice( } auto offsetValue = utils::variantToValue(*this, getLoc(), offset); - auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); + auto sizesValue = utils::variantToValue(*this, getLoc(), size); auto extractSliceOp = qtensor::ExtractSliceOp::create(*this, tensor, offsetValue, sizesValue); auto slicedTensor = extractSliceOp.getResult(); @@ -302,7 +302,7 @@ Value QCOProgramBuilder::qtensorInsert( Value QCOProgramBuilder::qtensorInsertSlice( Value source, Value dest, const std::variant& offset, - const std::variant& sizes) { + const std::variant& size) { checkFinalized(); auto sourceTensorType = llvm::dyn_cast(source.getType()); @@ -324,7 +324,7 @@ Value QCOProgramBuilder::qtensorInsertSlice( } auto offsetValue = utils::variantToValue(*this, getLoc(), offset); - auto sizesValue = utils::variantToValue(*this, getLoc(), sizes); + auto sizesValue = utils::variantToValue(*this, getLoc(), size); auto insertSliceOp = qtensor::InsertSliceOp::create(*this, source, dest, offsetValue, sizesValue); From 0c050312fc687b1506e2a26ce6b3b1a5697a3361 Mon Sep 17 00:00:00 2001 From: li-mingbao <74404929+li-mingbao@users.noreply.github.com> Date: Sun, 15 Mar 2026 20:11:03 +0100 Subject: [PATCH 085/108] Apply suggestions from code review Co-authored-by: Daniel Haag <121057143+denialhaag@users.noreply.github.com> Signed-off-by: li-mingbao <74404929+li-mingbao@users.noreply.github.com> --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 06005ef92a..1c16abf276 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -202,7 +202,7 @@ def QTensor_InsertOp : QTensorOp<"insert", [ 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 scalar into a ranked tensor as specified by the operation`s index. + This operation inserts a scalar into a ranked tensor as specified by the operation's index. Example: ```mlir @@ -232,7 +232,7 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ 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 argument. + operation's offset and size arguments. The insert_slice operation supports the following arguments: From ff9f27803fa3ae23c778a65920d5e5273e171596 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sun, 15 Mar 2026 20:38:49 +0100 Subject: [PATCH 086/108] simplify code and use 1D tensor as type --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 36 +++++++++--------- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 2 +- .../Dialect/QTensor/IR/Operations/AllocOp.cpp | 4 -- .../QTensor/IR/Operations/ExtractOp.cpp | 10 ++--- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 38 ++++++++----------- .../QTensor/IR/Operations/FromTensorOp.cpp | 2 +- .../QTensor/IR/Operations/InsertOp.cpp | 3 -- .../QTensor/IR/Operations/InsertSliceOp.cpp | 7 +--- 8 files changed, 40 insertions(+), 62 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 1c16abf276..cb89b444a0 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -59,7 +59,7 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ }]; let arguments = (ins ConfinedAttr:$size); - let results = (outs AnyStaticShapeTensor:$result); + let results = (outs 1DTensorOf<[AnyType]>:$result); let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; let builders = [ @@ -80,7 +80,7 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { ``` }]; - let arguments = (ins AnyRankedTensor:$tensor); + let arguments = (ins 1DTensorOf<[AnyType]>:$tensor); let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; let hasCanonicalizer = 1; @@ -106,7 +106,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ }]; let arguments = (ins Variadic:$elements); - let results = (outs AnyStaticShapeTensor:$result); + let results = (outs 1DTensorOf<[AnyType]>:$result); let assemblyFormat = "$elements attr-dict `:` type($result)"; let builders = [ @@ -137,8 +137,8 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ ``` }]; - let arguments = (ins AnyRankedTensor:$tensor, Index:$index); - let results = (outs AnyRankedTensor:$out_tensor, AnyType:$result); + let arguments = (ins 1DTensorOf<[AnyType]>:$tensor, Index:$index); + let results = (outs 1DTensorOf<[AnyType]>:$out_tensor, AnyType:$result); let assemblyFormat = "$tensor `[` $index `]` attr-dict `:` type($tensor)"; let hasFolder = 1; @@ -148,7 +148,7 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ Pure, TypesMatchWith<"returned tensor type matches input tensor", - "source", "out_source", "$_self"> + "tensor", "out_tensor", "$_self"> ]> { let summary = "Extract slice from tensor"; let description = [{ @@ -158,31 +158,31 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ The extract_slice operation supports the following arguments: - * source: the "base" tensor from which to extract a slice. + * 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 - %outSourceTensor, %extractedSlice = qtensor.extract_slice %sourceTensor[%c0][%c2] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> + %outTensor, %extractedSlice = qtensor.extract_slice %tensor[%c0][%c2] : tensor<3x!qco.qubit> to tensor<2x!qco.qubit> ``` }]; let arguments = (ins - AnyRankedTensor:$source, + 1DTensorOf<[AnyType]>:$tensor, Index:$offset, Index:$size ); - let results = (outs AnyRankedTensor:$out_source, AnyRankedTensor:$result); + let results = (outs 1DTensorOf<[AnyType]>:$out_tensor, 1DTensorOf<[AnyType]>:$result); let assemblyFormat = [{ - $source `[`$offset `]` `[`$size`]` - attr-dict `:` type($source) `to` type($result) + $tensor `[`$offset `]` `[`$size`]` + attr-dict `:` type($tensor) `to` type($result) }]; let builders = [ - OpBuilder<(ins "Value":$source, "Value":$offset, + OpBuilder<(ins "Value":$tensor, "Value":$offset, "Value":$size, CArg<"ArrayRef", "{}">:$attrs)> ]; @@ -211,9 +211,9 @@ def QTensor_InsertOp : QTensorOp<"insert", [ }]; let arguments = (ins AnyType:$scalar, - AnyRankedTensor:$dest, + 1DTensorOf<[AnyType]>:$dest, Index:$index); - let results = (outs AnyRankedTensor:$result); + let results = (outs 1DTensorOf<[AnyType]>:$result); let assemblyFormat = [{ $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) }]; @@ -249,12 +249,12 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ }]; let arguments = (ins - AnyRankedTensor:$source, - AnyRankedTensor:$dest, + 1DTensorOf<[AnyType]>:$source, + 1DTensorOf<[AnyType]>:$dest, Index:$offset, Index:$size ); - let results = (outs AnyRankedTensor:$result); + let results = (outs 1DTensorOf<[AnyType]>:$result); let assemblyFormat = [{ $source `into` $dest `[`$offset `]` `[`$size`]` attr-dict `:` type($source) `into` type($dest) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 62fdd0f17a..2e7c3ad4f8 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -268,7 +268,7 @@ std::pair QCOProgramBuilder::qtensorExtractSlice( auto extractSliceOp = qtensor::ExtractSliceOp::create(*this, tensor, offsetValue, sizesValue); auto slicedTensor = extractSliceOp.getResult(); - auto outTensor = extractSliceOp.getOutSource(); + auto outTensor = extractSliceOp.getOutTensor(); validTensors.insert(slicedTensor); updateTensorTracking(tensor, outTensor); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index 307915ff00..0a386abb6e 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -42,10 +42,6 @@ LogicalResult AllocOp::verify() { auto size = static_cast(getSize()); - if (resultType.getRank() != 1) { - return emitOpError("Result must be a 1-D tensor"); - } - if (resultType.getShape()[0] != size) { return emitOpError("Tensor length must match size attribute (") << size << "), but got " << resultType.getShape()[0]; diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index f3574b3a92..8ef85e9062 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -23,10 +23,6 @@ using namespace mlir::qtensor; LogicalResult ExtractOp::verify() { auto tensorType = getTensor().getType(); - if (tensorType.getRank() != 1) { - return emitOpError("Tensor must be 1-D"); - } - auto tensorDim = getTensor().getType().getDimSize(0); auto index = getConstantIntValue(getIndex()); @@ -35,10 +31,10 @@ LogicalResult ExtractOp::verify() { } if (index) { - if (index < 0) { + if (*index < 0) { return emitOpError("Index must be non-negative"); } - if (!ShapedType::isDynamic(tensorDim) && index >= tensorDim) { + if (!ShapedType::isDynamic(tensorDim) && *index >= tensorDim) { return emitOpError("Index exceeds tensor dimension"); } } @@ -55,7 +51,7 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { return nullptr; } - if (insertOp.getScalar().getType() != extractOp.getType(1)) { + if (insertOp.getScalar().getType() != extractOp.getResult().getType()) { return nullptr; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 0e5773427b..331f7bbdb9 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -25,33 +25,28 @@ using namespace mlir; using namespace mlir::qtensor; -void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value source, +void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value tensor, Value offset, Value size, ArrayRef attrs) { - auto sourceType = cast(source.getType()); + auto tensorType = cast(tensor.getType()); auto sizeValue = getConstantIntValue(size); auto resultType = RankedTensorType::get( {sizeValue ? *sizeValue : ShapedType::kDynamic}, - sourceType.getElementType(), sourceType.getEncoding()); + tensorType.getElementType(), tensorType.getEncoding()); result.addAttributes(attrs); - build(b, result, {source.getType(), resultType}, source, offset, size); + build(b, result, {tensor.getType(), resultType}, tensor, offset, size); } LogicalResult ExtractSliceOp::verify() { - auto sourceType = getSource().getType(); + auto tensorType = getTensor().getType(); auto resultType = getResult().getType(); - auto outSourceType = getOutSource().getType(); - if (sourceType.getRank() != 1 || resultType.getRank() != 1 || - outSourceType.getRank() != 1) { - return emitOpError("Tensors must be 1-D"); - } - auto srcDim = sourceType.getDimSize(0); + auto tensorDim = tensorType.getDimSize(0); auto constOffset = getConstantIntValue(getOffset()); auto constSize = getConstantIntValue(getSize()); - if (!llvm::isa(sourceType.getElementType())) { + if (!llvm::isa(tensorType.getElementType())) { return emitOpError("Elements of source tensor must be of qubit type"); } @@ -63,18 +58,15 @@ LogicalResult ExtractSliceOp::verify() { return emitOpError("Size must be positive"); } - if (constOffset && constSize && !ShapedType::isDynamic(srcDim)) { - if (*constOffset + *constSize > srcDim) { + if (constOffset && constSize && !ShapedType::isDynamic(tensorDim)) { + if (*constOffset + *constSize > tensorDim) { return emitOpError("Offset + Size exceeds source dimension"); } } - if (resultType.getElementType() != sourceType.getElementType()) { - return emitOpError("Result element type must match source element type"); - } - - if (outSourceType != sourceType) { - return emitOpError("Outsource tensor type must match source tensor type"); + if (resultType.getElementType() != tensorType.getElementType()) { + return emitOpError( + "Result element type must match input tensor element type"); } if (constSize && !ShapedType::isDynamic(resultType.getDimSize(0))) { @@ -94,13 +86,13 @@ LogicalResult ExtractSliceOp::verify() { static InsertSliceOp foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { auto insertSliceOp = - extractSliceOp.getSource().getDefiningOp(); + extractSliceOp.getTensor().getDefiningOp(); if (!insertSliceOp) { return nullptr; } - // Source types must match - if (insertSliceOp.getSource().getType() != extractSliceOp.getType(1)) { + if (insertSliceOp.getSource().getType() != + extractSliceOp.getResult().getType()) { return nullptr; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index ccfd036ce5..7288122063 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -27,7 +27,7 @@ using namespace mlir::qtensor; void FromElementsOp::build(OpBuilder& builder, OperationState& result, ValueRange elements) { assert(!elements.empty() && "Expected at least one element"); - Type resultType = RankedTensorType::get( + 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 index f0ed6fb996..5e4cf95c89 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -24,9 +24,6 @@ using namespace mlir::qtensor; LogicalResult InsertOp::verify() { auto destType = getDest().getType(); - if (destType.getRank() != 1) { - return emitOpError("Dest tensor must be 1-D"); - } auto dstDim = destType.getDimSize(0); auto index = getConstantIntValue(getIndex()); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 4d316237ae..86353a8412 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -26,9 +26,6 @@ using namespace mlir::qtensor; LogicalResult InsertSliceOp::verify() { auto sourceType = getSource().getType(); auto destType = getDest().getType(); - if (sourceType.getRank() != 1 || destType.getRank() != 1) { - return emitOpError("Tensors must be 1-D"); - } auto srcDim = sourceType.getDimSize(0); auto dstDim = destType.getDimSize(0); @@ -81,7 +78,7 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { return nullptr; } - if (extractSliceOp.getOutSource() != insertSliceOp.getDest()) { + if (extractSliceOp.getOutTensor() != insertSliceOp.getDest()) { return nullptr; } @@ -95,7 +92,7 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { return nullptr; } - return extractSliceOp.getSource(); + return extractSliceOp.getTensor(); } OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { From bb6d65c27001ffc5069f475a59d44dad17fbc0c2 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Sun, 15 Mar 2026 21:57:02 +0100 Subject: [PATCH 087/108] add qubit type constraint in tablegen to simplify verifiers --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 41 +++++++++++-------- .../Dialect/QTensor/IR/Operations/AllocOp.cpp | 5 --- .../QTensor/IR/Operations/DeallocOp.cpp | 9 ---- .../QTensor/IR/Operations/ExtractOp.cpp | 7 ---- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 21 ++-------- .../QTensor/IR/Operations/FromTensorOp.cpp | 15 ------- .../QTensor/IR/Operations/InsertOp.cpp | 14 +------ .../QTensor/IR/Operations/InsertSliceOp.cpp | 21 +--------- mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 1 + 9 files changed, 32 insertions(+), 102 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index cb89b444a0..686d2c5eef 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -31,6 +31,8 @@ def QTensorDialect : Dialect { In addition, alloc/dealloc operations are added to the dialect to support the bulk allocation and deallocation of tensors with linear types. }]; + let dependentDialects = ["::mlir::qco::QCODialect"]; + let cppNamespace = "::mlir::qtensor"; } @@ -41,6 +43,14 @@ def QTensorDialect : Dialect { class QTensorOp traits = []> : Op; +//===----------------------------------------------------------------------===// +// Type Constraints +//===----------------------------------------------------------------------===// + +def QubitType : + Type($_self)">, + "Elements must be of qubit type">; + //===----------------------------------------------------------------------===// // Operations //===----------------------------------------------------------------------===// @@ -59,7 +69,7 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ }]; let arguments = (ins ConfinedAttr:$size); - let results = (outs 1DTensorOf<[AnyType]>:$result); + let results = (outs 1DTensorOf<[QubitType]>:$result); let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; let builders = [ @@ -80,11 +90,10 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { ``` }]; - let arguments = (ins 1DTensorOf<[AnyType]>:$tensor); + let arguments = (ins 1DTensorOf<[QubitType]>:$tensor); let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; let hasCanonicalizer = 1; - let hasVerifier = 1; } def QTensor_FromElementsOp : QTensorOp<"from_elements", [ @@ -105,16 +114,14 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ ``` }]; - let arguments = (ins Variadic:$elements); - let results = (outs 1DTensorOf<[AnyType]>:$result); + let arguments = (ins Variadic:$elements); + let results = (outs 1DTensorOf<[QubitType]>:$result); let assemblyFormat = "$elements attr-dict `:` type($result)"; let builders = [ // Special case builder for when `elements` has size >=1. OpBuilder<(ins "ValueRange":$elements)> ]; - - let hasVerifier = 1; } def QTensor_ExtractOp : QTensorOp<"extract", [ @@ -137,8 +144,8 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ ``` }]; - let arguments = (ins 1DTensorOf<[AnyType]>:$tensor, Index:$index); - let results = (outs 1DTensorOf<[AnyType]>:$out_tensor, AnyType:$result); + 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; @@ -170,11 +177,11 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ }]; let arguments = (ins - 1DTensorOf<[AnyType]>:$tensor, + 1DTensorOf<[QubitType]>:$tensor, Index:$offset, Index:$size ); - let results = (outs 1DTensorOf<[AnyType]>:$out_tensor, 1DTensorOf<[AnyType]>:$result); + let results = (outs 1DTensorOf<[QubitType]>:$out_tensor, 1DTensorOf<[QubitType]>:$result); let assemblyFormat = [{ $tensor `[`$offset `]` `[`$size`]` @@ -210,10 +217,10 @@ def QTensor_InsertOp : QTensorOp<"insert", [ ``` }]; - let arguments = (ins AnyType:$scalar, - 1DTensorOf<[AnyType]>:$dest, + let arguments = (ins QubitType:$scalar, + 1DTensorOf<[QubitType]>:$dest, Index:$index); - let results = (outs 1DTensorOf<[AnyType]>:$result); + let results = (outs 1DTensorOf<[QubitType]>:$result); let assemblyFormat = [{ $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) }]; @@ -249,12 +256,12 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ }]; let arguments = (ins - 1DTensorOf<[AnyType]>:$source, - 1DTensorOf<[AnyType]>:$dest, + 1DTensorOf<[QubitType]>:$source, + 1DTensorOf<[QubitType]>:$dest, Index:$offset, Index:$size ); - let results = (outs 1DTensorOf<[AnyType]>:$result); + let results = (outs 1DTensorOf<[QubitType]>:$result); let assemblyFormat = [{ $source `into` $dest `[`$offset `]` `[`$size`]` attr-dict `:` type($source) `into` type($dest) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index 0a386abb6e..818b4886f1 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -35,11 +35,6 @@ void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { LogicalResult AllocOp::verify() { auto resultType = getResult().getType(); - - if (!llvm::isa(resultType.getElementType())) { - return emitOpError("Result element type must be of qubit type"); - } - auto size = static_cast(getSize()); if (resultType.getShape()[0] != size) { diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp index dab02ba466..22ba3cf078 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp @@ -8,10 +8,8 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -19,13 +17,6 @@ using namespace mlir; using namespace mlir::qtensor; -LogicalResult DeallocOp::verify() { - if (!llvm::isa(getTensor().getType().getElementType())) { - return emitOpError("Elements of tensor must be of qubit type"); - } - return success(); -} - namespace { /** diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 8ef85e9062..1490732913 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -8,10 +8,8 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -22,14 +20,9 @@ using namespace mlir; using namespace mlir::qtensor; LogicalResult ExtractOp::verify() { - auto tensorType = getTensor().getType(); auto tensorDim = getTensor().getType().getDimSize(0); auto index = getConstantIntValue(getIndex()); - if (!llvm::isa(tensorType.getElementType())) { - return emitOpError("Elements of tensor must be of qubit type"); - } - if (index) { if (*index < 0) { return emitOpError("Index must be non-negative"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 331f7bbdb9..5a1c8b0677 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -8,10 +8,8 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -39,17 +37,11 @@ void ExtractSliceOp::build(OpBuilder& b, OperationState& result, Value tensor, } LogicalResult ExtractSliceOp::verify() { - auto tensorType = getTensor().getType(); - auto resultType = getResult().getType(); - - auto tensorDim = tensorType.getDimSize(0); + auto tensorDim = getTensor().getType().getDimSize(0); + auto resultDim = getResult().getType().getDimSize(0); auto constOffset = getConstantIntValue(getOffset()); auto constSize = getConstantIntValue(getSize()); - if (!llvm::isa(tensorType.getElementType())) { - return emitOpError("Elements of source tensor must be of qubit type"); - } - if (constOffset && *constOffset < 0) { return emitOpError("Offset must be non-negative"); } @@ -64,13 +56,8 @@ LogicalResult ExtractSliceOp::verify() { } } - if (resultType.getElementType() != tensorType.getElementType()) { - return emitOpError( - "Result element type must match input tensor element type"); - } - - if (constSize && !ShapedType::isDynamic(resultType.getDimSize(0))) { - if (resultType.getDimSize(0) != *constSize) { + if (constSize && !ShapedType::isDynamic(resultDim)) { + if (resultDim != *constSize) { return emitOpError("Result tensor dimension must match size operand"); } } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 7288122063..338ad27fc6 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -8,10 +8,8 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -31,16 +29,3 @@ void FromElementsOp::build(OpBuilder& builder, OperationState& result, {static_cast(elements.size())}, elements.front().getType()); build(builder, result, resultType, elements); } - -LogicalResult FromElementsOp::verify() { - if (!llvm::isa(getResult().getType().getElementType())) { - return emitOpError("Result tensor must have qubit element type"); - } - - for (auto type : getElements().getTypes()) { - if (!llvm::isa(type)) { - return emitOpError("Elements of ValueRange must be of qubit type"); - } - } - return success(); -} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 5e4cf95c89..9ea15c3aeb 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -8,10 +8,8 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -23,19 +21,9 @@ using namespace mlir; using namespace mlir::qtensor; LogicalResult InsertOp::verify() { - auto destType = getDest().getType(); - - auto dstDim = destType.getDimSize(0); + auto dstDim = getDest().getType().getDimSize(0); auto index = getConstantIntValue(getIndex()); - if (!llvm::isa(getScalar().getType())) { - return emitOpError("Scalar must be of qubit type"); - } - - if (!llvm::isa(destType.getElementType())) { - return emitOpError("Elements of dest tensor must be of qubit type"); - } - if (index) { if (*index < 0) { return emitOpError("Index must be non-negative"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 86353a8412..62b9d7c3ec 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -8,10 +8,8 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include @@ -24,22 +22,11 @@ using namespace mlir; using namespace mlir::qtensor; LogicalResult InsertSliceOp::verify() { - auto sourceType = getSource().getType(); - auto destType = getDest().getType(); - - auto srcDim = sourceType.getDimSize(0); - auto dstDim = destType.getDimSize(0); + auto srcDim = getSource().getType().getDimSize(0); + auto dstDim = getDest().getType().getDimSize(0); auto constOffset = getConstantIntValue(getOffset()); auto constSize = getConstantIntValue(getSize()); - if (!llvm::isa(sourceType.getElementType())) { - return emitOpError("Elements of source tensor must be of qubit type"); - } - - if (!llvm::isa(destType.getElementType())) { - return emitOpError("Elements of dest tensor must be of qubit type"); - } - if (constOffset && *constOffset < 0) { return emitOpError("Offset must be non-negative"); } @@ -60,10 +47,6 @@ LogicalResult InsertSliceOp::verify() { } } - if (getResult().getType() != destType) { - return emitOpError("Result type must match dest type"); - } - return success(); } diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index b1bee05cc7..d345f95158 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -10,6 +10,7 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated #include From 0027839237fd7adbaffddc4b8ab2aaa4c5dd5027 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 09:15:12 +0100 Subject: [PATCH 088/108] remove redundant headers --- mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp | 1 - mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp | 2 -- 2 files changed, 3 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index 818b4886f1..72b040b90f 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -11,7 +11,6 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp index 338ad27fc6..9d3b6a70f9 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/FromTensorOp.cpp @@ -13,8 +13,6 @@ #include #include #include -#include -#include #include #include From 01bddadda8df88804e0b77d79add67ff13a3c43c Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 09:19:02 +0100 Subject: [PATCH 089/108] Update CHANGELOG.md --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ed8c0f998..e0c5140e68 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]) ([**@denialhaag**]) - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547]) ([**@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]) + ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**]) ### Changed @@ -336,6 +336,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 From 7f41930af79e683d1085bc14d0a6413e23101914 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 09:30:41 +0100 Subject: [PATCH 090/108] fix linter issues --- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h | 1 + mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index 7d4c6dff75..862708a26b 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -22,6 +22,7 @@ #pragma clang diagnostic pop #endif +#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" #include diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index d345f95158..b1bee05cc7 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -10,7 +10,6 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated #include From 0261825fa8bc0f663d298afe10dadcd7e0b1298c Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 10:11:08 +0100 Subject: [PATCH 091/108] apply coderabbit suggestions --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 11 +++++-- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 29 ++++++++++--------- .../Dialect/QTensor/IR/Operations/AllocOp.cpp | 3 +- 3 files changed, 26 insertions(+), 17 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 686d2c5eef..a475293e6e 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -51,6 +51,13 @@ def QubitType : Type($_self)">, "Elements must be of qubit type">; +def Static1DQubitTensor : Type< + And<[ + 1DTensorOf<[QubitType]>.predicate, + HasStaticShapePred + ]>, + "Tensor must be statically shaped and elements must be of qubit type">; + //===----------------------------------------------------------------------===// // Operations //===----------------------------------------------------------------------===// @@ -69,7 +76,7 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ }]; let arguments = (ins ConfinedAttr:$size); - let results = (outs 1DTensorOf<[QubitType]>:$result); + let results = (outs Static1DQubitTensor:$result); let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; let builders = [ @@ -115,7 +122,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ }]; let arguments = (ins Variadic:$elements); - let results = (outs 1DTensorOf<[QubitType]>:$result); + let results = (outs Static1DQubitTensor:$result); let assemblyFormat = "$elements attr-dict `:` type($result)"; let builders = [ diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 2e7c3ad4f8..44897edb92 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -229,12 +229,12 @@ QCOProgramBuilder::qtensorExtract(Value tensor, const std::variant& index) { checkFinalized(); - auto rankedTensorType = llvm::dyn_cast(tensor.getType()); + auto tensorType = llvm::dyn_cast(tensor.getType()); - if (!rankedTensorType) { - llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + if (!tensorType || tensorType.getRank() != 1) { + llvm::reportFatalUsageError("Tensor must be of 1-D RankedTensorType!"); } - if (!llvm::isa(rankedTensorType.getElementType())) { + if (!llvm::isa(tensorType.getElementType())) { llvm::reportFatalUsageError("Elements must be of QubitType!"); } @@ -256,8 +256,8 @@ std::pair QCOProgramBuilder::qtensorExtractSlice( auto tensorType = llvm::dyn_cast(tensor.getType()); - if (!tensorType) { - llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + 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!"); @@ -282,8 +282,8 @@ Value QCOProgramBuilder::qtensorInsert( auto tensorType = llvm::dyn_cast(tensor.getType()); - if (!tensorType) { - llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + 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!"); @@ -307,8 +307,9 @@ Value QCOProgramBuilder::qtensorInsertSlice( auto sourceTensorType = llvm::dyn_cast(source.getType()); - if (!sourceTensorType) { - llvm::reportFatalUsageError("Source must be of RankedTensorType!"); + if (!sourceTensorType || sourceTensorType.getRank() != 1) { + llvm::reportFatalUsageError( + "Source tensor must be of 1-D RankedTensorType!"); } if (!llvm::isa(sourceTensorType.getElementType())) { llvm::reportFatalUsageError("Source elements must be of QubitType!"); @@ -316,8 +317,8 @@ Value QCOProgramBuilder::qtensorInsertSlice( auto destTensorType = llvm::dyn_cast(dest.getType()); - if (!destTensorType) { - llvm::reportFatalUsageError("Dest must be of RankedTensorType!"); + if (!destTensorType || destTensorType.getRank() != 1) { + llvm::reportFatalUsageError("Dest tensor must be of 1-D RankedTensorType!"); } if (!llvm::isa(destTensorType.getElementType())) { llvm::reportFatalUsageError("Dest elements must be of QubitType!"); @@ -342,8 +343,8 @@ QCOProgramBuilder& QCOProgramBuilder::qtensorDealloc(Value tensor) { auto tensorType = llvm::dyn_cast(tensor.getType()); - if (!tensorType) { - llvm::reportFatalUsageError("Tensor must be of RankedTensorType!"); + 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!"); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index 72b040b90f..d337280970 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -11,6 +11,7 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include #include #include @@ -33,7 +34,7 @@ void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { } LogicalResult AllocOp::verify() { - auto resultType = getResult().getType(); + auto resultType = cast(getResult().getType()); auto size = static_cast(getSize()); if (resultType.getShape()[0] != size) { From 60bb305dafa962c6e1c4f076ac826ceeeb2a6f27 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 11:00:29 +0100 Subject: [PATCH 092/108] apply coderabbit feedback --- mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 1 + .../Dialect/QTensor/IR/Operations/AllocOp.cpp | 2 +- .../QTensor/IR/Operations/InsertOp.cpp | 3 +++ .../QTensor/IR/Operations/InsertSliceOp.cpp | 3 +++ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 8 ++++++- mlir/unittests/programs/qco_programs.cpp | 22 +++++++++++++++++++ mlir/unittests/programs/qco_programs.h | 7 ++++++ 7 files changed, 44 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt index 17c3149d32..774bd86635 100644 --- a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -19,6 +19,7 @@ add_mlir_dialect_library( LINK_LIBS PUBLIC MLIRIR + MLIRQCODialect MLIRArithDialect MLIRInferTypeOpInterface MLIRSideEffectInterfaces diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index d337280970..ff72a7b99f 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -11,11 +11,11 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" -#include #include #include #include #include +#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 9ea15c3aeb..5adad13ed6 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -8,6 +8,7 @@ * Licensed under the MIT License */ +#include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include @@ -94,6 +95,8 @@ struct CombineSubsequentInsertOp : public OpRewritePattern { return failure(); } + qco::DeallocOp::create(rewriter, prevInsertOp.getLoc(), + prevInsertOp.getScalar()); rewriter.replaceOpWithNewOp(insertOp, insertOp.getScalar(), prevInsertOp.getDest(), insertIndex); rewriter.eraseOp(prevInsertOp); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 62b9d7c3ec..626d0b901d 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -118,10 +118,13 @@ struct CombineSubsequentInsertSliceOp final !isSameIndex(prevSize, curSize)) { return failure(); } + DeallocOp::create(rewriter, prevInsertOp->getLoc(), + prevInsertOp.getSource()); rewriter.replaceOpWithNewOp( insertSliceOp, insertSliceOp.getSource(), prevInsertOp.getDest(), curOffset, curSize); rewriter.eraseOp(prevInsertOp); + return success(); } }; diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 99229be742..621230a811 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1095,5 +1095,11 @@ INSTANTIATE_TEST_SUITE_P( QCOTestCase{ "QTensorExtractSliceExtractInsertInsertSlice", MQT_NAMED_BUILDER(qtensorExtractSliceExtractInsertInsertSlice), - MQT_NAMED_BUILDER(qtensorAlloc)})); + MQT_NAMED_BUILDER(qtensorAlloc)}, + QCOTestCase{"QTensorInsertInsert", + MQT_NAMED_BUILDER(qtensorInsertInsert), + MQT_NAMED_BUILDER(qtensorInsert)}, + QCOTestCase{"QTensorInsertSliceInsertSlice", + MQT_NAMED_BUILDER(qtensorInsertSliceInsertSlice), + MQT_NAMED_BUILDER(qtensorInsertSlice)})); /// @} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index f8f84345cb..9b7f6c6f04 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2220,4 +2220,26 @@ void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b) { b.qtensorInsertSlice(slicedTensor1, extractSliceOutTensor1, 0, 2); } +void qtensorInsertInsert(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto qTemp = b.allocQubit(); + auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); + auto insertTensorTemp = b.qtensorInsert(qTemp, extractOutTensor, 0); + auto q1 = b.h(q0); + auto insertOutTensor = b.qtensorInsert(q1, insertTensorTemp, 0); +} + +void qtensorInsertSliceInsertSlice(QCOProgramBuilder& b) { + auto qtensor = b.qtensorAlloc(3); + auto qtensorTemp = b.qtensorAlloc(2); + 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 insertSliceTensorTemp = + b.qtensorInsertSlice(qtensorTemp, extractSliceOutTensor, 0, 2); + b.qtensorInsertSlice(insertOutTensor, insertSliceTensorTemp, 0, 2); +} + } // namespace mlir::qco diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 901bc1f5c7..68019e16e2 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1012,4 +1012,11 @@ void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b); /// the slice and the slice back to the tensor immediately. void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b); +/// Inserts two qubits to the same index back to back. +void qtensorInsertInsert(QCOProgramBuilder& b); + +/// Inserts two slices of qubits with the same size and offset into the same +/// tensor back to back. +void qtensorInsertSliceInsertSlice(QCOProgramBuilder& b); + } // namespace mlir::qco From b034a3b90398645b2acc3bb623faaf8990db981a Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 11:09:41 +0100 Subject: [PATCH 093/108] fix linter issue --- mlir/unittests/programs/qco_programs.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 9b7f6c6f04..1cca432e3d 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2226,7 +2226,7 @@ void qtensorInsertInsert(QCOProgramBuilder& b) { auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); auto insertTensorTemp = b.qtensorInsert(qTemp, extractOutTensor, 0); auto q1 = b.h(q0); - auto insertOutTensor = b.qtensorInsert(q1, insertTensorTemp, 0); + b.qtensorInsert(q1, insertTensorTemp, 0); } void qtensorInsertSliceInsertSlice(QCOProgramBuilder& b) { From 1937a75b0235f9fac95a9c9affdc7d3274c2376b Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Mon, 16 Mar 2026 11:23:01 +0100 Subject: [PATCH 094/108] fix grammar mistakes --- mlir/unittests/programs/qco_programs.h | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 68019e16e2..222c2023be 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -975,6 +975,8 @@ void nestedTrueIf(QCOProgramBuilder& b); /// the same condition. void nestedFalseIf(QCOProgramBuilder& b); +// --- QTensor Operations -------------------------------------------------- // + /// Allocates a tensor of size `3`. void qtensorAlloc(QCOProgramBuilder& b); @@ -996,20 +998,21 @@ void qtensorExtractSlice(QCOProgramBuilder& b); /// Inserts a slice into a tensor. void qtensorInsertSlice(QCOProgramBuilder& b); -/// Extracts a qubit from a tensor and insert it immediately. +/// Extracts a qubit from a tensor and inserts it immediately. void qtensorExtractInsert(QCOProgramBuilder& b); -/// Inserts a qubit into a tensor and extract it immediately. +/// Inserts a qubit into a tensor and extracts it immediately. void qtensorInsertExtract(QCOProgramBuilder& b); -/// Extracts a slice of qubits from a tensor and insert it immediately. +/// Extracts a slice of qubits from a tensor and inserts it immediately. void qtensorExtractSliceInsertSlice(QCOProgramBuilder& b); -/// Inserts a slice of qubits into a tensor and extract it immediately. +/// Inserts a slice of qubits into a tensor and extracts it immediately. void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b); -/// Extracts a slice of qubits, a qubit from the slice, insert the qubit back to -/// the slice and the slice back to the tensor immediately. +/// 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. void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b); /// Inserts two qubits to the same index back to back. From d265531af2d5870bcfaffd8173ffe33f9bfe43d8 Mon Sep 17 00:00:00 2001 From: li-mingbao <74404929+li-mingbao@users.noreply.github.com> Date: Wed, 18 Mar 2026 09:27:08 +0100 Subject: [PATCH 095/108] Apply suggestions from code review Co-authored-by: Lukas Burgholzer Signed-off-by: li-mingbao <74404929+li-mingbao@users.noreply.github.com> --- docs/mlir/index.md | 2 +- .../mlir/Dialect/QTensor/IR/QTensorDialect.h | 4 ---- .../include/mlir/Dialect/QTensor/IR/QTensorOps.td | 15 +++++++-------- mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 2 +- .../Dialect/QTensor/IR/Operations/DeallocOp.cpp | 4 ++-- .../Dialect/QTensor/IR/Operations/ExtractOp.cpp | 4 ++-- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 4 ++-- 7 files changed, 15 insertions(+), 20 deletions(-) diff --git a/docs/mlir/index.md b/docs/mlir/index.md index 42ff96def3..f928fd86cd 100644 --- a/docs/mlir/index.md +++ b/docs/mlir/index.md @@ -8,7 +8,7 @@ We define multiple dialects, each with its dedicated purpose: - The {doc}`QCO dialect ` uses value semantics and is mainly designed for running optimizations. -- The {doc}`QTensor dialect ` adds support for tensors with linear typing and is used in the QCO dialect to represent registers. +- 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. diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h index f16d153d1d..1d88688be8 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.h @@ -27,7 +27,3 @@ #define GET_TYPEDEF_CLASSES #include "mlir/Dialect/QTensor/IR/QTensorOpsTypes.h.inc" // IWYU pragma: export - -//===----------------------------------------------------------------------===// -// QTensor Dialect Helpers -//===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index a475293e6e..181dff3f4b 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -143,7 +143,6 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ 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 element from the input tensor at the given index. In addition, it also returns the updated input tensor as a result. - The index must be of `index` type. Example: ```mlir @@ -172,10 +171,10 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ 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 + - 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. + - size: the length of the slice to extract from the base tensor. Example: ```mlir @@ -250,10 +249,10 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ The insert_slice operation supports the following arguments: - * source: the tensor that is inserted. - * 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: the tensor that is inserted. + - 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: diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt index 774bd86635..4448bbd601 100644 --- a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -13,7 +13,7 @@ add_mlir_dialect_library( QTensorOps.cpp ${OPERATIONS} ADDITIONAL_HEADER_DIRS - ${PROJECT_SOURCE_DIR}/mlir/include/mlir/Dialect/QTensor + ${MQT_MLIR_SOURCE_INCLUDE_DIR}/mlir/Dialect/QTensor DEPENDS MLIRQTensorOpsIncGen LINK_LIBS diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp index 22ba3cf078..ec1f38f45e 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp @@ -28,10 +28,10 @@ struct RemoveAllocDeallocPair final : OpRewritePattern { LogicalResult matchAndRewrite(DeallocOp op, PatternRewriter& rewriter) const override { - // Check if the predecessor is an qtensor::AllocOp + // Check if the predecessor is a qtensor::AllocOp auto tensor = op.getTensor(); auto allocOp = tensor.getDefiningOp(); - if (!allocOp || !tensor.hasOneUse()) { + if (!allocOp) { return failure(); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 1490732913..609f1b7b97 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -60,8 +60,8 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { LogicalResult ExtractOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { if (auto insertOp = foldExtractAfterInsert(*this)) { - results.push_back(insertOp.getDest()); - results.push_back(insertOp.getScalar()); + results.emplace_back(insertOp.getDest()); + results.emplace_back(insertOp.getScalar()); return success(); } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index 5a1c8b0677..aa18dfcb16 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -99,8 +99,8 @@ foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { LogicalResult ExtractSliceOp::fold(FoldAdaptor /*adaptor*/, SmallVectorImpl& results) { if (auto insertOp = foldExtractAfterInsertSlice(*this)) { - results.push_back(insertOp.getDest()); - results.push_back(insertOp.getSource()); + results.emplace_back(insertOp.getDest()); + results.emplace_back(insertOp.getSource()); return success(); } From 46d06a8492b9c940504dd7d9d589c2b7905b8e99 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 10:20:22 +0100 Subject: [PATCH 096/108] update the descriptions --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 62 ++++++++++++++++--- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 23 ++++--- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 3a43d38eab..daaa39d5e3 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -209,6 +209,12 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Allocate a qubit tensor + * + * @details + * Allocates an one-dimensional tensor of !qco.qubit types with the given size + * if the size is statically known, otherwise the tensor has dynamic size. The + * resulting tensor is added to the tracking. + * * @param size Number of qubits (must be positive) * @return The allocated tensor * @@ -224,7 +230,14 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Allocate a qubit tensor from a list of qubit values - * @param elements Inserted Qubits + * + * @details + * Consumes the input qubits and creates an 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: @@ -239,7 +252,14 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Extract a qubit from a tensor - * @param tensor Source 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) * @@ -256,7 +276,14 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Extract a qubit slice from a tensor - * @param tensor Source 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) @@ -277,8 +304,15 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Insert a qubit into a tensor - * @param scalar The scalar qubit that is inserted - * @param tensor The tensor where the qubit is inserted + * + * @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 * @@ -295,8 +329,16 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Insert a qubit slice into a tensor - * @param sourceTensor The slice that is inserted - * @param destTensor The tensor where the slice is inserted + * + * @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 @@ -316,6 +358,12 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Explicitly deallocate a tensor + * + * @details + * Explicitly deallocates the given tensor and the value it holds. Qubits or + * tensors of qubits that were extracted from the tensor but not inserted back + * again need to be deallocated separately. + * * @param tensor Tensor to deallocate (must be valid/unconsumed) * @return Reference to this builder for method chaining * diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 181dff3f4b..bd828c8a78 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -23,12 +23,12 @@ include "mlir/Interfaces/TilingInterface.td" def QTensorDialect : Dialect { let name = "qtensor"; - let summary = "The QTensor dialect for tensors with linear typing."; + 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 tensor types. - 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 tensors with linear types. + 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"]; @@ -90,6 +90,7 @@ def QTensor_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 @@ -113,7 +114,8 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ 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 tensor using the operands as its values. + This operation creates a one-dimensional qubit tensor using the operands as its input. The also causes + the input qubits to be consumed. Example: ```mlir @@ -141,7 +143,7 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ 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 element from the input tensor + 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: @@ -166,7 +168,7 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ 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 ranked tensor and returns the extracted tensor specified by the + operation of the tensor dialect. It reads a one-dimensional qubit 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: @@ -215,7 +217,8 @@ def QTensor_InsertOp : QTensorOp<"insert", [ 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 scalar into a ranked tensor as specified by the operation's index. + 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 @@ -245,11 +248,11 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ 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. + operation's offset and size arguments. This insertion consumes `source` tensor of qubits. The insert_slice operation supports the following arguments: - - source: the tensor that is inserted. + - 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 From 6cc90e6f259e628b980c5c04a2abbbed9227ceb4 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 10:31:00 +0100 Subject: [PATCH 097/108] apply code review suggestions and simplify code --- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 2 - .../QTensor/IR/Operations/ExtractOp.cpp | 13 ++--- .../QTensor/IR/Operations/ExtractSliceOp.cpp | 9 +--- .../QTensor/IR/Operations/InsertOp.cpp | 44 +-------------- .../QTensor/IR/Operations/InsertSliceOp.cpp | 54 +------------------ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 11 +--- mlir/unittests/programs/qco_programs.cpp | 22 -------- mlir/unittests/programs/qco_programs.h | 7 --- 8 files changed, 14 insertions(+), 148 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index bd828c8a78..3c01aa3ae6 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -234,7 +234,6 @@ def QTensor_InsertOp : QTensorOp<"insert", [ $scalar `into` $dest `[` $index `]` attr-dict `:` type($dest) }]; - let hasCanonicalizer = 1; let hasFolder = 1; let hasVerifier = 1; } @@ -276,7 +275,6 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ attr-dict `:` type($source) `into` type($dest) }]; - let hasCanonicalizer = 1; let hasFolder = 1; let hasVerifier = 1; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index 609f1b7b97..fc030e08e5 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -44,16 +44,17 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { return nullptr; } - if (insertOp.getScalar().getType() != extractOp.getResult().getType()) { - return nullptr; - } + auto isSame = [](Value a, Value b) { + return getAsOpFoldResult(a) == getAsOpFoldResult(b); + }; - auto insertIndex = insertOp.getIndex(); - auto extractIndex = extractOp.getIndex(); + Value insertIndex = insertOp.getIndex(); + Value extractIndex = extractOp.getIndex(); - if (!isSameIndex(insertIndex, extractIndex)) { + if (getAsOpFoldResult(insertIndex) != getAsOpFoldResult(extractIndex)) { return nullptr; } + return insertOp; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp index aa18dfcb16..3c7648b04d 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractSliceOp.cpp @@ -78,18 +78,13 @@ foldExtractAfterInsertSlice(ExtractSliceOp extractSliceOp) { return nullptr; } - if (insertSliceOp.getSource().getType() != - extractSliceOp.getResult().getType()) { - return nullptr; - } - auto insertOffset = insertSliceOp.getOffset(); auto extractOffset = extractSliceOp.getOffset(); auto insertSize = insertSliceOp.getSize(); auto extractSize = extractSliceOp.getSize(); - if (!isSameIndex(insertOffset, extractOffset) || - !isSameIndex(insertSize, extractSize)) { + if (getAsOpFoldResult(insertOffset) != getAsOpFoldResult(extractOffset) || + getAsOpFoldResult(insertSize) != getAsOpFoldResult(extractSize)) { return nullptr; } diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 5adad13ed6..1b73731198 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -50,14 +50,11 @@ static Value foldInsertAfterExtract(InsertOp insertOp) { if (insertOp.getDest() != extractOp.getOutTensor()) { return nullptr; } - if (extractOp.getTensor().getType() != insertOp.getDest().getType()) { - return nullptr; - } auto insertIndex = insertOp.getIndex(); auto extractIndex = extractOp.getIndex(); - if (!isSameIndex(insertIndex, extractIndex)) { + if (getAsOpFoldResult(insertIndex) != getAsOpFoldResult(extractIndex)) { return nullptr; } @@ -71,42 +68,3 @@ OpFoldResult InsertOp::fold(FoldAdaptor /*adaptor*/) { return {}; } - -namespace { - -/** - * @brief Combine subsequent insert operations with the same index. - */ -struct CombineSubsequentInsertOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOp insertOp, - PatternRewriter& rewriter) const final { - auto prevInsertOp = insertOp.getDest().getDefiningOp(); - - if (!prevInsertOp) { - return failure(); - } - - auto insertIndex = insertOp.getIndex(); - auto prevInsertIndex = prevInsertOp.getIndex(); - - if (!isSameIndex(insertIndex, prevInsertIndex)) { - return failure(); - } - - qco::DeallocOp::create(rewriter, prevInsertOp.getLoc(), - prevInsertOp.getScalar()); - rewriter.replaceOpWithNewOp(insertOp, insertOp.getScalar(), - prevInsertOp.getDest(), insertIndex); - rewriter.eraseOp(prevInsertOp); - return success(); - } -}; - -} // namespace - -void InsertOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 626d0b901d..15c71f0c35 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -70,8 +70,8 @@ static Value foldInsertAfterExtractSlice(InsertSliceOp insertSliceOp) { auto insertSize = insertSliceOp.getSize(); auto extractSize = extractSliceOp.getSize(); - if (!isSameIndex(insertOffset, extractOffset) || - !isSameIndex(insertSize, extractSize)) { + if (getAsOpFoldResult(insertOffset) != getAsOpFoldResult(extractOffset) || + getAsOpFoldResult(insertSize) != getAsOpFoldResult(extractSize)) { return nullptr; } @@ -85,53 +85,3 @@ OpFoldResult InsertSliceOp::fold(FoldAdaptor /*adaptor*/) { return {}; } - -namespace { - -/** - * @brief Combine subsequent insertSlice operations with the same offset and - * size. - */ -struct CombineSubsequentInsertSliceOp final - : public OpRewritePattern { - - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp, - PatternRewriter& rewriter) const override { - auto prevInsertOp = insertSliceOp.getDest().getDefiningOp(); - if (!prevInsertOp) { - return failure(); - } - - if (prevInsertOp.getSource().getType() != - insertSliceOp.getSource().getType()) { - return failure(); - } - - auto prevOffset = prevInsertOp.getOffset(); - auto curOffset = insertSliceOp.getOffset(); - auto prevSize = prevInsertOp.getSize(); - auto curSize = insertSliceOp.getSize(); - - if (!isSameIndex(prevOffset, curOffset) || - !isSameIndex(prevSize, curSize)) { - return failure(); - } - DeallocOp::create(rewriter, prevInsertOp->getLoc(), - prevInsertOp.getSource()); - rewriter.replaceOpWithNewOp( - insertSliceOp, insertSliceOp.getSource(), prevInsertOp.getDest(), - curOffset, curSize); - rewriter.eraseOp(prevInsertOp); - - return success(); - } -}; - -} // namespace - -void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet& results, - MLIRContext* context) { - results.add(context); -} diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 621230a811..c31b963753 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include @@ -50,7 +49,7 @@ class QCOTest : public testing::TestWithParam { // Register all necessary dialects DialectRegistry registry; registry.insert(); + qtensor::QTensorDialect>(); context = std::make_unique(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); @@ -1095,11 +1094,5 @@ INSTANTIATE_TEST_SUITE_P( QCOTestCase{ "QTensorExtractSliceExtractInsertInsertSlice", MQT_NAMED_BUILDER(qtensorExtractSliceExtractInsertInsertSlice), - MQT_NAMED_BUILDER(qtensorAlloc)}, - QCOTestCase{"QTensorInsertInsert", - MQT_NAMED_BUILDER(qtensorInsertInsert), - MQT_NAMED_BUILDER(qtensorInsert)}, - QCOTestCase{"QTensorInsertSliceInsertSlice", - MQT_NAMED_BUILDER(qtensorInsertSliceInsertSlice), - MQT_NAMED_BUILDER(qtensorInsertSlice)})); + MQT_NAMED_BUILDER(qtensorAlloc)})); /// @} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 1cca432e3d..f8f84345cb 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2220,26 +2220,4 @@ void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b) { b.qtensorInsertSlice(slicedTensor1, extractSliceOutTensor1, 0, 2); } -void qtensorInsertInsert(QCOProgramBuilder& b) { - auto qtensor = b.qtensorAlloc(3); - auto qTemp = b.allocQubit(); - auto [extractOutTensor, q0] = b.qtensorExtract(qtensor, 0); - auto insertTensorTemp = b.qtensorInsert(qTemp, extractOutTensor, 0); - auto q1 = b.h(q0); - b.qtensorInsert(q1, insertTensorTemp, 0); -} - -void qtensorInsertSliceInsertSlice(QCOProgramBuilder& b) { - auto qtensor = b.qtensorAlloc(3); - auto qtensorTemp = b.qtensorAlloc(2); - 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 insertSliceTensorTemp = - b.qtensorInsertSlice(qtensorTemp, extractSliceOutTensor, 0, 2); - b.qtensorInsertSlice(insertOutTensor, insertSliceTensorTemp, 0, 2); -} - } // namespace mlir::qco diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 222c2023be..e3676e3313 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1015,11 +1015,4 @@ void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b); /// immediately. void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b); -/// Inserts two qubits to the same index back to back. -void qtensorInsertInsert(QCOProgramBuilder& b); - -/// Inserts two slices of qubits with the same size and offset into the same -/// tensor back to back. -void qtensorInsertSliceInsertSlice(QCOProgramBuilder& b); - } // namespace mlir::qco From a8cbaf6487459f6fc6fa2b0e6a3a99a7e9b02ecd Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 10:31:20 +0100 Subject: [PATCH 098/108] remove dead code --- mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp index fc030e08e5..27e8de6995 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/ExtractOp.cpp @@ -44,10 +44,6 @@ static InsertOp foldExtractAfterInsert(ExtractOp extractOp) { return nullptr; } - auto isSame = [](Value a, Value b) { - return getAsOpFoldResult(a) == getAsOpFoldResult(b); - }; - Value insertIndex = insertOp.getIndex(); Value extractIndex = extractOp.getIndex(); From 4799bb87aee9d2375adc88bfa4569274b03ef385 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 10:33:11 +0100 Subject: [PATCH 099/108] remove redundant helper function --- .../mlir/Dialect/QTensor/IR/QTensorOps.h | 13 ------------ mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp | 21 ------------------- 2 files changed, 34 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h index 862708a26b..677873c6e6 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.h @@ -30,16 +30,3 @@ #define GET_OP_CLASSES #include "mlir/Dialect/QTensor/IR/QTensorOps.h.inc" // IWYU pragma: export - -namespace mlir::qtensor { - -/** - * @brief Check if two values of IndexType are identical. - * - * @param index1 The first IndexType value. - * @param index2 The second IndexType value. - * @return True if both values are equal. - */ -bool isSameIndex(TypedValue index1, TypedValue index2); - -} // namespace mlir::qtensor diff --git a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp index b1bee05cc7..c26ed12d62 100644 --- a/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp +++ b/mlir/lib/Dialect/QTensor/IR/QTensorOps.cpp @@ -12,13 +12,6 @@ #include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: associated -#include -#include -#include -#include -#include -#include - // The following headers are needed for some template instantiations. // IWYU pragma: begin_keep #include @@ -28,20 +21,6 @@ using namespace mlir; using namespace mlir::qtensor; -namespace mlir::qtensor { - -bool isSameIndex(TypedValue index1, TypedValue index2) { - if (index1 == index2) { - return true; - } - - auto val1 = getConstantIntValue(index1); - auto val2 = getConstantIntValue(index2); - - return val1 && val2 && *val1 == *val2; -} -} // namespace mlir::qtensor - //===----------------------------------------------------------------------===// // Dialect //===----------------------------------------------------------------------===// From d7bc3e24478403cb9f6b4dda6ccc1ac03fbb0bcf Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 11:01:17 +0100 Subject: [PATCH 100/108] add dynamic sizes for allocOp --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 13 +++---- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 15 ++++---- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 14 ++++---- .../Dialect/QTensor/IR/Operations/AllocOp.cpp | 34 +++++++++++++------ 4 files changed, 44 insertions(+), 32 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index daaa39d5e3..f93b556031 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -211,9 +211,10 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @brief Allocate a qubit tensor * * @details - * Allocates an one-dimensional tensor of !qco.qubit types with the given size - * if the size is statically known, otherwise the tensor has dynamic size. The - * resulting tensor is added to the tracking. + * 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 @@ -223,16 +224,16 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * auto tensor = builder.qtensorAlloc(3); * ``` * ```mlir - * %tensor = qtensor.alloc(3) : tensor<3x!qco.qubit> + * %tensor = qtensor.alloc(%c3) : tensor<3x!qco.qubit> * ``` */ - Value qtensorAlloc(int64_t size); + Value qtensorAlloc(const std::variant& size); /** * @brief Allocate a qubit tensor from a list of qubit values * * @details - * Consumes the input qubits and creates an one-dimensional tensor of + * 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. diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 3c01aa3ae6..f94406562d 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -66,21 +66,22 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ MemoryEffects<[MemAlloc]>]> { let summary = "Tensor alloc operation"; let description = [{ - Allocates a qubit tensor with the given size and returns the allocated tensor. - The qubits are initialized to the |0> state. + 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(3) : tensor<3x!qco.qubit> + %qtensor = qtensor.alloc(%c3) : tensor<3x!qco.qubit> ``` }]; - let arguments = (ins ConfinedAttr:$size); - let results = (outs Static1DQubitTensor:$result); + let arguments = (ins Index:$size); + let results = (outs 1DTensorOf<[QubitType]>:$result); let assemblyFormat = "`(`$size`)` attr-dict `:` type($result)"; let builders = [ - OpBuilder<(ins "int64_t":$size)> + OpBuilder<(ins "Value":$size)> ]; let hasVerifier = 1; @@ -247,7 +248,7 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ 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 `source` tensor of qubits. + operation's offset and size arguments. This insertion consumes the `source` tensor of qubits. The insert_slice operation supports the following arguments: diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 44897edb92..bdd5563df0 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -190,14 +190,12 @@ void QCOProgramBuilder::updateTensorTracking(Value inputTensor, // QTensor Operations //===----------------------------------------------------------------------===// -Value QCOProgramBuilder::qtensorAlloc(int64_t size) { +Value QCOProgramBuilder::qtensorAlloc( + const std::variant& size) { checkFinalized(); + auto sizeValue = utils::variantToValue(*this, getLoc(), size); - if (size <= 0) { - llvm::reportFatalUsageError("Size must be positive"); - } - - auto allocOp = qtensor::AllocOp::create(*this, size); + auto allocOp = qtensor::AllocOp::create(*this, sizeValue); auto result = allocOp.getResult(); validTensors.insert(result); return result; @@ -325,9 +323,9 @@ Value QCOProgramBuilder::qtensorInsertSlice( } auto offsetValue = utils::variantToValue(*this, getLoc(), offset); - auto sizesValue = utils::variantToValue(*this, getLoc(), size); + auto sizeValue = utils::variantToValue(*this, getLoc(), size); auto insertSliceOp = qtensor::InsertSliceOp::create(*this, source, dest, - offsetValue, sizesValue); + offsetValue, sizeValue); auto outTensor = insertSliceOp.getResult(); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index ff72a7b99f..be81806731 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -11,6 +11,7 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include #include #include @@ -19,27 +20,38 @@ #include #include -#include +#include using namespace mlir; using namespace mlir::qtensor; -void AllocOp::build(OpBuilder& builder, OperationState& result, int64_t size) { - assert(size > 0 && "qtensor.alloc size must be positive"); +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({size}, qco::QubitType::get(builder.getContext())); - build(builder, result, resultType, - IntegerAttr::get(builder.getIntegerType(64), size)); + 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 size = static_cast(getSize()); - - if (resultType.getShape()[0] != size) { - return emitOpError("Tensor length must match size attribute (") - << size << "), but got " << resultType.getShape()[0]; + auto sizeValue = getConstantIntValue(getSize()); + auto resultSize = resultType.getShape()[0]; + + 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(); From 7414ebd082c8707d813796e11b50ff7981806531 Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 11:04:13 +0100 Subject: [PATCH 101/108] apply coderabbit feedback --- mlir/lib/Dialect/QTensor/IR/CMakeLists.txt | 5 +++-- mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt index 4448bbd601..98aa36b0c5 100644 --- a/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/IR/CMakeLists.txt @@ -17,12 +17,13 @@ add_mlir_dialect_library( DEPENDS MLIRQTensorOpsIncGen LINK_LIBS - PUBLIC + PRIVATE MLIRIR - MLIRQCODialect MLIRArithDialect MLIRInferTypeOpInterface MLIRSideEffectInterfaces + PUBLIC + MLIRQCODialect DISABLE_INSTALL) mqt_mlir_target_use_project_options(MLIRQTensorDialect) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp index ec1f38f45e..90f076ede1 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/DeallocOp.cpp @@ -28,7 +28,7 @@ struct RemoveAllocDeallocPair final : OpRewritePattern { LogicalResult matchAndRewrite(DeallocOp op, PatternRewriter& rewriter) const override { - // Check if the predecessor is a qtensor::AllocOp + // Check whether the tensor is directly defined by a qtensor::AllocOp. auto tensor = op.getTensor(); auto allocOp = tensor.getDefiningOp(); if (!allocOp) { From 2a5a25c68e34812d0a353a31c5e57e3ea379a4ea Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 11:58:35 +0100 Subject: [PATCH 102/108] address coderabbit comments and linter issues --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 4 +- .../mlir/Dialect/QTensor/IR/QTensorOps.td | 2 +- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 74 ++----------------- .../Dialect/QTensor/IR/Operations/AllocOp.cpp | 5 +- .../QTensor/IR/Operations/InsertOp.cpp | 2 - .../QTensor/IR/Operations/InsertSliceOp.cpp | 1 - 6 files changed, 16 insertions(+), 72 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index f93b556031..5ae5aa1193 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1353,7 +1353,9 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { llvm::DenseSet validQubits; /** - * @brief Validate that a tensor value is valid and unconsumed + * @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) */ diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index f94406562d..797c5d521b 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -169,7 +169,7 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ 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 and returns the extracted tensor of qubits specified by the + 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: diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index bdd5563df0..2a5b6b4d66 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -172,6 +172,14 @@ void QCOProgramBuilder::validateTensorValue(Value tensor) const { 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, @@ -190,17 +198,6 @@ void QCOProgramBuilder::updateTensorTracking(Value inputTensor, // 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(); @@ -227,15 +224,6 @@ QCOProgramBuilder::qtensorExtract(Value tensor, const std::variant& index) { checkFinalized(); - 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!"); - } - auto indexValue = utils::variantToValue(*this, getLoc(), index); auto extractOp = qtensor::ExtractOp::create(*this, tensor, indexValue); auto qubit = extractOp.getResult(); @@ -252,15 +240,6 @@ std::pair QCOProgramBuilder::qtensorExtractSlice( const std::variant& size) { checkFinalized(); - 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!"); - } - auto offsetValue = utils::variantToValue(*this, getLoc(), offset); auto sizesValue = utils::variantToValue(*this, getLoc(), size); auto extractSliceOp = @@ -278,15 +257,6 @@ Value QCOProgramBuilder::qtensorInsert( Value scalar, Value tensor, const std::variant& index) { checkFinalized(); - 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!"); - } - auto indexValue = utils::variantToValue(*this, getLoc(), index); auto insertOp = qtensor::InsertOp::create(*this, scalar, tensor, indexValue); @@ -303,25 +273,6 @@ Value QCOProgramBuilder::qtensorInsertSlice( const std::variant& size) { checkFinalized(); - auto sourceTensorType = llvm::dyn_cast(source.getType()); - - if (!sourceTensorType || sourceTensorType.getRank() != 1) { - llvm::reportFatalUsageError( - "Source tensor must be of 1-D RankedTensorType!"); - } - if (!llvm::isa(sourceTensorType.getElementType())) { - llvm::reportFatalUsageError("Source elements must be of QubitType!"); - } - - auto destTensorType = llvm::dyn_cast(dest.getType()); - - if (!destTensorType || destTensorType.getRank() != 1) { - llvm::reportFatalUsageError("Dest tensor must be of 1-D RankedTensorType!"); - } - if (!llvm::isa(destTensorType.getElementType())) { - llvm::reportFatalUsageError("Dest elements must be of QubitType!"); - } - auto offsetValue = utils::variantToValue(*this, getLoc(), offset); auto sizeValue = utils::variantToValue(*this, getLoc(), size); auto insertSliceOp = qtensor::InsertSliceOp::create(*this, source, dest, @@ -339,15 +290,6 @@ Value QCOProgramBuilder::qtensorInsertSlice( QCOProgramBuilder& QCOProgramBuilder::qtensorDealloc(Value tensor) { checkFinalized(); - 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!"); - } - validateTensorValue(tensor); validTensors.erase(tensor); diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp index be81806731..898b8b6412 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/AllocOp.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -20,7 +21,6 @@ #include #include -#include using namespace mlir; using namespace mlir::qtensor; @@ -42,6 +42,9 @@ LogicalResult AllocOp::verify() { 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 ") diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp index 1b73731198..982d4a6335 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertOp.cpp @@ -8,13 +8,11 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include #include #include -#include #include #include diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index 15c71f0c35..d2601878c8 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -14,7 +14,6 @@ #include #include #include -#include #include #include From 2afb9e68e8c6947041ba0100e11270003a148e1d Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 12:05:35 +0100 Subject: [PATCH 103/108] add back accidently removed functioN --- mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 2a5b6b4d66..05b55ff6dc 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -198,6 +198,17 @@ void QCOProgramBuilder::updateTensorTracking(Value inputTensor, // 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(); From b31cf8d68c728f9c88551a2fa220fb7cd197d88f Mon Sep 17 00:00:00 2001 From: Li-ming Bao Date: Wed, 18 Mar 2026 12:22:25 +0100 Subject: [PATCH 104/108] add additional tests at different indices and offsets --- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 30 ++++++++--- mlir/unittests/programs/qco_programs.cpp | 54 +++++++++++++++---- mlir/unittests/programs/qco_programs.h | 36 +++++++++---- 3 files changed, 94 insertions(+), 26 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index c31b963753..a6c7396304 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1079,18 +1079,32 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(qtensorExtractSlice)}, QCOTestCase{"QTensorInsertSlice", MQT_NAMED_BUILDER(qtensorInsertSlice), MQT_NAMED_BUILDER(qtensorInsertSlice)}, - QCOTestCase{"QTensorExtractInsert", - MQT_NAMED_BUILDER(qtensorExtractInsert), + QCOTestCase{"QTensorExtractInsertSameIndex", + MQT_NAMED_BUILDER(qtensorExtractInsertSameIndex), MQT_NAMED_BUILDER(qtensorAlloc)}, - QCOTestCase{"QTensorInsertExtract", - MQT_NAMED_BUILDER(qtensorInsertExtract), + QCOTestCase{"QTensorExtractInsertIndexMismatch", + MQT_NAMED_BUILDER(qtensorExtractInsertIndexMismatch), + MQT_NAMED_BUILDER(qtensorExtractInsertIndexMismatch)}, + QCOTestCase{"QTensorInsertExtractSameIndex", + MQT_NAMED_BUILDER(qtensorInsertExtractSameIndex), MQT_NAMED_BUILDER(qtensorInsert)}, - QCOTestCase{"QTensorExtractSliceInsertSlice", - MQT_NAMED_BUILDER(qtensorExtractSliceInsertSlice), + QCOTestCase{"QTensorInsertExtractIndexMismatch", + MQT_NAMED_BUILDER(qtensorInsertExtractIndexMismatch), + MQT_NAMED_BUILDER(qtensorInsertExtractIndexMismatch)}, + QCOTestCase{"QTensorExtractSliceInsertSliceSameOffset", + MQT_NAMED_BUILDER(qtensorExtractSliceInsertSliceSameOffset), MQT_NAMED_BUILDER(qtensorAlloc)}, - QCOTestCase{"QTensorInsertSliceExtractSlice", - MQT_NAMED_BUILDER(qtensorInsertSliceExtractSlice), + 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), diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index f8f84345cb..1ef16df6d0 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -2175,29 +2175,42 @@ void qtensorInsertSlice(QCOProgramBuilder& b) { b.qtensorInsertSlice(insertOutTensor, extractSliceOutTensor, 0, 2); } -void qtensorExtractInsert(QCOProgramBuilder& b) { +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 qtensorExtractSliceInsertSlice(QCOProgramBuilder& b) { +void qtensorExtractSliceInsertSliceOffsetMismatch(QCOProgramBuilder& b) { auto qtensor = b.qtensorAlloc(3); auto [extractSliceOutTensor, slicedTensor] = b.qtensorExtractSlice(qtensor, 0, 2); - b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 0, 2); + b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 1, 2); } -void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b) { +void qtensorExtractSliceInsertSliceSameOffset(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); + b.qtensorInsertSlice(slicedTensor, extractSliceOutTensor, 0, 2); } -void qtensorInsertExtract(QCOProgramBuilder& b) { +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); @@ -2206,7 +2219,21 @@ void qtensorInsertExtract(QCOProgramBuilder& b) { b.qtensorInsert(q2, extractOutTensor1, 0); } -void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b) { +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); @@ -2220,4 +2247,13 @@ void qtensorInsertSliceExtractSlice(QCOProgramBuilder& b) { 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 e3676e3313..ba26e42969 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -998,21 +998,39 @@ void qtensorExtractSlice(QCOProgramBuilder& b); /// Inserts a slice into a tensor. void qtensorInsertSlice(QCOProgramBuilder& b); -/// Extracts a qubit from a tensor and inserts it immediately. -void qtensorExtractInsert(QCOProgramBuilder& b); +/// Extracts a qubit from a tensor and inserts it immediately at a different +/// index. +void qtensorExtractInsertIndexMismatch(QCOProgramBuilder& b); -/// Inserts a qubit into a tensor and extracts it immediately. -void qtensorInsertExtract(QCOProgramBuilder& b); +/// Extracts a qubit from a tensor and inserts it immediately at the same index. +void qtensorExtractInsertSameIndex(QCOProgramBuilder& b); -/// Extracts a slice of qubits from a tensor and inserts it immediately. -void qtensorExtractSliceInsertSlice(QCOProgramBuilder& b); +/// Inserts a qubit into a tensor and extracts it immediately at a different +/// index. +void qtensorInsertExtractIndexMismatch(QCOProgramBuilder& b); -/// Inserts a slice of qubits into a tensor and extracts it immediately. -void qtensorInsertSliceExtractSlice(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. +/// immediately at the same index and offset. void qtensorExtractSliceExtractInsertInsertSlice(QCOProgramBuilder& b); } // namespace mlir::qco From 2e7195b592d99b7f1c7ab2d4bd3c4b5908e1061c Mon Sep 17 00:00:00 2001 From: Daniel Haag <121057143+denialhaag@users.noreply.github.com> Date: Thu, 19 Mar 2026 00:48:32 +0100 Subject: [PATCH 105/108] Address the Rabbit's comments --- .../include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 9 +++++---- mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 5ae5aa1193..f2cb7e4cfa 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -361,9 +361,10 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * @brief Explicitly deallocate a tensor * * @details - * Explicitly deallocates the given tensor and the value it holds. Qubits or - * tensors of qubits that were extracted from the tensor but not inserted back - * again need to be deallocated separately. + * 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 @@ -1230,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) diff --git a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp index d2601878c8..1a7b6526ab 100644 --- a/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp +++ b/mlir/lib/Dialect/QTensor/IR/Operations/InsertSliceOp.cpp @@ -41,7 +41,7 @@ LogicalResult InsertSliceOp::verify() { } if (constOffset && constSize && !ShapedType::isDynamic(dstDim)) { - if (*constOffset + *constSize > dstDim) { + if (*constSize > dstDim || *constOffset > dstDim - *constSize) { return emitOpError("Offset + Size exceeds destination dimension"); } } From 83c936a970e24a252a565bf55eb771bb95895cfb Mon Sep 17 00:00:00 2001 From: Daniel Haag <121057143+denialhaag@users.noreply.github.com> Date: Thu, 19 Mar 2026 01:04:52 +0100 Subject: [PATCH 106/108] Move dialect definition to separate TableGen file --- docs/mlir/QTensor.md | 3 +- .../mlir/Dialect/QTensor/IR/CMakeLists.txt | 3 +- .../mlir/Dialect/QTensor/IR/QTensorDialect.td | 30 +++++++++++++++++ .../mlir/Dialect/QTensor/IR/QTensorOps.td | 33 ++++--------------- 4 files changed, 39 insertions(+), 30 deletions(-) create mode 100644 mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td diff --git a/docs/mlir/QTensor.md b/docs/mlir/QTensor.md index ae63cb6eae..b12a912fbf 100644 --- a/docs/mlir/QTensor.md +++ b/docs/mlir/QTensor.md @@ -2,7 +2,6 @@ tocdepth: 3 --- -```{include} Dialects/MLIRQTensorDialect.md - +```{include} Dialects/QTensorDialect.md :heading-offset: 1 ``` diff --git a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt index 1fedc2350e..488b88c1c6 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt @@ -7,4 +7,5 @@ # Licensed under the MIT License add_mlir_dialect(QTensorOps qtensor) -add_mlir_doc(QTensorOps MLIRQTensorDialect Dialects/ -gen-dialect-doc) + +add_mlir_doc(QTensorOps QTensorDialect Dialects/ -gen-dialect-doc) 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..c2a00064ea --- /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.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 797c5d521b..4a90777f1e 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -6,8 +6,11 @@ // // Licensed under the MIT License -#ifndef QTENSOROPS -#define QTENSOROPS +#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" @@ -16,26 +19,6 @@ include "mlir/Interfaces/ShapedOpInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/Interfaces/TilingInterface.td" -//===----------------------------------------------------------------------===// -// Dialect -//===----------------------------------------------------------------------===// - -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"; -} - //===----------------------------------------------------------------------===// // Base Operation Classes //===----------------------------------------------------------------------===// @@ -47,10 +30,6 @@ class QTensorOp traits = []> // Type Constraints //===----------------------------------------------------------------------===// -def QubitType : - Type($_self)">, - "Elements must be of qubit type">; - def Static1DQubitTensor : Type< And<[ 1DTensorOf<[QubitType]>.predicate, @@ -280,4 +259,4 @@ def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ let hasVerifier = 1; } -#endif // QTENSOROPS +#endif // MLIR_DIALECT_QTENSOR_IR_QTENSOROPS_TD From 755a2ce2bac87b7cd0ba291e98d3bf03ea3cbc64 Mon Sep 17 00:00:00 2001 From: Daniel Haag <121057143+denialhaag@users.noreply.github.com> Date: Thu, 19 Mar 2026 01:45:20 +0100 Subject: [PATCH 107/108] Fix generation of QTensorDialect.md --- .../mlir/Dialect/QTensor/IR/CMakeLists.txt | 2 +- .../include/mlir/Dialect/QTensor/IR/QTensorOps.td | 15 +++++++-------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt index 488b88c1c6..0d66effe7c 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/QTensor/IR/CMakeLists.txt @@ -8,4 +8,4 @@ add_mlir_dialect(QTensorOps qtensor) -add_mlir_doc(QTensorOps QTensorDialect Dialects/ -gen-dialect-doc) +add_mlir_doc(QTensorOps QTensorDialect Dialects/ -gen-dialect-doc -dialect=qtensor) diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 4a90777f1e..1442b9c2f8 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -41,8 +41,7 @@ def Static1DQubitTensor : Type< // Operations //===----------------------------------------------------------------------===// -def QTensor_AllocOp : QTensorOp<"alloc", [ - MemoryEffects<[MemAlloc]>]> { +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. @@ -66,7 +65,7 @@ def QTensor_AllocOp : QTensorOp<"alloc", [ let hasVerifier = 1; } -def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { +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. @@ -84,7 +83,7 @@ def QTensor_DeallocOp : QTensorOp<"dealloc", [MemoryEffects<[MemFree]>]> { let hasCanonicalizer = 1; } -def QTensor_FromElementsOp : QTensorOp<"from_elements", [ +def FromElementsOp : QTensorOp<"from_elements", [ Pure, TypesMatchWith<"operand types match result element type", "result", "elements", "SmallVector(" @@ -113,7 +112,7 @@ def QTensor_FromElementsOp : QTensorOp<"from_elements", [ ]; } -def QTensor_ExtractOp : QTensorOp<"extract", [ +def ExtractOp : QTensorOp<"extract", [ Pure, TypesMatchWith<"result type matches element type of tensor", "tensor", "result", @@ -140,7 +139,7 @@ def QTensor_ExtractOp : QTensorOp<"extract", [ let hasVerifier = 1; } -def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ +def ExtractSliceOp : QTensorOp<"extract_slice", [ Pure, TypesMatchWith<"returned tensor type matches input tensor", "tensor", "out_tensor", "$_self"> @@ -186,7 +185,7 @@ def QTensor_ExtractSliceOp : QTensorOp<"extract_slice", [ let hasVerifier = 1; } -def QTensor_InsertOp : QTensorOp<"insert", [ +def InsertOp : QTensorOp<"insert", [ Pure, TypesMatchWith<"result type matches type of dest", "dest", "result", @@ -218,7 +217,7 @@ def QTensor_InsertOp : QTensorOp<"insert", [ let hasVerifier = 1; } -def QTensor_InsertSliceOp : QTensorOp<"insert_slice", [ +def InsertSliceOp : QTensorOp<"insert_slice", [ Pure, TypesMatchWith<"expected result type to match dest type", "dest", "result", "$_self"> From 86d631b97df5812b79fb109cb8be994f4922fd9f Mon Sep 17 00:00:00 2001 From: Daniel Haag <121057143+denialhaag@users.noreply.github.com> Date: Thu, 19 Mar 2026 02:14:46 +0100 Subject: [PATCH 108/108] Address the Rabbit's new comments --- mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h | 8 ++++---- mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td | 6 +++--- mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td | 4 ++-- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 1 + 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index f2cb7e4cfa..a09c49ea7a 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1300,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 */ diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td index c2a00064ea..75122363fa 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorDialect.td @@ -17,9 +17,9 @@ def QTensorDialect : Dialect { 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. + 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"]; diff --git a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td index 1442b9c2f8..9322eace45 100644 --- a/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td +++ b/mlir/include/mlir/Dialect/QTensor/IR/QTensorOps.td @@ -93,8 +93,8 @@ def FromElementsOp : QTensorOp<"from_elements", [ 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. The also causes - the input qubits to be consumed. + 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 diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 0d2a79e499..ad4981d1b7 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -77,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());