Update llvm version

.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

@@ -187,9 +187,10 @@ jobs:
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
 
-        export TARGET_FILE=mlir/mlir-hlo/mhlo/transforms/CMakeLists.txt
-        export PATCH_FILE=mlir/patches/mhlo-Add-PassesIncGen-in-transforms-CMakeList.patch
-        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi
+        pushd $GITHUB_WORKSPACE/mlir/mlir-hlo
+        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
+        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-add-back-necessary-passes.patch
+        popd
 
         cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
@@ -215,14 +216,19 @@ jobs:
       if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
+
+        export TARGET_FILE=mlir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
+        export PATCH_FILE=mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
+        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi
+
         cmake -S mlir/Enzyme/enzyme -B $GITHUB_WORKSPACE/enzyme-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
               -DLLVM_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm" \
               -DENZYME_STATIC_LIB=ON \
               -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

@@ -210,9 +210,10 @@ jobs:
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
 
-        export TARGET_FILE=mlir/mlir-hlo/mhlo/transforms/CMakeLists.txt
-        export PATCH_FILE=mlir/patches/mhlo-Add-PassesIncGen-in-transforms-CMakeList.patch
-        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi
+        pushd $GITHUB_WORKSPACE/mlir/mlir-hlo
+        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
+        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-add-back-necessary-passes.patch
+        popd
 
         cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
@@ -238,14 +239,19 @@ jobs:
       if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
+
+        export TARGET_FILE=mlir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
+        export PATCH_FILE=mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
+        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi
+
         cmake -S mlir/Enzyme/enzyme -B $GITHUB_WORKSPACE/enzyme-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
               -DLLVM_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm" \
               -DENZYME_STATIC_LIB=ON \
               -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

@@ -185,9 +185,10 @@ jobs:
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
 
-        export TARGET_FILE=mlir/mlir-hlo/mhlo/transforms/CMakeLists.txt
-        export PATCH_FILE=mlir/patches/mhlo-Add-PassesIncGen-in-transforms-CMakeList.patch
-        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi
+        pushd $GITHUB_WORKSPACE/mlir/mlir-hlo
+        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
+        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-add-back-necessary-passes.patch
+        popd
 
         cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
@@ -212,13 +213,17 @@ jobs:
     - name: Build Enzyme
       if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
       run: |
+        export TARGET_FILE=mlir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
+        export PATCH_FILE=mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
+        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi
+
         cmake -S mlir/Enzyme/enzyme -B $GITHUB_WORKSPACE/enzyme-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
               -DLLVM_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm" \
               -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
 
     - name: Build LLVM
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'
@@ -194,7 +194,7 @@ jobs:
       uses: actions/cache@v4
       with:
         path: mhlo-build
-        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-ci-build-${{ matrix.compiler }}
+        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-ci-build-${{ matrix.compiler }}-0
 
     - name: Get Cached LLVM Source
       id: cache-llvm-source
@@ -351,7 +351,7 @@ jobs:
       uses: actions/cache/restore@v4
       with:
         path: mhlo-build
-        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-ci-build-${{ matrix.compiler }}
+        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-ci-build-${{ matrix.compiler }}-0
         fail-on-cache-miss: true
 
     - name: Get Cached Enzyme Source

.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

