Migrate the entire bufferization pipeline to one-shot bufferization (#1751)

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

Now that we have migrated all the individual dialects, we should migrate
the entire bufferization pipeline.

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

See more context in #1027. 

Upstream changes in llvm were required for this bufferization update. As
a result, the llvm version and mlir-hlo version were updated to
```
mhlo=25b008569f413d76cfa8f481f3a84e82b89c47f4
llvm=5f74671c85877e03622e8d308aee15ed73ccee7c
```
These are the versions tracked by jax 0.4.32. 
These are the earliest jax-tagged versions with complete upstream
bufferization changes.

**Related GitHub Issues:**
[sc-71487]

---------

Co-authored-by: Tzung-Han Juang <[email protected]>

Author: paul0403

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

.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

.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

.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

doc/releases/changelog-dev.md

@@ -213,6 +213,7 @@
   [(#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)

frontend/catalyst/pipelines.py

@@ -213,32 +213,33 @@ def get_quantum_compilation_stage(options: CompileOptions) -> List[str]:
     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",
-        "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
-        "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
         "func.func(buffer-hoisting)",
         "func.func(buffer-loop-hoisting)",
+        "func.func(promote-buffers-to-stack)",
         "func.func(buffer-deallocation)",
         "convert-arraylist-to-memref",
         "convert-bufferization-to-memref",
@@ -247,6 +248,7 @@ def get_bufferization_stage(_options: CompileOptions) -> List[str]:
         # "cse",
         "cp-global-memref",
     ]
+
     return bufferization
 
 

mlir/CMakeLists.txt

@@ -51,7 +51,6 @@ set(ALL_MHLO_PASSES
   HloToLinalgUtils
   MhloToLinalg
   MhloToStablehlo
-  MhloQuantToIntConversion
   StablehloToMhlo
 )
 

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:

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();

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);