Skip to content

🚸 Remove One-Tensor Constraint From Mapping Pass#1748

Merged
burgholzer merged 10 commits into
mainfrom
enh/lift-one-tensor-constraint-mapping
May 28, 2026
Merged

🚸 Remove One-Tensor Constraint From Mapping Pass#1748
burgholzer merged 10 commits into
mainfrom
enh/lift-one-tensor-constraint-mapping

Conversation

@MatthiasReumann
Copy link
Copy Markdown
Collaborator

@MatthiasReumann MatthiasReumann commented May 27, 2026

Description

This pull request is a follow-up of #1716 and lifts the assumption that the quantum computation consists of a single quantum tensor. By doing so, we support programs such as the qpeexact ones generated by MQT-Bench.

Changes

  • 🚸 Improve usability by supporting multiple tensors in the mapping pass

Checklist

  • The pull request only contains commits that are focused and relevant to this change.
  • I have added appropriate tests that cover the new/changed functionality.
  • I have updated the documentation to reflect these changes.
  • I have added entries to the changelog for any noteworthy additions, changes, fixes, or removals.
  • I have added migration instructions to the upgrade guide (if needed).
  • The changes follow the project's style guidelines and introduce no new warnings.
  • The changes are fully tested and pass the CI checks.
  • I have reviewed my own code changes.

If PR contains AI-assisted content:

  • I have disclosed the use of AI tools in the PR description as per our AI Usage Guidelines.
  • AI-assisted commits include an Assisted-by: [Model Name] via [Tool Name] footer.
  • I confirm that I have personally reviewed and understood all AI-generated content, and accept full responsibility for it.

@MatthiasReumann MatthiasReumann added this to the MLIR Support milestone May 27, 2026
@MatthiasReumann MatthiasReumann self-assigned this May 27, 2026
@MatthiasReumann MatthiasReumann added c++ Anything related to C++ code MLIR Anything related to MLIR labels May 27, 2026
@MatthiasReumann
Copy link
Copy Markdown
Collaborator Author

@coderabbitai full review

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 27, 2026

✅ Actions performed

Full review triggered.

@codecov
Copy link
Copy Markdown

codecov Bot commented May 27, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 27, 2026

Review Change Stack

📝 Walkthrough

Walkthrough

TensorIterator gains an isFinal_ state flag to track when iteration reaches terminal operations (dealloc/yield). The mapping pass is refactored to handle multiple qtensor::AllocOp regions independently using TensorIterator, replacing single-chain assumptions. Tests are updated to validate multi-tensor processing and enforce extract-before-insert ordering.

Changes

TensorIterator and Mapping Pass Multi-AllocOp Support

Layer / File(s) Summary
TensorIterator final-state tracking
mlir/include/mlir/Dialect/QTensor/Utils/TensorIterator.h, mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp, mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp
TensorIterator adds a private isFinal_ member initialized to false in both default and tensor-based constructors. The forward() method checks isFinal_ and transitions to sentinel state before following the SSA chain; the backward() method sets isFinal_ when reactivating from sentinel. Unit tests validate forward and backward traversal with the new state transitions.
Mapping pass multi-allocop iteration and placement
mlir/lib/Dialect/QCO/Transforms/CMakeLists.txt, mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp
CMake build adds MLIRQTensorUtils dependency. getComputation iterates over all qtensor::AllocOp operations using TensorIterator to collect extracted qubits while enforcing extract-before-insert ordering. place processes each allocation via make_early_inc_range with TypeSwitch dispatch for ExtractOp, InsertOp, and DeallocOp, then erases the processed allocation.
Mapping pass test cases for multi-tensor support
mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp
Negative test NoTwoTensors replaced with NoExtractAfterInsert to validate that mixed extract-insert ordering is rejected. The Sabre positive test refactored to construct two separate tensors (tensorUp with 4 qubits, tensorDown with 2 qubits) instead of a single 6-qubit tensor, and to extract and reinsert qubits from their respective source tensors.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

  • munich-quantum-toolkit/core#1730: Introduces the initial TensorIterator implementation that this PR extends with isFinal_ state tracking and sentinel-transition logic.
  • munich-quantum-toolkit/core#1716: Modifies MappingPass::getComputation and placement logic with allocop-count validation that this PR refactors to support multiple allocations via TensorIterator.

Suggested labels

feature, usability

Suggested reviewers

  • burgholzer

Poem

🐰 The tensors now iterate with isFinal_ grace,
Multiple allocops dance through each place,
Extract before insert, the order is tight,
TensorIterator guides the mapping just right! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 36.36% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and specifically summarizes the main change: removing the one-tensor constraint from the mapping pass, which is the primary objective of this PR.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description check ✅ Passed The PR description includes a clear summary explaining the change (removing single-tensor assumption), motivation (supporting multiple tensors like qpeexact), and a completed checklist. However, the description lacks specific issue references or details about dependencies.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
✨ Simplify code
  • Create PR with simplified code
  • Commit simplified code in branch enh/lift-one-tensor-constraint-mapping

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp (1)

