Migrate gradient dialect to new one-shot bufferization (#1740)

**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.

The `Quantum` dialect was migrated in #1686 .
The `Catalyst` dialect was migrated in #1708 .

Note that #1139 refactors the gradient dialect's bufferization into
preprocess, bufferization, and postprocess.
Only the middle bufferization stage is supposed to be replaced by
one-shot bufferization.

**Description of the Change:**
Migrate `Gradient` dialect to one-shot bufferization.

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

[sc-71487]

---------

Co-authored-by: Tzung-Han Juang <tzunghan.juang@gmail.com>

Author: paul0403

doc/releases/changelog-dev.md

@@ -212,6 +212,7 @@
   [(#1027)](https://github.com/PennyLaneAI/catalyst/pull/1027)
   [(#1686)](https://github.com/PennyLaneAI/catalyst/pull/1686)
   [(#1708)](https://github.com/PennyLaneAI/catalyst/pull/1708)
+  [(#1740)](https://github.com/PennyLaneAI/catalyst/pull/1740)
 
 * Redundant `OptionalAttr` is removed from `adjoint` argument in `QuantumOps.td` TableGen file
   [(#1746)](https://github.com/PennyLaneAI/catalyst/pull/1746)

frontend/catalyst/pipelines.py

@@ -219,7 +219,7 @@ def get_bufferization_stage(_options: CompileOptions) -> List[str]:
         "one-shot-bufferize{dialect-filter=memref}",
         "inline",
         "gradient-preprocess",
-        "gradient-bufferize",
+        "one-shot-bufferize{dialect-filter=gradient unknown-type-conversion=identity-layout-map}",
         "scf-bufferize",
         "convert-tensor-to-linalg",  # tensor.pad
         "convert-elementwise-to-linalg",  # Must be run before --arith-bufferize

mlir/include/Gradient/IR/GradientOps.h

@@ -21,10 +21,10 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
-
-#include "Gradient/IR/GradientInterfaces.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 
 #include "Gradient/IR/GradientDialect.h"
+#include "Gradient/IR/GradientInterfaces.h"
 
 #define GET_OP_CLASSES
 #include "Gradient/IR/GradientOps.h.inc"

mlir/include/Gradient/IR/GradientOps.td

@@ -17,6 +17,7 @@
 
 include "mlir/Interfaces/FunctionInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/IR/BuiltinAttributes.td"
 include "mlir/IR/OpBase.td"
@@ -306,7 +307,7 @@ def ForwardOp : Gradient_Op<"forward",
      OptionalAttr<DictArrayAttr>: $res_attrs
   );
 
-  let regions = (region AnyRegion: $body);
+  let regions = (region MaxSizedRegion<1>: $body);
 
   let builders = [OpBuilder<(ins
     "mlir::StringRef":$name, "mlir::FunctionType":$type,
@@ -362,7 +363,7 @@ def ReverseOp : Gradient_Op<"reverse",
      OptionalAttr<DictArrayAttr>: $res_attrs
   );
 
-  let regions = (region AnyRegion: $body);
+  let regions = (region MaxSizedRegion<1>: $body);
 
   let builders = [OpBuilder<(ins
     "mlir::StringRef":$name, "mlir::FunctionType":$type,
@@ -388,7 +389,7 @@ def ReverseOp : Gradient_Op<"reverse",
 }
 
 def ReturnOp : Gradient_Op<"return",
-    [Terminator, ParentOneOf<["ForwardOp", "ReverseOp"]>]> {
+    [ReturnLike, Terminator, ParentOneOf<["ForwardOp", "ReverseOp"]>]> {
 
   let summary = "Return tapes or nothing";
 

mlir/include/Gradient/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 gradient {
+
+void registerBufferizableOpInterfaceExternalModels(mlir::DialectRegistry &registry);
+
+}
+
+} // namespace catalyst

mlir/include/Gradient/Transforms/Passes.td

@@ -17,18 +17,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-def GradientBufferizationPass : Pass<"gradient-bufferize"> {
-    let summary = "Bufferize tensors in quantum operations.";
-
-    let dependentDialects = [
-        "bufferization::BufferizationDialect",
-        "memref::MemRefDialect",
-        "index::IndexDialect"
-    ];
-
-    let constructor = "catalyst::createGradientBufferizationPass()";
-}
-
 def GradientPreprocessingPass : Pass<"gradient-preprocess"> {
     let summary = "Insert Func.CallOp for ForwardOp and ReverseOp.";
 

mlir/include/Gradient/Transforms/Patterns.h

@@ -21,7 +21,6 @@
 namespace catalyst {
 namespace gradient {
 
-void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populatePreprocessingPatterns(mlir::RewritePatternSet &);
 void populatePostprocessingPatterns(mlir::RewritePatternSet &);
 void populateLoweringPatterns(mlir::RewritePatternSet &);

mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -22,7 +22,6 @@
 #include "Catalyst/Transforms/BufferizableOpInterfaceImpl.h"
 
 using namespace mlir;
-using namespace mlir::bufferization;
 using namespace catalyst;
 
 /**

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -37,7 +37,6 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createDisentangleSWAPPass);
     mlir::registerPass(catalyst::createEmitCatalystPyInterfacePass);
     mlir::registerPass(catalyst::createGEPInboundsPass);
-    mlir::registerPass(catalyst::createGradientBufferizationPass);
     mlir::registerPass(catalyst::createGradientConversionPass);
     mlir::registerPass(catalyst::createGradientPreprocessingPass);
     mlir::registerPass(catalyst::createGradientPostprocessingPass);

mlir/lib/Driver/CompilerDriver.cpp

@@ -64,6 +64,7 @@
 #include "Driver/Support.h"
 #include "Gradient/IR/GradientDialect.h"
 #include "Gradient/IR/GradientInterfaces.h"
+#include "Gradient/Transforms/BufferizableOpInterfaceImpl.h"
 #include "Gradient/Transforms/Passes.h"
 #include "Ion/IR/IonDialect.h"
 #include "MBQC/IR/MBQCDialect.h"
@@ -966,6 +967,7 @@ int QuantumDriverMainFromCL(int argc, char **argv)
 
     // Register bufferization interfaces
     catalyst::registerBufferizableOpInterfaceExternalModels(registry);
+    catalyst::gradient::registerBufferizableOpInterfaceExternalModels(registry);
     catalyst::quantum::registerBufferizableOpInterfaceExternalModels(registry);
 
     // Register and parse command line options.