PennyLaneAI · paul0403 · May 26, 2025 · May 12, 2025 · May 12, 2025 · May 12, 2025
diff --git a/.dep-versions b/.dep-versions
@@ -1,7 +1,15 @@
 # Always update the version check in catalyst.__init__ when changing the JAX version.
+
+#############
+# We track mlir submodule versions from jax 0.4.32 for now
+# These are the earliest versions with complete upstream bufferization changes
+# Versions are retrieved from
+# python3 .github/workflows/set_dep_versions.py 0.4.32
+#############
+
 jax=0.6.0
-mhlo=89a891c986650c33df76885f5620e0a92150d90f
-llvm=3a8316216807d64a586b971f51695e23883331f7
+mhlo=25b008569f413d76cfa8f481f3a84e82b89c47f4
+llvm=5f74671c85877e03622e8d308aee15ed73ccee7c
 enzyme=v0.0.149
 
 # Always remove custom PL/LQ versions before release.

diff --git a/.github/workflows/build-wheel-linux-arm64.yaml b/.github/workflows/build-wheel-linux-arm64.yaml
@@ -222,7 +222,7 @@ jobs:
               -DCMAKE_CXX_VISIBILITY_PRESET=default \
               -DCMAKE_CXX_FLAGS="-fuse-ld=lld"
 
-        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-19
+        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-20
 
     - name: Save Enzyme Build
       id: save-enzyme-build

diff --git a/.github/workflows/build-wheel-linux-x86_64.yaml b/.github/workflows/build-wheel-linux-x86_64.yaml
@@ -245,7 +245,7 @@ jobs:
               -DCMAKE_CXX_VISIBILITY_PRESET=default \
               -DCMAKE_CXX_FLAGS="-fuse-ld=lld"
 
-        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-19
+        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-20
 
     - name: Save Enzyme Build
       id: save-enzyme-build

diff --git a/.github/workflows/build-wheel-macos-arm64.yaml b/.github/workflows/build-wheel-macos-arm64.yaml
@@ -218,7 +218,7 @@ jobs:
               -DENZYME_STATIC_LIB=ON \
               -DCMAKE_CXX_VISIBILITY_PRESET=default
 