148-153: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Add a round-trip assertion (rewind then forward) to cover isFinal_ state reset.

These additions check sentinel idempotence, but they still miss the regression where -- from sentinel leaves a stale final-state and the next ++ jumps to sentinel again.

Suggested test additions
   --it;
   ASSERT_EQ(it.operation(), tensor0.getDefiningOp()); // qtensor.alloc
   ASSERT_EQ(it.tensor(), tensor0);

+  ++it;
+  ASSERT_EQ(it.operation(), tensor1.getDefiningOp()); // qtensor.extract
+  ASSERT_EQ(it.tensor(), tensor1);
+
   --it;
   ASSERT_EQ(it.operation(), tensor0.getDefiningOp()); // qtensor.alloc
   ASSERT_EQ(it.tensor(), tensor0);
@@
   --recIt;
   ASSERT_EQ(recIt.operation(), nullptr);
   ASSERT_EQ(recIt.tensor(), tensorElse0);

+  ++recIt;
+  ASSERT_EQ(recIt.operation(), tensorElse1.getDefiningOp()); // qtensor.extract
+  ASSERT_EQ(recIt.tensor(), tensorElse1);
+
   --recIt;
   ASSERT_EQ(recIt.operation(), nullptr);
   ASSERT_EQ(recIt.tensor(), tensorElse0);

Also applies to: 222-223

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp` around lines
148 - 153, The test currently checks that incrementing an iterator at sentinel
is idempotent but misses the regression where decrementing from sentinel leaves
the iterator's internal isFinal_ flag stuck; add a round-trip check that does
--it then ++it and asserts the iterator does not immediately re-enter
std::default_sentinel (use the existing iterator variable it and the ++/--
operators) to verify isFinal_ is reset, and apply the same round-trip assertion
at the other similar spot in the test (the block around the other ++it
assertion).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp`:
- Around line 471-489: The code flattens all qtensor.extract into a single
global wires list and never frees the assigned program indices, so
getComputation()/place() cumulatively count qubits across disjoint tensors;
update the logic to track per-alloc lifetimes and reuse/release prog indices
when the matching qtensor.insert is encountered: iterate allocs from
func.getOps<qtensor::AllocOp>() with TensorIterator/ExtractOp and InsertOp,
maintain a per-alloc active set (or reference-counted mapping) of wires->prog
assignments instead of a single global wires vector, on seeing an InsertOp
remove/release the corresponding prog index(s) so place() and comp->size()
reflect the peak concurrent qubits rather than the sum of all extracts, and
ensure runOnOperation() compares device.nqubits() against the maximum
live-program-qubit count computed this way.

In `@mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp`:
- Around line 47-51: When marking the iterator as a sentinel (setting
isSentinel_ = true) make sure to clear the terminal flag by resetting isFinal_ =
false so rewinding doesn't leave the iterator in a stale "final" state; update
the places that set isSentinel_ (the increment/advance logic/operator++ and the
corresponding rewind/decrement logic/operator--) to also clear isFinal_ when
transitioning states so a subsequent ++/-- behaves correctly.

---

Outside diff comments:
In `@mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp`:
- Around line 148-153: The test currently checks that incrementing an iterator
at sentinel is idempotent but misses the regression where decrementing from
sentinel leaves the iterator's internal isFinal_ flag stuck; add a round-trip
check that does --it then ++it and asserts the iterator does not immediately
re-enter std::default_sentinel (use the existing iterator variable it and the
++/-- operators) to verify isFinal_ is reset, and apply the same round-trip
assertion at the other similar spot in the test (the block around the other ++it
assertion).
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 955d21bd-c373-467e-8aa8-cef6d20949d1

📥 Commits

Reviewing files that changed from the base of the PR and between af09841 and c862449.

📒 Files selected for processing (6)
  • mlir/include/mlir/Dialect/QTensor/Utils/TensorIterator.h
  • mlir/lib/Dialect/QCO/Transforms/CMakeLists.txt
  • mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp
  • mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp
  • mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp
  • mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp

Comment thread mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp
Comment thread mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp
Copy link
Copy Markdown
Member

@burgholzer burgholzer left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Besides one broken docstring, this looks very clean to me!

Comment thread mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp Outdated
@burgholzer burgholzer added the enhancement Improvement of existing feature label May 28, 2026
@burgholzer burgholzer merged commit b3a94a2 into main May 28, 2026
56 of 59 checks passed
@burgholzer burgholzer deleted the enh/lift-one-tensor-constraint-mapping branch May 28, 2026 10:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

c++ Anything related to C++ code enhancement Improvement of existing feature MLIR Anything related to MLIR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants