Migrate quantum dialect to new one-shot bufferization (#1686)

**Context:**
This work is based on #1027 

As part of the mlir update, the bufferization of the custom catalyst
dialects need to migrate to the new one-shot bufferization interface, as
opposed to the old pattern-rewrite style bufferization passes. See more
context in #1027.

As an example, here is how the new bufferization interface is used for
mlir's core `arith` dialect:

https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h

https://github.com/llvm/llvm-project/blob/7ee0097b486b31be8b9a1750b2cd47580efd9587/mlir/lib/Dialect/Arith/IR/ArithDialect.cpp#L54


**Description of the Change:**
On the current mlir commit we track, both the old and new bufferization
styles exist. The old pattern rewrite style is deprecated.

To ease the workflow organization, we migrate one dialect at a time.
This PR migrates the `Quantum` dialect's bufferization to the new
one-shot interface.

Note that the new one-shot interface is supposed to be called only once
in the pipeline. However, because we haven't migrated all the dialects
yet, we simply swap out the old `--quantum--bufferize` pass in-place,
with the new one-shot bufferization pass running on the quantum dialect
only.


**Benefits:**
Align with mlir practices; one step closer to updating mlir.


[sc-71487]

---------

Co-authored-by: Tzung-Han Juang <[email protected]>
Co-authored-by: paul0403 <[email protected]>
Co-authored-by: Paul <[email protected]>
Co-authored-by: David Ittah <[email protected]>

Author: 4 people

doc/releases/changelog-dev.md

@@ -110,6 +110,12 @@
 * Improved the definition of `YieldOp` in the quantum dialect by removing `AnyTypeOf`
   [(#1696)](https://github.com/PennyLaneAI/catalyst/pull/1696)
 
+* The bufferization of custom catalyst dialects has been migrated to the new one-shot
+  bufferization interface in mlir.
+  The new mlir bufferization interface is required by jax 0.4.29 or higher.
+  [(#1027)](https://github.com/PennyLaneAI/catalyst/pull/1027)
+  [(#1686)](https://github.com/PennyLaneAI/catalyst/pull/1686)
+
 <h3>Documentation 📝</h3>
 
 <h3>Contributors ✍️</h3>
@@ -119,6 +125,7 @@ This release contains contributions from (in alphabetical order):
 Joey Carter,
 Sengthai Heng,
 David Ittah,
+Tzung-Han Juang,
 Christina Lee,
 Erick Ochoa Lopez,
 Paul Haochen Wang.

frontend/catalyst/pipelines.py

@@ -230,7 +230,7 @@ def get_bufferization_stage(_options: CompileOptions) -> List[str]:
         "catalyst-bufferize",  # Must be run before -- func.func(linalg-bufferize)
         "func.func(linalg-bufferize)",
         "func.func(tensor-bufferize)",
-        "quantum-bufferize",
+        "one-shot-bufferize{dialect-filter=quantum}",
         "func-bufferize",
         "func.func(finalizing-bufferize)",
         "canonicalize",  # Remove dead memrefToTensorOp's

mlir/include/Quantum/Transforms/BufferizableOpInterfaceImpl.h

@@ -0,0 +1,27 @@
+// Copyright 2024-2025 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+using namespace mlir;
+
+namespace catalyst {
+
+namespace quantum {
+
+void registerBufferizableOpInterfaceExternalModels(mlir::DialectRegistry &registry);
+
+}
+
+} // namespace catalyst

mlir/include/Quantum/Transforms/Passes.td

@@ -17,17 +17,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-def QuantumBufferizationPass : Pass<"quantum-bufferize"> {
-    let summary = "Bufferize tensors in quantum operations.";
-
-    let dependentDialects = [
-        "bufferization::BufferizationDialect",
-        "memref::MemRefDialect"
-    ];
-
-    let constructor = "catalyst::createQuantumBufferizationPass()";
-}
-
 def QuantumConversionPass : Pass<"convert-quantum-to-llvm"> {
     let summary = "Perform a dialect conversion from Quantum to LLVM (QIR).";
 

mlir/lib/Catalyst/Transforms/AsyncUtils.cpp

@@ -215,7 +215,7 @@ std::optional<LLVM::LLVMFuncOp> AsyncUtils::getCalleeSafe(LLVM::CallOp callOp)
 bool AsyncUtils::isFunctionNamed(LLVM::LLVMFuncOp funcOp, llvm::StringRef expectedName)
 {
     llvm::StringRef observedName = funcOp.getSymName();
-    return observedName.equals(expectedName);
+    return observedName == expectedName;
 }
 
 bool AsyncUtils::isMlirAsyncRuntimeCreateValue(LLVM::LLVMFuncOp funcOp)

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -50,7 +50,6 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createMemrefToLLVMWithTBAAPass);
     mlir::registerPass(catalyst::createMitigationLoweringPass);
     mlir::registerPass(catalyst::createQnodeToAsyncLoweringPass);
-    mlir::registerPass(catalyst::createQuantumBufferizationPass);
     mlir::registerPass(catalyst::createQuantumConversionPass);
     mlir::registerPass(catalyst::createRegisterInactiveCallbackPass);
     mlir::registerPass(catalyst::createRemoveChainedSelfInversePass);

mlir/lib/Driver/CompilerDriver.cpp

@@ -70,6 +70,7 @@
 #include "Mitigation/Transforms/Passes.h"
 #include "QEC/IR/QECDialect.h"
 #include "Quantum/IR/QuantumDialect.h"
+#include "Quantum/Transforms/BufferizableOpInterfaceImpl.h"
 #include "Quantum/Transforms/Passes.h"
 
 #include "Enzyme.h"
@@ -962,6 +963,9 @@ int QuantumDriverMainFromCL(int argc, char **argv)
     registerAllCatalystDialects(registry);
     registerLLVMTranslations(registry);
 
+    // Register bufferization interfaces
+    catalyst::quantum::registerBufferizableOpInterfaceExternalModels(registry);
+
     // Register and parse command line options.
     std::string inputFilename, outputFilename;
     std::string helpStr = "Catalyst Command Line Interface options. \n"

mlir/lib/Driver/Pipelines.cpp

@@ -16,6 +16,7 @@
 #include "Catalyst/Transforms/Passes.h"
 #include "Gradient/Transforms/Passes.h"
 #include "Mitigation/Transforms/Passes.h"
+#include "Quantum/IR/QuantumDialect.h"
 #include "Quantum/Transforms/Passes.h"
 #include "mhlo/transforms/passes.h"
 #include "mlir/InitAllDialects.h"
@@ -79,7 +80,9 @@ void createBufferizationPipeline(OpPassManager &pm)
     pm.addPass(catalyst::createCatalystBufferizationPass());
     pm.addNestedPass<mlir::func::FuncOp>(mlir::createLinalgBufferizePass());
     pm.addNestedPass<mlir::func::FuncOp>(mlir::tensor::createTensorBufferizePass());
-    pm.addPass(catalyst::createQuantumBufferizationPass());
+    mlir::bufferization::OneShotBufferizationOptions quantum_buffer_options;
+    quantum_buffer_options.opFilter.allowDialect<catalyst::quantum::QuantumDialect>();
+    pm.addPass(mlir::bufferization::createOneShotBufferizePass(quantum_buffer_options));
     pm.addPass(mlir::func::createFuncBufferizePass());
     pm.addNestedPass<mlir::func::FuncOp>(mlir::bufferization::createFinalizingBufferizePass());
     pm.addPass(mlir::createCanonicalizerPass());

mlir/lib/QEC/Transforms/CommuteCliffordPastPPM.cpp

@@ -13,10 +13,11 @@
 // limitations under the License.
 
 #define DEBUG_TYPE "merge_ppr_ppm"
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/Debug.h"
 
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+//#include "mlir/Analysis/TopologicalSortUtils.h"  // enable when updating llvm
 #include "mlir/Transforms/TopologicalSortUtils.h"
 
 #include "QEC/IR/QECDialect.h"

mlir/lib/QEC/Transforms/CommuteCliffordTPPR.cpp

@@ -15,7 +15,7 @@
 #define DEBUG_TYPE "commute_ppr"
 
 #include "llvm/Support/Debug.h"
-
+//#include "mlir/Analysis/TopologicalSortUtils.h" // enable when updating llvm
 #include "mlir/Transforms/TopologicalSortUtils.h"
 
 #include "QEC/IR/QECDialect.h"