From eb637f6231cda49d9978a20a110dd26a347be450 Mon Sep 17 00:00:00 2001
From: Simon Hofmann <simon.t.hofmann@tum.de>
Date: Tue, 1 Sep 2026 00:14:03 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Reject=20nonlinear=20QTensor=20s?=
 =?UTF-8?q?hrink=20chains=20safely?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../QTensor/Transforms/ShrinkRegisters.cpp    |  3 ++
 .../Transforms/test_qtensor_transforms.cpp    | 40 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/mlir/lib/Dialect/QTensor/Transforms/ShrinkRegisters.cpp b/mlir/lib/Dialect/QTensor/Transforms/ShrinkRegisters.cpp
index 550d37583c..0d79fc2cea 100644
--- a/mlir/lib/Dialect/QTensor/Transforms/ShrinkRegisters.cpp
+++ b/mlir/lib/Dialect/QTensor/Transforms/ShrinkRegisters.cpp
@@ -62,6 +62,9 @@ struct TensorAccess {
     SmallVectorImpl<TensorAccess>& accesses, DeallocOp& deallocOp) {
   auto tensor = allocOp.getResult();
   while (true) {
+    if (!tensor.hasOneUse()) {
+      return failure();
+    }
     auto* user = *tensor.getUsers().begin();
 
     if (auto currentDealloc = dyn_cast<DeallocOp>(user)) {
diff --git a/mlir/unittests/Dialect/QTensor/Transforms/test_qtensor_transforms.cpp b/mlir/unittests/Dialect/QTensor/Transforms/test_qtensor_transforms.cpp
index b99dae7351..259bbd7f15 100644
--- a/mlir/unittests/Dialect/QTensor/Transforms/test_qtensor_transforms.cpp
+++ b/mlir/unittests/Dialect/QTensor/Transforms/test_qtensor_transforms.cpp
@@ -20,6 +20,7 @@
 #include "mlir/Dialect/QTensor/Transforms/Passes.h"
 
 #include <gtest/gtest.h>
+#include <llvm/Support/raw_ostream.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/IR/BuiltinAttributes.h>
@@ -34,6 +35,7 @@
 #include <mlir/Support/LLVM.h>
 
 #include <cstdint>
+#include <string>
 
 using namespace mlir;
 
@@ -118,4 +120,42 @@ TEST(QTensorTransformsTest, HugeDeclaredTensorUsesSparseShrinkPlan) {
             ArrayRef<int64_t>{1});
 }
 
+TEST(QTensorTransformsTest, RejectsNonLinearChainWithoutMutation) {
+  DialectRegistry registry;
+  registry.insert<arith::ArithDialect, func::FuncDialect, mqt::MQTDialect,
+                  qco::QCODialect, qtensor::QTensorDialect>();
+  MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+
+  auto moduleOp = parseSourceString<ModuleOp>(R"mlir(
+    module {
+      func.func @main() {
+        %c1 = arith.constant 1 : index
+        %c3 = arith.constant 3 : index
+        %reg = qtensor.alloc(%c3) : tensor<3x!qco.qubit>
+        %rest, %qubit = qtensor.extract %reg[%c1]
+            : tensor<3x!qco.qubit>
+        qtensor.dealloc %rest : tensor<3x!qco.qubit>
+        qtensor.dealloc %rest : tensor<3x!qco.qubit>
+        return
+      }
+    }
+  )mlir",
+                                              &context);
+  ASSERT_TRUE(moduleOp);
+  ASSERT_TRUE(succeeded(verify(*moduleOp)));
+
+  std::string before;
+  llvm::raw_string_ostream(before) << *moduleOp;
+
+  PassManager manager(&context);
+  manager.addPass(qtensor::createShrinkQTensorToFitPass());
+  ASSERT_TRUE(succeeded(manager.run(*moduleOp)));
+  ASSERT_TRUE(succeeded(verify(*moduleOp)));
+
+  std::string after;
+  llvm::raw_string_ostream(after) << *moduleOp;
+  EXPECT_EQ(after, before);
+}
+
 } // namespace