Migrate catalyst dialect to new one-shot bufferization

doc/releases/changelog-dev.md

@@ -137,6 +137,7 @@
   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)
+  [(#1708)](https://github.com/PennyLaneAI/catalyst/pull/1708)
 
 <h3>Documentation 📝</h3>
 

frontend/catalyst/pipelines.py

@@ -97,7 +97,7 @@ class CompileOptions:
 
     def __post_init__(self):
         # Check that async runs must not be seeded
-        if self.async_qnodes and self.seed != None:
+        if self.async_qnodes and self.seed is not None:
             raise CompileError(
                 """
                 Seeding has no effect on asynchronous QNodes,
@@ -107,7 +107,7 @@ def __post_init__(self):
             )
 
         # Check that seed is 32-bit unsigned int
-        if (self.seed != None) and (self.seed < 0 or self.seed > 2**32 - 1):
+        if (self.seed is not None) and (self.seed < 0 or self.seed > 2**32 - 1):
             raise ValueError(
                 """
                 Seed must be an unsigned 32-bit integer!
@@ -227,12 +227,12 @@ def get_bufferization_stage(_options: CompileOptions) -> List[str]:
         "empty-tensor-to-alloc-tensor",
         "func.func(bufferization-bufferize)",
         "func.func(tensor-bufferize)",
-        "catalyst-bufferize",  # Must be run before -- func.func(linalg-bufferize)
+        # Catalyst dialect's bufferization must be run before --func.func(linalg-bufferize)
+        "one-shot-bufferize{dialect-filter=catalyst}",
         "func.func(linalg-bufferize)",
         "func.func(tensor-bufferize)",
         "one-shot-bufferize{dialect-filter=quantum}",
         "func-bufferize",
-        "func.func(finalizing-bufferize)",
         "canonicalize",  # Remove dead memrefToTensorOp's
         "gradient-postprocess",
         # introduced during gradient-bufferize of callbacks

mlir/include/Catalyst/Transforms/BufferizableOpInterfaceImpl.h

@@ -0,0 +1,23 @@
+// 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 {
+
+void registerBufferizableOpInterfaceExternalModels(mlir::DialectRegistry &registry);
+
+} // namespace catalyst

mlir/include/Catalyst/Transforms/Passes.td

@@ -27,18 +27,6 @@ def DetensorizeSCFPass : Pass<"detensorize-scf"> {
     let constructor = "catalyst::createDetensorizeSCFPass()";
 }
 
-def CatalystBufferizationPass : Pass<"catalyst-bufferize"> {
-    let summary = "Bufferize tensors in catalyst utility ops.";
-
-    let dependentDialects = [
-        "bufferization::BufferizationDialect",
-        "memref::MemRefDialect",
-        "index::IndexDialect"
-    ];
-
-    let constructor = "catalyst::createCatalystBufferizationPass()";
-}
-
 def ArrayListToMemRefPass : Pass<"convert-arraylist-to-memref"> {
     let summary = "Lower array list operations to memref operations.";
     let description = [{

mlir/include/Catalyst/Transforms/Patterns.h

@@ -21,8 +21,6 @@
 
 namespace catalyst {
 
-void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
-
 void populateScatterPatterns(mlir::RewritePatternSet &);
 
 void populateHloCustomCallPatterns(mlir::RewritePatternSet &);

mlir/include/Quantum/Transforms/Patterns.h

@@ -21,8 +21,6 @@
 namespace catalyst {
 namespace quantum {
 
-void populateBufferizationLegality(mlir::TypeConverter &, mlir::ConversionTarget &);
-void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateQIRConversionPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateAdjointPatterns(mlir::RewritePatternSet &);
 void populateSelfInversePatterns(mlir::RewritePatternSet &);

mlir/lib/Catalyst/IR/CatalystDialect.cpp

@@ -14,6 +14,7 @@
 
 #include "Catalyst/IR/CatalystDialect.h"
 #include "Catalyst/IR/CatalystOps.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h" // needed for generated type parser
 #include "mlir/Interfaces/FunctionImplementation.h"
@@ -40,6 +41,9 @@ void CatalystDialect::initialize()
 #define GET_OP_LIST
 #include "Catalyst/IR/CatalystOps.cpp.inc"
         >();
+
+    declarePromisedInterfaces<bufferization::BufferizableOpInterface, PrintOp, CustomCallOp,
+                              CallbackCallOp, CallbackOp>();
 }
 
 //===----------------------------------------------------------------------===//