✨ Add TensorIterator

CHANGELOG.md

@@ -18,7 +18,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 - ✨ Add conversions between `jeff` and QCO ([#1479], [#1548], [#1565], [#1637], [#1676], [#1706]) ([**@denialhaag**], [**@burgholzer**])
 - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716]) ([**@MatthiasReumann**], [**@burgholzer**])
 - ✨ 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], [#1542], [#1548], [#1550], [#1554], [#1567], [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1717])
+  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1567], [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1717], [#1730])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
 
 ### Changed
@@ -403,6 +403,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 <!-- PR links -->
 
 [#1737]: https://github.com/munich-quantum-toolkit/core/pull/1737
+[#1730]: https://github.com/munich-quantum-toolkit/core/pull/1730
 [#1720]: https://github.com/munich-quantum-toolkit/core/pull/1720
 [#1719]: https://github.com/munich-quantum-toolkit/core/pull/1719
 [#1718]: https://github.com/munich-quantum-toolkit/core/pull/1718

mlir/include/mlir/Dialect/QTensor/Utils/TensorIterator.h

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+ * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+ * All rights reserved.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Licensed under the MIT License
+ */
+
+#pragma once
+
+#include <mlir/IR/Builders.h>
+#include <mlir/IR/BuiltinTypes.h>
+#include <mlir/IR/Operation.h>
+#include <mlir/IR/Value.h>
+
+#include <iterator>
+
+namespace mlir::qtensor {
+
+/**
+ * @brief A bidirectional_iterator traversing the tensor chain.
+ **/
+class [[nodiscard]] TensorIterator {
+public:
+  using iterator_category = std::bidirectional_iterator_tag;
+  using difference_type = std::ptrdiff_t;
+  using value_type = Operation*;
+
+  TensorIterator() : op_(nullptr), tensor_(nullptr), isSentinel_(false) {}
+  explicit TensorIterator(TypedValue<RankedTensorType> tensor)
+      : op_(tensor.getDefiningOp()), tensor_(tensor), isSentinel_(false) {}
+
+  /// @returns the operation the iterator points to.
+  [[nodiscard]] Operation* operation() const { return op_; }
+
+  /// @returns the operation the iterator points to.
+  [[nodiscard]] Operation* operator*() const { return operation(); }
+
+  /// @returns the tensor the iterator points to.
+  [[nodiscard]] TypedValue<RankedTensorType> tensor() const;
+
+  TensorIterator& operator++() {
+    forward();
+    return *this;
+  }
+
+  TensorIterator operator++(int) {
+    auto tmp = *this;
+    operator++();
+    return tmp;
+  }
+
+  TensorIterator& operator--() {
+    backward();
+    return *this;
+  }
+
+  TensorIterator operator--(int) {
+    auto tmp = *this;
+    operator--();
+    return tmp;
+  }
+
+  bool operator==(const TensorIterator& other) const {
+    return other.tensor_ == tensor_ && other.op_ == op_ &&
+           other.isSentinel_ == isSentinel_;
+  }
+
+  bool operator==([[maybe_unused]] std::default_sentinel_t s) const {
+    return isSentinel_;
+  }
+
+private:
+  /// @brief Move to the next operation on the tensor def-use chain.
+  void forward();
+
+  /// @brief Move to the previous operation on the tensor def-use chain.
+  void backward();
+
+  Operation* op_;
+  TypedValue<RankedTensorType> tensor_;
+  bool isSentinel_;
+};
+} // namespace mlir::qtensor

mlir/lib/Dialect/QCO/Utils/WireIterator.cpp

@@ -39,8 +39,8 @@ void WireIterator::forward() {
   }
 
   // Find the user-operation of the qubit SSA value.
-  assert(qubit_.getNumUses() == 1 && "expected linear typing");
-  op_ = *(qubit_.getUsers().begin());
+  assert(qubit_.hasOneUse() && "expected linear typing");
+  op_ = *(qubit_.user_begin());
 
   // A sink/insert defines the end of the qubit wire (dynamic and static).
   if (isa<SinkOp, YieldOp, qtensor::InsertOp>(op_)) {
@@ -57,8 +57,8 @@ void WireIterator::forward() {
         .Case<MeasureOp>([&](MeasureOp op) { qubit_ = op.getQubitOut(); })
         .Case<ResetOp>([&](ResetOp op) { qubit_ = op.getQubitOut(); })
         .Default([&](Operation* op) {
-          report_fatal_error("unknown op in def-use chain: " +
-                             op->getName().getStringRef());
+          llvm::reportFatalInternalError("unknown op in def-use chain: " +
+                                         op->getName().getStringRef());
         });
   }
 }
@@ -90,8 +90,8 @@ void WireIterator::backward() {
       .Case<MeasureOp>([&](MeasureOp op) { qubit_ = op.getQubitIn(); })
       .Case<ResetOp>([&](ResetOp op) { qubit_ = op.getQubitIn(); })
       .Default([&](Operation* op) {
-        report_fatal_error("unknown op in def-use chain: " +
-                           op->getName().getStringRef());
+        llvm::reportFatalInternalError("unknown op in def-use chain: " +
+                                       op->getName().getStringRef());
       });
 
   // Get the operation that produces the qubit value.

mlir/lib/Dialect/QTensor/CMakeLists.txt

@@ -8,3 +8,4 @@
 
 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(Utils)

mlir/lib/Dialect/QTensor/Utils/CMakeLists.txt

@@ -0,0 +1,44 @@
+# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+# All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+#
+# Licensed under the MIT License
+
+file(GLOB_RECURSE UTILS_CPP "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
+
+add_mlir_dialect_library(
+  MLIRQTensorUtils
+  ${UTILS_CPP}
+  ADDITIONAL_HEADER_DIRS
+  ${PROJECT_SOURCE_DIR}/mlir/include/mlir/Dialect/QTensor
+  DEPENDS
+  MLIRQTensorOpsIncGen
+  LINK_LIBS
+  PUBLIC
+  MLIRQTensorDialect)
+
+mqt_mlir_target_use_project_options(MLIRQTensorUtils)
+
+# collect header files
+file(GLOB_RECURSE UTILS_HEADERS_SOURCE
+     "${MQT_MLIR_SOURCE_INCLUDE_DIR}/mlir/Dialect/QTensor/Utils/*.h")
+file(GLOB_RECURSE UTILS_HEADERS_BUILD
+     "${MQT_MLIR_BUILD_INCLUDE_DIR}/mlir/Dialect/QTensor/Utils/*.inc")
+
+# add public headers using file sets
+target_sources(
+  MLIRQTensorUtils
+  PUBLIC FILE_SET
+         HEADERS
+         BASE_DIRS
+         ${MQT_MLIR_SOURCE_INCLUDE_DIR}
+         FILES
+         ${UTILS_HEADERS_SOURCE}
+         FILE_SET
+         HEADERS
+         BASE_DIRS
+         ${MQT_MLIR_BUILD_INCLUDE_DIR}
+         FILES
+         ${UTILS_HEADERS_BUILD})

mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp

@@ -0,0 +1,144 @@
+/*
+ * 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/Utils/TensorIterator.h"
+
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QTensor/IR/QTensorOps.h"
+
+#include <llvm/ADT/STLExtras.h>
+#include <llvm/ADT/TypeSwitch.h>
+#include <llvm/Support/ErrorHandling.h>
+#include <mlir/Dialect/SCF/IR/SCF.h>
+#include <mlir/IR/Builders.h>
+#include <mlir/IR/Value.h>
+#include <mlir/Support/LLVM.h>
+
+#include <cassert>
+#include <iterator>
+
+namespace mlir::qtensor {
+TypedValue<RankedTensorType> TensorIterator::tensor() const {
+  if (op_ == nullptr) {
+    return tensor_;
+  }
+
+  // The following operations don't have an OpResult.
+  if (isa<DeallocOp, scf::YieldOp, qco::YieldOp>(op_)) {
+    return nullptr;
+  }
+
+  return tensor_;
+}
+
+void TensorIterator::forward() {
+  // If the iterator is a sentinel already, there is nothing to do.
+  if (isSentinel_) {
+    return;
+  }
+
+  // Find the user-operation of the tensor SSA value.
+  assert(tensor_.hasOneUse() && "expected linear typing");
+  op_ = *(tensor_.user_begin());
+
+  // The following operations define the end of the tensor's life-chain.
+  if (isa<DeallocOp, scf::YieldOp, qco::YieldOp>(op_)) {
+    isSentinel_ = true;
+    return;
+  }
+
+  // Find the output from the input tensor SSA value.
+  if (!(isa<AllocOp, FromElementsOp>(op_))) {
+    TypeSwitch<Operation*>(op_)
+        .Case<ExtractOp>([&](ExtractOp op) { tensor_ = op.getOutTensor(); })
+        .Case<InsertOp>([&](InsertOp op) { tensor_ = op.getResult(); })
+        .Case<scf::ForOp>([&](scf::ForOp op) {
+          tensor_ = cast<TypedValue<RankedTensorType>>(
+              op.getTiedLoopResult(&*(tensor_.use_begin())));
+        })
+        .Case<qco::IfOp>([&](qco::IfOp op) {
+          auto it = llvm::find(op.getQubits(), tensor_);
+          assert(it != op.getQubits().end());
+          const auto idx = std::distance(op.getQubits().begin(), it);
+          tensor_ = cast<TypedValue<RankedTensorType>>(op.getResults()[idx]);
+        })
+        .Default([&](Operation* op) {
+          report_fatal_error("unknown op in def-use chain: " +
+                             op->getName().getStringRef());
+        });
+  }
+}
+
+void TensorIterator::backward() {
+  // If the iterator is a sentinel, reactivate the iterator.
+  if (isSentinel_) {
+    isSentinel_ = false;
+    return;
+  }
+
+  // If the op is a nullptr, the tensor value is a block argument and thus the
+  // beginning of the tensor's life-chain.
+  if (op_ == nullptr) {
+    return;
+  }
+
+  // For these operations, tensor_ is an OpOperand. Hence, only get the def-op.
+  if (isa<DeallocOp, scf::YieldOp, qco::YieldOp>(op_)) {
+    op_ = tensor_.getDefiningOp();
+    return;
+  }
+
+  // Allocations and FromElements define the start of the tensor's life-chain.
+  // Consequently, stop and early exit.
+  if (isa<AllocOp, FromElementsOp>(op_)) {
+    return;
+  }
+
+  // Find the input from the output tensor SSA value.
+  TypeSwitch<Operation*>(op_)
+      .Case<ExtractOp>([&](ExtractOp op) { tensor_ = op.getTensor(); })
+      .Case<InsertOp>([&](InsertOp op) { tensor_ = op.getDest(); })
+      .Case<scf::ForOp>([&](scf::ForOp op) {
+        if (auto res = dyn_cast<OpResult>(tensor_)) {
+          OpOperand* operand = op.getTiedLoopInit(res);
+          tensor_ = cast<TypedValue<RankedTensorType>>(operand->get());
+          return;
+        }
+
+        llvm::reportFatalInternalError(
+            "expected scf.for result for tied init lookup");
+      })
+      .Case<qco::IfOp>([&](qco::IfOp op) {
+        if (auto res = dyn_cast<OpResult>(tensor_)) {
+          auto it = llvm::find(op.getResults(), res);
+          assert(it != op->result_end());
+          const auto idx = std::distance(op.result_begin(), it);
+          tensor_ = cast<TypedValue<RankedTensorType>>(op.getQubits()[idx]);
+          return;
+        }
+
+        llvm::reportFatalInternalError(
+            "expected scf.for result for tied init lookup");
+      })
+      .Default([&](Operation* op) {
+        llvm::reportFatalInternalError("unknown op in def-use chain: " +
+                                       op->getName().getStringRef());
+      });
+
+  // Get the operation that produces the tensor value.
+  // If the current tensor SSA value is a BlockArgument (no defining op), the
+  // operation will be a nullptr.
+  op_ = tensor_.getDefiningOp();
+}
+
+static_assert(std::bidirectional_iterator<TensorIterator>);
+static_assert(std::sentinel_for<std::default_sentinel_t, TensorIterator>,
+              "std::default_sentinel_t must be a sentinel for TensorIterator.");
+} // namespace mlir::qtensor

mlir/unittests/Dialect/QTensor/CMakeLists.txt

@@ -7,3 +7,4 @@
 # Licensed under the MIT License
 
 add_subdirectory(IR)
+add_subdirectory(Utils)

mlir/unittests/Dialect/QTensor/Utils/CMakeLists.txt

@@ -0,0 +1,16 @@
+# 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
+
+set(qtensor_utils_target mqt-core-mlir-unittest-qtensor-utils)
+add_executable(${qtensor_utils_target} test_tensoriterator.cpp)
+target_link_libraries(${qtensor_utils_target} PRIVATE GTest::gtest_main MLIRQTensorDialect
+                                                      MLIRQTensorUtils MLIRQCOProgramBuilder)
+mqt_mlir_configure_unittest_target(${qtensor_utils_target})
+
+gtest_discover_tests(${qtensor_utils_target} PROPERTIES LABELS mqt-mlir-unittests DISCOVERY_TIMEOUT
+                                                        60)