✨ Add pass and patterns for Classical Control Replacement#1756
✨ Add pass and patterns for Classical Control Replacement#1756DRovara wants to merge 24 commits into
Conversation
…toolkit/core into mlir/measurement-lifting
|
@CodeRabbit full review |
✅ Actions performedFull review triggered. |
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughThis PR introduces dead gate elimination patterns and two circuit optimization passes (measurement lifting and classical control replacement) for the QCO quantum dialect. Changes span operation trait annotations, builder API extensions, canonicalization patterns, optimization pass implementations, and comprehensive test coverage. ChangesDead Gate Elimination and Canonicalization
Builder API Extensions
Measurement and Control-Flow Optimization Passes
Test Infrastructure and Documentation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches✨ Simplify code
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Infer (1.2.0)mlir/lib/Dialect/QCO/IR/Operations/MeasureOp.cppmlir/lib/Dialect/QCO/IR/Operations/MeasureOp.cpp:11:10: fatal error: 'mlir/Dialect/QCO/IR/QCOOps.h' file not found ... [truncated 1160 characters] ... ude" mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cppmlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp:11:10: fatal error: 'mlir/Dialect/QCO/Builder/QCOProgramBuilder.h' file not found ... [truncated 1220 characters] ... e" mlir/lib/Dialect/QCO/IR/Operations/ResetOp.cppmlir/lib/Dialect/QCO/IR/Operations/ResetOp.cpp:11:10: fatal error: 'mlir/Dialect/QCO/IR/QCOOps.h' file not found ... [truncated 2200 characters] ... led from ClangFrontend__CTrans.CTrans_funct.instruction_log.(fun) in file "src/clang/cTrans.ml", line 4784, characters 10-1023
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. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
mlir/include/mlir/Dialect/QCO/IR/QCOOps.td (1)
1071-1077:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDrop
Purefromqco.ctrl,qco.inv, andqco.if(region wrappers).
Purein MLIR expands toAlwaysSpeculatable+NoMemoryEffect(viaAlwaysSpeculatableImplTrait), which asserts speculatability and no memory effects. These wrapper ops already model their effects viaRecursiveMemoryEffects, so keepingPureoverstates effects independently of what’s inside the nested region. KeepRecursiveMemoryEffects, but removePurefromCtrlOp,InvOp, andIfOp(incl. ranges covering the corresponding defs around 1071-1077, 1150-1154, and 1223-1232).🤖 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/include/mlir/Dialect/QCO/IR/QCOOps.td` around lines 1071 - 1077, The QCO region wrapper ops declare the Pure trait which expands to AlwaysSpeculatable + NoMemoryEffect and incorrectly duplicates/overrides the region-modelled effects; remove the Pure trait from the op definitions for CtrlOp, InvOp, and IfOp (the TD defs named CtrlOp, InvOp, IfOp) so they only keep RecursiveMemoryEffects (and other existing traits like UnitaryOpInterface, AttrSizedOperandSegments, AttrSizedResultSegments, SameOperandsAndResultType, SameOperandsAndResultShape, SingleBlockImplicitTerminator) to ensure effects are solely driven by the nested region.mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp (1)
1105-1165:⚠️ Potential issue | 🟠 Major | ⚡ Quick winEnforce the configured
mainreturn signature before finalizing.After
initialize(TypeRange), the no-argfinalize()still hardcodes ani64zero, andfinalize(ValueRange)never checks arity/types againstmain.initialize({})or any non-i64signature now silently builds an invalidfunc.return.Suggested fix
OwningOpRef<ModuleOp> QCOProgramBuilder::finalize() { checkFinalized(); - - auto exitCode = intConstant(0); - return finalize({exitCode}); + func::FuncOp mainFunc = nullptr; + for (auto op : cast<ModuleOp>(module).getOps<func::FuncOp>()) { + if (op.getName() == "main") { + mainFunc = op; + break; + } + } + if (!mainFunc) { + llvm::reportFatalUsageError("Could not find main function"); + } + + auto resultTypes = mainFunc.getFunctionType().getResults(); + if (resultTypes.empty()) { + return finalize({}); + } + if (resultTypes.size() == 1 && resultTypes.front().isSignlessInteger(64)) { + auto exitCode = intConstant(0); + return finalize({exitCode}); + } + llvm::reportFatalUsageError( + "Use finalize(returnValues) when main has custom return types"); } OwningOpRef<ModuleOp> QCOProgramBuilder::finalize(ValueRange returnValues) { checkFinalized(); @@ if (!mainFunc) { llvm::reportFatalUsageError("Could not find main function"); } + + auto resultTypes = mainFunc.getFunctionType().getResults(); + if (returnValues.size() != resultTypes.size()) { + llvm::reportFatalUsageError( + "Return values must match the main function return arity"); + } + for (auto [value, type] : llvm::zip_equal(returnValues, resultTypes)) { + if (value.getType() != type) { + llvm::reportFatalUsageError( + "Return values must match the main function return types"); + } + } @@ func::ReturnOp::create(*this, returnValues);🤖 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/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp` around lines 1105 - 1165, The finalize paths currently emit a func::ReturnOp without validating or matching the main function's declared result types (initialize(TypeRange) may set a non-i64 or different arity), so update QCOProgramBuilder::finalize() and QCOProgramBuilder::finalize(ValueRange) to first query mainFunc.getType().getResults(), assert that returnValues.size() equals that result list size and that each returnValues[i].getType() == expectedResults[i], and if the no-arg finalize() is used produce appropriate default constants per expectedResults type (not always an i64 zero) instead of unconditionally calling intConstant(0); then create the func::ReturnOp with validated/matched values (and emit a clear error via llvm::reportFatalUsageError if arity/types mismatch). Ensure references: finalize(), finalize(ValueRange), initialize(TypeRange), func::ReturnOp::create, mainFunc.getType().getResults(), and intConstant are used when implementing the checks and per-type default constant construction.
🤖 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 `@CHANGELOG.md`:
- Line 14: Replace the incorrect changelog bullet "✨ Add Dead Gate Elimination
Pattern ([`#1755`]) ([**`@DRovara`**])" with an entry that references this PR
(`#1756`) and accurately describes the implemented feature (measurement lifting /
classical-control replacement), e.g., update the PR number to 1756 and change
the description to mention "Measurement lifting and classical-control
replacement" so the released notes match the actual code changes.
In `@mlir/include/mlir/Dialect/QCO/QCOUtils.h`:
- Around line 240-276: The header's inline function checkAndRemoveDeadGate uses
QCO-specific types (SinkOp, UnitaryOpInterface, MeasureOp, IfOp, Reset) but
doesn't include their declarations, making the header order-dependent; fix by
making the header self-contained: add a direct include of the QCO ops/declaring
header (e.g., include "mlir/Dialect/QCO/QCOOps.h") at the top of
mlir/include/mlir/Dialect/QCO/QCOUtils.h so dyn_cast/isa and method calls
resolve, or alternatively move the implementation of checkAndRemoveDeadGate out
of the header into a .cpp that includes QCOOps.h if you prefer to avoid
additional header dependencies.
In `@mlir/lib/Dialect/QCO/Transforms/Optimizations/ReplaceClassicalControls.cpp`:
- Around line 143-155: The lambda passed to IfOp::create currently always builds
a CtrlOp (via CtrlOp::create) even when newControls (derived from toKeep.size())
is empty, which yields invalid IR because qco::CtrlOp requires >=1 control;
modify the lambda to branch on newControls.empty(): if non-empty, keep the
existing flow (create CtrlOp, inline newCtrl.getRegion() into op.getRegion(),
return newCtrl.getOutputQubits()); if empty, instead directly clone or move the
original unitary body into the IfOp region (use rewriter to inline/clone the
original region/body into the new IfOp region) and return the appropriate output
qubits, ensuring no zero-control CtrlOp is constructed (affects
ReplaceClassicalControls.cpp paths like replaceClassicalControlsOnlyControl and
replaceClassicalControlsTwoOfTwoControls).
- Around line 58-63: The isDiagonal predicate omits RzOp so controlled-rz paths
are not recognized; update the isDiagonal function (the dyn_cast<InvOp> branch
and the isa<> return) to include RzOp among the recognized diagonal ops (i.e.,
add RzOp to the isa<ZOp, SOp, TOp, POp, SdgOp, TdgOp, IdOp> list) so
controlled-`rz` will take the diagonal-swap path.
In
`@mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_replace_classical_controls.cpp`:
- Around line 377-380: Remove the leftover debug printing by deleting the
module->dump() call at the end of the test in
test_qco_replace_classical_controls.cpp; keep the existing assertions
runReplaceClassicalControlsPass(module.get()).succeeded() and
runCanonicalizerPass(reference.get()).succeeded() but remove the module->dump()
line so the test no longer prints the transformed module to CI logs.
---
Outside diff comments:
In `@mlir/include/mlir/Dialect/QCO/IR/QCOOps.td`:
- Around line 1071-1077: The QCO region wrapper ops declare the Pure trait which
expands to AlwaysSpeculatable + NoMemoryEffect and incorrectly
duplicates/overrides the region-modelled effects; remove the Pure trait from the
op definitions for CtrlOp, InvOp, and IfOp (the TD defs named CtrlOp, InvOp,
IfOp) so they only keep RecursiveMemoryEffects (and other existing traits like
UnitaryOpInterface, AttrSizedOperandSegments, AttrSizedResultSegments,
SameOperandsAndResultType, SameOperandsAndResultShape,
SingleBlockImplicitTerminator) to ensure effects are solely driven by the nested
region.
In `@mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp`:
- Around line 1105-1165: The finalize paths currently emit a func::ReturnOp
without validating or matching the main function's declared result types
(initialize(TypeRange) may set a non-i64 or different arity), so update
QCOProgramBuilder::finalize() and QCOProgramBuilder::finalize(ValueRange) to
first query mainFunc.getType().getResults(), assert that returnValues.size()
equals that result list size and that each returnValues[i].getType() ==
expectedResults[i], and if the no-arg finalize() is used produce appropriate
default constants per expectedResults type (not always an i64 zero) instead of
unconditionally calling intConstant(0); then create the func::ReturnOp with
validated/matched values (and emit a clear error via llvm::reportFatalUsageError
if arity/types mismatch). Ensure references: finalize(), finalize(ValueRange),
initialize(TypeRange), func::ReturnOp::create, mainFunc.getType().getResults(),
and intConstant are used when implementing the checks and per-type default
constant construction.
🪄 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: 404b26d5-75c2-4dfc-86eb-990ccae1574b
📒 Files selected for processing (17)
CHANGELOG.mdmlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.hmlir/include/mlir/Dialect/QCO/IR/QCODialect.tdmlir/include/mlir/Dialect/QCO/IR/QCOOps.tdmlir/include/mlir/Dialect/QCO/QCOUtils.hmlir/include/mlir/Dialect/QCO/Transforms/Passes.tdmlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cppmlir/lib/Dialect/QCO/IR/Operations/MeasureOp.cppmlir/lib/Dialect/QCO/IR/Operations/ResetOp.cppmlir/lib/Dialect/QCO/IR/QCOOps.cppmlir/lib/Dialect/QCO/IR/SCF/IfOp.cppmlir/lib/Dialect/QCO/Transforms/Optimizations/MeasurementLifting.cppmlir/lib/Dialect/QCO/Transforms/Optimizations/ReplaceClassicalControls.cppmlir/unittests/Dialect/QCO/IR/test_qco_ir.cppmlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txtmlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_measurement_lifting.cppmlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_replace_classical_controls.cpp
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Description
Part 2 out of 3 of the Qubit Reuse Implementation
This PR ports the previous implementation of the Classical Control Replacement pass to the new dialect infrastructure.
This path has the goal of replacing controlled gates with "classically controlled gates" whenever a classical value for the control qubit is available because it has just been measured.
In those cases, we remove the control from the
CtrlOpand instead introduce aqco.ifOp.This is one of 3 PRs that will, in total, constitute qubit reuse:
Requires #1705 to be merged first.
AI Notice: Gemini 3.1 and Claude Sonnet 4.6 were used for the generation of some code.
Checklist
If PR contains AI-assisted content:
Assisted-by: [Model Name] via [Tool Name]footer.