Allow quantum registers returned from functions during buffer deallocation

.github/workflows/build-wheel-linux-x86_64.yaml

@@ -175,25 +175,11 @@ jobs:
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
-        cmake -S mlir/llvm-project/llvm -B llvm-build -G Ninja \
-              -DCMAKE_BUILD_TYPE=Release \
-              -DLLVM_BUILD_EXAMPLES=OFF \
-              -DLLVM_TARGETS_TO_BUILD="host" \
-              -DLLVM_ENABLE_PROJECTS="mlir" \
-              -DLLVM_ENABLE_ASSERTIONS=ON \
-              -DLLVM_INSTALL_UTILS=ON \
-              -DLLVM_ENABLE_ZLIB=FORCE_ON \
-              -DLLVM_ENABLE_ZSTD=OFF \
-              -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
-              -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
-              -DCMAKE_CXX_VISIBILITY_PRESET=default \
-              -DLLVM_ENABLE_LLD=ON
-
-        # 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" cmake --build llvm-build --target check-mlir
+        PYTHON=$(which python${{ matrix.python_version }}) \
+        LLVM_BUILD_DIR="llvm-build" \
+        LLVM_TARGETS="check-mlir" \
+        ENABLE_ZLIB=FORCE_ON \
+        make llvm
 
     - name: Save LLVM Build
       id: save-llvm-build
