Migrate quantum dialect to new one-shot bufferization

doc/releases/changelog-dev.md

@@ -71,6 +71,9 @@
 * The utility function `EnsureFunctionDeclaration` is refactored into the `Utils` of the `Catalyst` dialect, instead of being duplicated in each individual dialect.
   [(#1683)](https://github.com/PennyLaneAI/catalyst/pull/1683)
 
+* 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.
+  [(#1686)](https://github.com/PennyLaneAI/catalyst/pull/1686)
+
 <h3>Documentation 📝</h3>
 
 <h3>Contributors ✍️</h3>
@@ -80,6 +83,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/IR/QuantumOps.td

@@ -63,15 +63,15 @@ def NamedObservableAttr : EnumAttr<QuantumDialect, NamedObservable, "named_obser
 
 class Quantum_Op<string mnemonic, list<Trait> traits = []> : Op<QuantumDialect, mnemonic, traits>;
 
-def InitializeOp : Quantum_Op<"init"> {
+def InitializeOp : Quantum_Op<"init"> { //, [NoMemoryEffect]> {
     let summary = "Initialize the quantum runtime.";
 
     let assemblyFormat = [{
         attr-dict
     }];
 }
 
-def FinalizeOp : Quantum_Op<"finalize"> {
+def FinalizeOp : Quantum_Op<"finalize"> { //, [NoMemoryEffect]> {
     let summary = "Teardown the quantum runtime.";
 
     let assemblyFormat = [{

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"

mlir/lib/Quantum/IR/QuantumDialect.cpp

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/DialectImplementation.h" // needed for generated type parser
 #include "llvm/ADT/TypeSwitch.h"           // needed for generated type parser
 
@@ -43,6 +44,10 @@ void QuantumDialect::initialize()
 #define GET_OP_LIST
 #include "Quantum/IR/QuantumOps.cpp.inc"
         >();
+
+    declarePromisedInterfaces<bufferization::BufferizableOpInterface, QubitUnitaryOp, HermitianOp,
+                              HamiltonianOp, SampleOp, CountsOp, ProbsOp, StateOp, SetStateOp,
+                              SetBasisStateOp>();
 }
 
 //===----------------------------------------------------------------------===//