@@ -252,23 +252,22 @@ Note how the value ``%q2`` links the two operations together from definition ``(
 across several other instructions.
 
 As seen in the `pattern rewriter documentation <https://mlir.llvm.org/docs/PatternRewriter/#defining-patterns>`_,
-a new rewrite pattern can be defined as a C++ class as follows, where we will focus on the ``match``
-and ``rewrite`` methods (refer to the link for the full class and up to date information):
+a new rewrite pattern can be defined as a C++ class as follows, where we will focus on the
+``matchAndRewrite`` method (refer to the link for the full class and up to date information):
 
 .. code-block:: cpp
 
     struct QubitUnitaryFusion : public OpRewritePattern<QubitUnitaryOp>
     {
         ...
 
-        LogicalResult match(QubitUnitaryOp op) const override {
-            // The ``match`` method returns ``success()`` if the pattern is a match, failure
-            // otherwise.
-        }
-
-        void rewrite(QubitUnitaryOp op, PatternRewriter &rewriter) {
-            // The ``rewrite`` method performs mutations on the IR rooted at ``op`` using
-            // the provided rewriter. All mutations must go through the provided rewriter.
+        LogicalResult matchAndRewrite(QubitUnitaryOp op, PatternRewriter &rewriter) const override {
+            // The `matchAndRewrite` method performs both the pattern matching and the mutation 
+            // on the IR rooted at `op` using the provided rewriter.
+            // All mutations must go through the provided rewriter and IR mutation should only
+            // take place after the match is deemed successful. 
+            // matchAndRewrite must return "success" if and only if the IR was modified.
+            // The root operation is required to either be: updated in-place, replaced, or erased.
         }
 
         ...
@@ -286,11 +285,11 @@ the second is a list of qubits):
 
     QubitUnitary(*, QubitUnitary(*, *))
 
-Let's implement it in C++:
+Let's add the pattern-matching logic to the ``matchAndRewrite`` method:
 
 .. code-block:: cpp
 
-    LogicalResult match(QubitUnitaryOp op) const override
+    LogicalResult matchAndRewrite(QubitUnitaryOp op, PatternRewriter &rewriter) const override
     {
         ValueRange qbs = op.getInQubits();
         Operation *parent = qbs[0].getDefiningOp();
@@ -314,6 +313,9 @@ Let's implement it in C++:
                 return failure();
             }
 
+        // Rewrite logic
+        // ... We have matched the pattern, now rewrite the IR here
+
         return success();
     }
 
@@ -351,8 +353,8 @@ MLIR will automatically generate canonical ``get*`` methods for attributes like
 ``out_qubits``, and ``matrix``. When in doubt it's best to have a look at the generated C++ files in
 the build folder, named ``QuantumOps.h.inc`` and ``QuantumOps.cpp.inc`` in this instance.
 
-Alright, now that we have the matching part, let's implement the actual transformation via the
-``rewrite`` method. All we need to do is replace the original pattern with the following:
+Alright, now that we have the matching part, let's add the actual transformation to the
+``matchAndRewrite`` method. All we need to do is replace the original pattern with the following:
 
 .. code-block::
 
@@ -362,8 +364,13 @@ In C++ it will look as follows:
 
 .. code-block:: cpp
 
-    void rewrite(QubitUnitaryOp op, PatternRewriter &rewriter) const override
+    LogicalResult matchAndRewrite(QubitUnitaryOp op, PatternRewriter &rewriter) const override
     {
+
+        // Pattern matching logic
+        // ... match the pattern
+
+        // Rewrite logic
         ValueRange qbs = op.getInQubits();
         QubitUnitaryOp parentOp = cast<QubitUnitaryOp>(qbs[0].getDefiningOp());
 
@@ -410,11 +417,13 @@ In C++ it will look as follows:
         // The second unitary is not needed anymore
         // Whoever uses the second unitary, use the first one instead!
         op.replaceAllUsesWith(parentOp);
+
+        return success();
     }
 
 When writing transformations, the rewriter is the most important tool we have. It can create new
 operations for us, delete others, or change the place in the IR where we are choosing to make
-changes (also called the insertion point). Let's have look at some of these elements:
+changes (also called the insertion point). Let's have a look at some of these elements:
 
 - **Constructing new operations**:
 
@@ -512,15 +521,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::
@@ -565,21 +574,30 @@ gradient ops that specify the finite-difference method, indicated via the ``"fd"
 
 .. code-block:: cpp
 
-    LogicalResult FiniteDiffLowering::match(GradOp op)
+    LogicalResult FiniteDiffLowering::matchAndRewrite(GradOp op, PatternRewriter &rewriter)
     {
-        if (op.getMethod() == "fd")
-            return success();
+        // Pattern matching logic
+        if (op.getMethod() != "fd")
+            return failure();
 
-        return failure();
+        // Rewrite logic
+        // ...
+
+        return success();
     }
 
 For the rewriting part we'll want to introduce a few new elements, such as looking up symbols
 (function names), creating new functions, and changing the insertion point.
 
 .. code-block:: cpp
 
-    void FiniteDiffLowering::rewrite(GradOp op, PatternRewriter &rewriter)
+    LogicalResult FiniteDiffLowering::matchAndRewrite(GradOp op, PatternRewriter &rewriter)
     {
+        // Pattern matching logic
+        if (op.getMethod() != "fd")
+            return failure();
+
+        // Rewrite logic
         // First let's find the function the grad operation is referencing.
         func::FuncOp callee =
             SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(op, op.getCalleeAttr());
@@ -609,6 +627,8 @@ For the rewriting part we'll want to introduce a few new elements, such as looki
             // Populate the function body.
             populateFiniteDiffMethod(rewriter, op, gradFn);
         }
+
+        return success();
     }
 
 Symbols are string references to IR objects, which rather than containing a physical reference or
@@ -711,18 +731,20 @@ Alright, our function should now look something like this:
         func.return %dx, %dy, %dz : f64, f64, f64
     }
 
-Finally, we have to amend our rewrite function to invoke the new function we created and delete the
+Finally, we have to amend our ``matchAndRewrite`` function to invoke the new function we created and delete the
 ``GradOp`` from the IR:
 
 .. code-block:: cpp
 
-    void FiniteDiffLowering::rewrite(GradOp op, PatternRewriter &rewriter)
+    LogicalResult FiniteDiffLowering::matchAndRewrite(GradOp op, PatternRewriter &rewriter)
     {
         ...
             populateFiniteDiffMethod(rewriter, op, gradFn);
         }
 
         rewriter.replaceOpWithNewOp<func::CallOp>(op, gradFn, op.getArgOperands());
+
+        return success();
     }
 
 Note how we can create a new operation, take its results, and use those to replace another operation