As part of the llvm update, the old --buffer-deallocation pass is removed.

Intended replacement is --buffer-deallocation-pipeline.
[mlir][bufferization] Remove buffer-deallocation pass llvm/llvm-project#126366,
https://discourse.llvm.org/t/psa-bufferization-new-buffer-deallocation-pipeline/73375.
However, we encountered a few issues. See detailed discord discussion at https://discord.com/channels/636084430946959380/642426447167881246/1376919538301403276

The heart of the new deallocation pass, --ownership-based-buffer-deallocation, blanket-ly fails for ops with unknown memory effects: https://github.com/llvm/llvm-project/blob/da4958ae2b384c2a027cf20c67b7e211d39fcbfe/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp#L523
To solve this issue, the suggestion was to add memory effects to custom operations.
The migration guide also suggests that bufferizable ops no longer implement the bufferizesToAllocation method, so we remove them.

This was supposed to be done alongside the llvm update in #1752;
However, soon it became clear that this migration to the new buffer deallocation is very complicated, and should be its own story.

The llvm update in #1752 thus did not finish this migration.
This PR records the work that was already done on this back in #1752.

Author: paul0403

.dep-versions

@@ -1,16 +1,10 @@
 # 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
-#############
-
+# To update JAX version alongside compatible dependency tags, run the following script:
+# python3 .github/workflows/set_dep_versions.py {JAX_version}
 jax=0.6.0
-mhlo=25b008569f413d76cfa8f481f3a84e82b89c47f4
-llvm=5f74671c85877e03622e8d308aee15ed73ccee7c
-enzyme=v0.0.149
+mhlo=617a9361d186199480c080c9e8c474a5e30c22d1
+llvm=179d30f8c3fddd3c85056fd2b8e877a4a8513158
+enzyme=v0.0.180
 
 # 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-20
+        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-21
 
     - 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-20
+        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-21
 
     - 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-20
+        cmake --build $GITHUB_WORKSPACE/enzyme-build --target EnzymeStatic-21
 
     - name: Save Enzyme Build
       id: save-enzyme-build

.github/workflows/check-catalyst.yaml

@@ -146,7 +146,7 @@ jobs:
         sudo apt-get update
         sudo apt-get install -y python3 python3-pip cmake ninja-build clang lld
         python3 --version | grep ${{ needs.constants.outputs.primary_python_version }}
-        python3 -m pip install numpy pybind11
+        python3 -m pip install numpy pybind11 nanobind==2.4
 
     - name: Build LLVM
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'

.github/workflows/set_dep_versions.py

@@ -32,11 +32,11 @@
 assert os.path.isfile(dep_versions_path)
 assert os.path.isfile(catalyst_init_path)
 
-url = f"https://raw.githubusercontent.com/google/jax/jaxlib-v{jax_version}/WORKSPACE"
+url = f"https://raw.githubusercontent.com/jax-ml/jax/jax-v{jax_version}/WORKSPACE"
 response = requests.get(url)
 match = re.search(r'strip_prefix = "xla-([a-zA-Z0-9]*)"', response.text)
 if not match:
-    url = f"https://raw.githubusercontent.com/google/jax/jaxlib-v{jax_version}/third_party/xla/workspace.bzl"
+    url = f"https://raw.githubusercontent.com/jax-ml/jax/jax-v{jax_version}/third_party/xla/workspace.bzl"
     response = requests.get(url)
     match = re.search(r'XLA_COMMIT = "([a-zA-Z0-9]*)"', response.text)
 xla_commit = match.group(1)
@@ -67,21 +67,16 @@
 response = requests.get(url).json()
 hlo_commit = response["items"][0]["sha"]
 
-existing_text = open(dep_versions_path, "r", encoding="UTF-8").read()
-match = re.search(r"enzyme=([a-zA-Z0-9]*)", existing_text)
-enzyme_commit = match.group(1)
-
-with open(dep_versions_path, "w", encoding="UTF-8") as f:
-    f.write(
-        f"""\
-jax={jax_version}
-mhlo={hlo_commit}
-llvm={llvm_commit}
-enzyme={enzyme_commit}
-"""
-    )
-
 quote = '"'