-        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-19
+        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-20
 
     - name: Save Enzyme Build
       id: save-enzyme-build

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
@@ -185,6 +185,8 @@
   [(#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)
+  [(#1751)](https://github.com/PennyLaneAI/catalyst/pull/1751)
 
 * Redundant `OptionalAttr` is removed from `adjoint` argument in `QuantumOps.td` TableGen file
   [(#1746)](https://github.com/PennyLaneAI/catalyst/pull/1746)

diff --git a/frontend/catalyst/pipelines.py b/frontend/catalyst/pipelines.py
@@ -213,27 +213,27 @@
     return list(filter(partial(is_not, None), quantum_compilation))
 
 
-def get_bufferization_stage(_options: CompileOptions) -> List[str]:
+def get_bufferization_stage(options: CompileOptions) -> List[str]:
     """Returns the list of passes that performs bufferization"""
+
+    bufferization_options = """bufferize-function-boundaries
+        allow-return-allocs-from-loops
+        function-boundary-type-conversion=identity-layout-map
+        unknown-type-conversion=identity-layout-map""".replace(
+        "\n", " "
+    )
+    if options.async_qnodes:
+        bufferization_options += " copy-before-write"
+
     bufferization = [
-        "one-shot-bufferize{dialect-filter=memref}",
         "inline",
-        "gradient-preprocess",
-        "gradient-bufferize",
-        "scf-bufferize",
         "convert-tensor-to-linalg",  # tensor.pad
-        "convert-elementwise-to-linalg",  # Must be run before --arith-bufferize
-        "arith-bufferize",
-        "empty-tensor-to-alloc-tensor",
-        "func.func(bufferization-bufferize)",
-        "func.func(tensor-bufferize)",
-        # Catalyst dialect's bufferization must be run before --func.func(linalg-bufferize)
-        "one-shot-bufferize{dialect-filter=catalyst unknown-type-conversion=identity-layout-map}",
-        "func.func(linalg-bufferize)",
-        "func.func(tensor-bufferize)",
-        "one-shot-bufferize{dialect-filter=quantum}",
-        "func-bufferize",
-        "func.func(finalizing-bufferize)",
+        "convert-elementwise-to-linalg",  # Must be run before --one-shot-bufferize
+        "gradient-preprocess",
+        "eliminate-empty-tensors",
+        ####################
+        "one-shot-bufferize{" + bufferization_options + "}",
+        ####################
         "canonicalize",  # Remove dead memrefToTensorOp's
         "gradient-postprocess",
         # introduced during gradient-bufferize of callbacks
@@ -247,6 +247,7 @@
         # "cse",
         "cp-global-memref",
     ]
+
     return bufferization
 
 

diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
@@ -51,7 +51,6 @@ set(ALL_MHLO_PASSES
   HloToLinalgUtils
   MhloToLinalg
   MhloToStablehlo
-  MhloQuantToIntConversion
   StablehloToMhlo
 )
 

diff --git a/mlir/Makefile b/mlir/Makefile
@@ -85,7 +85,10 @@ llvm:
 
 	# TODO: when updating LLVM, test to see if mlir/unittests/Bytecode/BytecodeTest.cpp:55 is passing
 	# and remove filter. This tests fails on CI/CD not locally.
-	LIT_FILTER_OUT="Bytecode|tosa-to-tensor" cmake --build $(LLVM_BUILD_DIR) --target $(LLVM_TARGETS)
+	# Note: the upstream lit test llvm-project/mlir/test/python/execution_engine.py requries
+	# the python package `ml_dtypes`. We don't actually use the execution engine, so we skip the
+	# test to reduce unnecessary dependencies.
+	LIT_FILTER_OUT="Bytecode|tosa-to-tensor|execution_engine" cmake --build $(LLVM_BUILD_DIR) --target $(LLVM_TARGETS)
 
 .PHONY: mhlo
 mhlo: TARGET_FILE := $(MK_DIR)/mlir-hlo/mhlo/transforms/CMakeLists.txt
@@ -130,7 +133,7 @@ enzyme:
 		-DCMAKE_CXX_VISIBILITY_PRESET=$(SYMBOL_VISIBILITY) \
 		-DCMAKE_POLICY_DEFAULT_CMP0116=NEW
 
-	cmake --build $(ENZYME_BUILD_DIR) --target EnzymeStatic-19
+	cmake --build $(ENZYME_BUILD_DIR) --target EnzymeStatic-20
 
 .PHONY: plugin
 plugin:

diff --git a/mlir/include/Gradient/IR/GradientOps.h b/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"
diff --git a/mlir/include/Gradient/IR/GradientOps.td b/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"
@@ -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";
 

diff --git a/mlir/include/Gradient/Transforms/BufferizableOpInterfaceImpl.h b/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
diff --git a/mlir/include/Gradient/Transforms/Passes.td b/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.";
 

diff --git a/mlir/include/Gradient/Transforms/Patterns.h b/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 &);

diff --git a/mlir/lib/Catalyst/Transforms/DetectQNodes.cpp b/mlir/lib/Catalyst/Transforms/DetectQNodes.cpp
@@ -925,7 +925,7 @@ struct AddExceptionHandlingPass : impl::AddExceptionHandlingPassBase<AddExceptio
 
         GreedyRewriteConfig config;
         config.strictMode = GreedyRewriteStrictness::ExistingOps;
-        config.enableRegionSimplification = false;
+        config.enableRegionSimplification = mlir::GreedySimplifyRegionLevel::Disabled;
 
         if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns1), config))) {
             signalPassFailure();

diff --git a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp
@@ -383,7 +383,7 @@ struct AnnotateWithFullyQualifiedNamePass
         // Do not fold to save in compile time.
         GreedyRewriteConfig config;
         config.strictMode = GreedyRewriteStrictness::ExistingOps;
-        config.enableRegionSimplification = false;
+        config.enableRegionSimplification = mlir::GreedySimplifyRegionLevel::Disabled;
 
         RewritePatternSet annotate(context);
         auto root = getOperation();
@@ -409,7 +409,7 @@ struct InlineNestedSymbolTablePass : PassWrapper<InlineNestedSymbolTablePass, Op
 
         GreedyRewriteConfig config;
         config.strictMode = GreedyRewriteStrictness::ExistingOps;
-        config.enableRegionSimplification = false;
+        config.enableRegionSimplification = mlir::GreedySimplifyRegionLevel::Disabled;
 
         RewritePatternSet renameFunctions(context);
 

diff --git a/mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp b/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);

diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/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"
@@ -737,7 +738,8 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &
         TimingScope translateTiming = timing.nest("Translate");
         llvmModule =
             timer::timer(translateModuleToLLVMIR, "translateModuleToLLVMIR",
-                         /* add_endl */ false, *mlirModule, llvmContext, "LLVMDialectModule");
+                         /* add_endl */ false, *mlirModule, llvmContext, "LLVMDialectModule",
+                         /* disableVerification */ true);
         if (!llvmModule) {
             CO_MSG(options, Verbosity::Urgent, "Failed to translate LLVM module\n");
             return failure();
@@ -966,6 +968,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.

diff --git a/mlir/lib/Driver/Pipelines.cpp b/mlir/lib/Driver/Pipelines.cpp
@@ -15,6 +15,7 @@
 #include "Driver/Pipelines.h"
 #include "Catalyst/IR/CatalystDialect.h"
 #include "Catalyst/Transforms/Passes.h"
+#include "Gradient/IR/GradientDialect.h"
 #include "Gradient/Transforms/Passes.h"
 #include "Mitigation/Transforms/Passes.h"
 #include "Quantum/IR/QuantumDialect.h"
@@ -66,35 +67,24 @@ void createQuantumCompilationPipeline(OpPassManager &pm)
 }
 void createBufferizationPipeline(OpPassManager &pm)
 {
-    mlir::bufferization::OneShotBufferizationOptions options;
-    options.opFilter.allowDialect<mlir::bufferization::BufferizationDialect>();
-    pm.addPass(mlir::bufferization::createOneShotBufferizePass(options));
     pm.addPass(mlir::createInlinerPass());
-    pm.addPass(catalyst::createGradientPreprocessingPass());
-    pm.addPass(catalyst::createGradientBufferizationPass());
-    pm.addPass(mlir::createSCFBufferizePass());
     pm.addPass(mlir::createConvertTensorToLinalgPass());
     pm.addPass(mlir::createConvertElementwiseToLinalgPass());
-    pm.addPass(mlir::arith::createArithBufferizePass());
-    pm.addPass(mlir::bufferization::createEmptyTensorToAllocTensorPass());
-    pm.addNestedPass<mlir::func::FuncOp>(mlir::bufferization::createBufferizationBufferizePass());
-    pm.addNestedPass<mlir::func::FuncOp>(mlir::tensor::createTensorBufferizePass());
-    mlir::bufferization::OneShotBufferizationOptions catalyst_buffer_options;
-    catalyst_buffer_options.opFilter.allowDialect<catalyst::CatalystDialect>();
-    catalyst_buffer_options.unknownTypeConverterFn =
-        [=](Value value, Attribute memorySpace,
-            const mlir::bufferization::BufferizationOptions &options) {
-            auto tensorType = cast<TensorType>(value.getType());
-            return bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
-        };
-    pm.addPass(mlir::bufferization::createOneShotBufferizePass(catalyst_buffer_options));
-    pm.addNestedPass<mlir::func::FuncOp>(mlir::createLinalgBufferizePass());
-    pm.addNestedPass<mlir::func::FuncOp>(mlir::tensor::createTensorBufferizePass());
-    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(catalyst::createGradientPreprocessingPass());
+    pm.addPass(mlir::bufferization::createEmptyTensorEliminationPass());
+    ///////////
+    mlir::bufferization::OneShotBufferizationOptions options;
+    options.bufferizeFunctionBoundaries = true;
+    options.allowReturnAllocsFromLoops = true;
+    options.setFunctionBoundaryTypeConversion(
+        mlir::bufferization::LayoutMapOption::IdentityLayoutMap);
+    options.unknownTypeConverterFn = [=](Value value, Attribute memorySpace,
+                                         const mlir::bufferization::BufferizationOptions &options) {
+        auto tensorType = cast<TensorType>(value.getType());
+        return bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
+    };
+    pm.addPass(mlir::bufferization::createOneShotBufferizePass(options));
+    //////////////
     pm.addPass(mlir::createCanonicalizerPass());
     pm.addPass(catalyst::createGradientPostprocessingPass());
     pm.addNestedPass<mlir::func::FuncOp>(mlir::bufferization::createBufferHoistingPass());

diff --git a/mlir/lib/Driver/Timer.hpp b/mlir/lib/Driver/Timer.hpp
@@ -25,6 +25,7 @@
 #include <string>
 #include <string_view>
 #include <thread>
+#include <utility> // std::forward
 
 #include <ctime>
 

diff --git a/mlir/lib/Gradient/IR/GradientDialect.cpp b/mlir/lib/Gradient/IR/GradientDialect.cpp
@@ -12,11 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Interfaces/FunctionImplementation.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 #include "Gradient/IR/GradientDialect.h"
 #include "Gradient/IR/GradientOps.h"
-#include "mlir/Interfaces/FunctionImplementation.h"
 
 using namespace mlir;
 using namespace catalyst::gradient;
@@ -50,6 +51,9 @@ void GradientDialect::initialize()
 #include "Gradient/IR/GradientOps.cpp.inc"
         >();
     addInterface<GradientInlinerInterface>();
+
+    declarePromisedInterfaces<bufferization::BufferizableOpInterface, AdjointOp, BackpropOp,
+                              ForwardOp, ReverseOp>();
 }
 
 //===----------------------------------------------------------------------===//