@@ -486,7 +472,7 @@ jobs:
       if: github.event.schedule == '35 4 * * 4'
       run: |
         python${{ matrix.python_version }} -m pip install standalone_plugin_wheel/wheel/*.whl --no-deps
-    
+
     - name: Run Python Pytest Tests
       run: |
         python${{ matrix.python_version }} -m pytest frontend/test/pytest -n auto

.github/workflows/build-wheel-macos-arm64.yaml

@@ -146,25 +146,11 @@ jobs:
     - name: Build LLVM / MLIR
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'
       run: |
-        cmake -S mlir/llvm-project/llvm -B llvm-build -G Ninja \
-              -DCMAKE_BUILD_TYPE=Release \
-              -DLLVM_BUILD_EXAMPLES=OFF \
-              -DLLVM_TARGETS_TO_BUILD="host" \
-              -DLLVM_ENABLE_PROJECTS="mlir" \
-              -DLLVM_ENABLE_ASSERTIONS=ON \
-              -DLLVM_INSTALL_UTILS=ON \
-              -DLLVM_ENABLE_ZLIB=FORCE_ON \
-              -DLLVM_ENABLE_ZSTD=OFF \
-              -DLLVM_ENABLE_LLD=OFF \
-              -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
-              -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
-              -DCMAKE_CXX_VISIBILITY_PRESET=default
-
-        # 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" cmake --build llvm-build --target check-mlir
+        PYTHON=$(which python${{ matrix.python_version }}) \
+        LLVM_BUILD_DIR="llvm-build" \
+        LLVM_TARGETS="check-mlir" \
+        ENABLE_ZLIB=FORCE_ON \
+        make llvm
 
     - name: Save LLVM Build
       id: save-llvm-build

.github/workflows/build-wheel-macos-x86_64.yaml

@@ -136,25 +136,11 @@ jobs:
     - name: Build LLVM / MLIR
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'
       run: |
-        cmake -S mlir/llvm-project/llvm -B llvm-build -G Ninja \
-              -DCMAKE_BUILD_TYPE=Release \
-              -DLLVM_BUILD_EXAMPLES=OFF \
-              -DLLVM_TARGETS_TO_BUILD="host" \
-              -DLLVM_ENABLE_PROJECTS="mlir" \
-              -DLLVM_ENABLE_ASSERTIONS=ON \
-              -DLLVM_INSTALL_UTILS=ON \
-              -DLLVM_ENABLE_ZLIB=FORCE_ON \
-              -DLLVM_ENABLE_ZSTD=OFF \
-              -DLLVM_ENABLE_LLD=OFF \
-              -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
-              -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
-              -DCMAKE_CXX_VISIBILITY_PRESET=default
-
-        # 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" cmake --build llvm-build --target check-mlir
+        PYTHON=$(which python${{ matrix.python_version }}) \
+        LLVM_BUILD_DIR="llvm-build" \
+        LLVM_TARGETS="check-mlir" \
+        ENABLE_ZLIB=FORCE_ON \
+        make llvm
 
     - name: Save LLVM Build
       id: save-llvm-build

.github/workflows/scripts/linux_arm64/rh8/build_llvm.sh

@@ -34,19 +34,7 @@ export PATH=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/bin:/o
 /usr/bin/python3 -m pip install pennylane pybind11 PyYAML cmake ninja
 
 # Build LLVM
-cmake -S /catalyst/mlir/llvm-project/llvm -B /catalyst/llvm-build -G Ninja \
-    -DCMAKE_BUILD_TYPE=Release \
-    -DLLVM_BUILD_EXAMPLES=OFF \
-    -DLLVM_TARGETS_TO_BUILD="host" \
-    -DLLVM_ENABLE_PROJECTS="mlir" \
-    -DLLVM_ENABLE_ASSERTIONS=ON \
-    -DLLVM_INSTALL_UTILS=ON \
-    -DLLVM_ENABLE_ZLIB=FORCE_ON \
-    -DLLVM_ENABLE_ZSTD=OFF \
-    -DLLVM_ENABLE_LLD=ON \
-    -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-    -DPython3_EXECUTABLE=/usr/bin/python3 \
-    -DPython3_NumPy_INCLUDE_DIRS=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/lib/python${PYTHON_VERSION}/site-packages/numpy/core/include \
-    -DCMAKE_CXX_VISIBILITY_PRESET=default
-
-LIT_FILTER_OUT="Bytecode|tosa-to-tensor" cmake --build /catalyst/llvm-build --target check-mlir llvm-symbolizer
+export LLVM_BUILD_DIR="/catalyst/llvm-build"
+export LLVM_TARGETS="check-mlir"
+export ENABLE_ZLIB=FORCE_ON
+make llvm

mlir/Makefile

@@ -12,8 +12,6 @@ ENZYME_BUILD_DIR?=$(MK_DIR)/Enzyme/build
 RT_BUILD_DIR?=$(MK_DIR)/../runtime/build
 ENABLE_ASAN?=OFF
 BUILD_TYPE?=Release
-TARGET_FILE=$(MK_DIR)/mlir-hlo/mhlo/transforms/CMakeLists.txt
-PATCH_FILE=$(MK_DIR)/patches/mhlo-Add-PassesIncGen-in-transforms-CMakeList.patch
 LLVM_EXTERNAL_LIT ?= $(LLVM_BUILD_DIR)/bin/llvm-lit
 
 ifeq ($(shell uname), Darwin)
@@ -36,6 +34,8 @@ USE_SANITIZER_NAMES=""
 USE_SANITIZER_FLAGS=""
 endif
 
+LLVM_TARGETS ?= check-mlir llvm-symbolizer
+
 .PHONY: help
 help:
 	@echo "Please use \`make <target>' where <target> is one of"
@@ -53,8 +53,15 @@ help:
 all: llvm mhlo enzyme dialects plugin
 
 .PHONY: llvm
+llvm: TARGET_FILE := $(MK_DIR)/llvm-project/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+llvm: PATCH_FILE := $(MK_DIR)/patches/mlir-buffer-deallocation.patch
 llvm:
 	@echo "build LLVM and MLIR enabling Python bindings"
+	# Patch in MLIR buffer deallocation bugfix
+	# TODO: remove once https://github.com/llvm/llvm-project/pull/121582 is merged & the dep updated
+	@if patch --dry-run -p1 -N $(TARGET_FILE) $(PATCH_FILE) > /dev/null 2>&1; then \
+		patch -p1 $(TARGET_FILE) $(PATCH_FILE); \
+	fi
 	cmake -G Ninja -S llvm-project/llvm -B $(LLVM_BUILD_DIR) \
 		-DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \
 		-DLLVM_BUILD_EXAMPLES=OFF \
@@ -75,10 +82,12 @@ llvm:
 		-DCMAKE_CXX_VISIBILITY_PRESET=$(SYMBOL_VISIBILITY)
 
 	# TODO: when updating LLVM, test to see if mlir/unittests/Bytecode/BytecodeTest.cpp:55 is passing
-	# and remove filter
-	LIT_FILTER_OUT="Bytecode" cmake --build $(LLVM_BUILD_DIR) --target check-mlir llvm-symbolizer
+	# 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)
 
 .PHONY: mhlo
+mhlo: TARGET_FILE := $(MK_DIR)/mlir-hlo/mhlo/transforms/CMakeLists.txt
+mhlo: PATCH_FILE := $(MK_DIR)/patches/mhlo-Add-PassesIncGen-in-transforms-CMakeList.patch
 mhlo:
 	@echo "build MLIR-HLO"
 	# Patch MHLO cmake dependency

mlir/patches/mlir-buffer-deallocation.patch

@@ -0,0 +1,14 @@
+diff --git a/mlir/llvm-project/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/llvm-project/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+index a0a81d4add..7b7be9e577 100644
+--- a/mlir/llvm-project/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
++++ b/mlir/llvm-project/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+@@ -308,6 +308,9 @@ private:
+
+     // Add new allocs and additional clone operations.
+     for (Value value : valuesToFree) {
++      if (!isa<BaseMemRefType>(value.getType())) {
++          continue;
++      }
+       if (failed(isa<BlockArgument>(value)
+                      ? introduceBlockArgCopy(cast<BlockArgument>(value))
+                      : introduceValueCopyForRegionResult(value)))