-cmd = f"sed -i 's/_jaxlib_version = {quote}\([0-9.]\+\){quote}/_jaxlib_version = {quote}{jax_version}{quote}/g' {catalyst_init_path}"
-res = os.system(cmd)
-assert res == 0
+# Update each version using sed
+cmds = [
+    f"sed -i '' 's/^jax=.*/jax={jax_version}/' {dep_versions_path}",
+    f"sed -i '' 's/^mhlo=.*/mhlo={hlo_commit}/' {dep_versions_path}",
+    f"sed -i '' 's/^llvm=.*/llvm={llvm_commit}/' {dep_versions_path}",
+    # Update jaxlib version in __init__.py
+    rf"sed -i '' 's/_jaxlib_version = {quote}\([0-9.]\+\){quote}/_jaxlib_version = {quote}{jax_version}{quote}/g' {catalyst_init_path}",
+]
+
+for cmd in cmds:
+    res = os.system(cmd)
+    assert res == 0

doc/dev/transforms.rst

@@ -512,15 +512,15 @@ and other function operations, which themselves can contain other operations, an
             quantumPatterns.add<QubitUnitaryFusion>(ctx);
 
             // Apply patterns in an iterative and greedy manner.
-            if (failed(applyPatternsAndFoldGreedily(op, std::move(quantumPatterns)))) {
+            if (failed(applyPatternsGreedily(op, std::move(quantumPatterns)))) {
                 return signalPassFailure();
             }
         }
     };
 
 To apply patterns we need a `pattern applicator <https://mlir.llvm.org/docs/PatternRewriter/#common-pattern-drivers>`_.
 There a few in MLIR but typically you can just use the greedy pattern rewrite driver
-(``applyPatternsAndFoldGreedily``), which will iterative over the IR and apply patterns until a
+(``applyPatternsGreedily``), which will iterative over the IR and apply patterns until a
 fixed point is reached.
 
 .. note::

doc/releases/changelog-dev.md

@@ -175,6 +175,14 @@
     [(#1671)](https://github.com/PennyLaneAI/catalyst/pull/1671)
     [(#1681)](https://github.com/PennyLaneAI/catalyst/pull/1681)
 
+* (Compiler developers only) The version of LLVM, MHLO and Enzyme used by Catalyst is
+  updated to track those in jax 0.6.0.
+  [(#1752)](https://github.com/PennyLaneAI/catalyst/pull/1752)
+
+  The LLVM version is updated to commit 179d30f8c3fddd3c85056fd2b8e877a4a8513158.
+  The MHLO version is updated to commit 617a9361d186199480c080c9e8c474a5e30c22d1.
+  The Enzyme version is updated to v0.0.180.
+
 * The clang-format and clang-tidy versions used by Catalyst have been updated to v20.
   [(#1721)](https://github.com/PennyLaneAI/catalyst/pull/1721)
 
@@ -257,7 +265,7 @@
 * Improved the definition of `YieldOp` in the quantum dialect by removing `AnyTypeOf`
   [(#1696)](https://github.com/PennyLaneAI/catalyst/pull/1696)
 
-* The assembly format of `MeasureOp` in the `Quantum` dialect and `MeasureInBasisOp` in the `MBQC` dialect now contains the `postselect` attribute. 
+* The assembly format of `MeasureOp` in the `Quantum` dialect and `MeasureInBasisOp` in the `MBQC` dialect now contains the `postselect` attribute.
   [(#1732)](https://github.com/PennyLaneAI/catalyst/pull/1732)
 
 * The bufferization of custom catalyst dialects has been migrated to the new one-shot

frontend/catalyst/pipelines.py

@@ -240,7 +240,7 @@ def get_bufferization_stage(options: CompileOptions) -> List[str]:
         "func.func(buffer-hoisting)",
         "func.func(buffer-loop-hoisting)",
         "func.func(promote-buffers-to-stack)",
-        "func.func(buffer-deallocation)",
+        "buffer-deallocation-pipeline",
         "convert-arraylist-to-memref",
         "convert-bufferization-to-memref",
         "canonicalize",  # Must be after convert-bufferization-to-memref

mlir/CMakeLists.txt

@@ -47,7 +47,6 @@ set(ALL_MHLO_PASSES
   StablehloPasses
   MhloToArithmeticConversion
   MhloToMemrefConversion
-  MhloToStandard
   HloToLinalgUtils
   MhloToLinalg
   MhloToStablehlo