[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

erick-xanadu · 2025-01-03T16:52:23Z

Hello,

I know that the buffer-deallocation pass is deprecated, but just in case anyone else is still relying on it, we found this issue with buffer-deallocation. The following test case segfaults on main:

func.func public @jit_main() {
    %c0 = arith.constant 0 : i1

    %4 = scf.while (%arg1 = %c0) : (i1) -> (i1) {
      scf.condition(%c0) %arg1 : i1
    } do {
    ^bb0(%arg1: i1):
      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<i32>

      %7 = func.call @jitted_fun_0(%alloc_1, %arg1) : (memref<i32>, i1) -> (i1)

      scf.yield %7#0: i1
    }

    return
}

func.func private @jitted_fun_0(%arg1: memref<i32>, %arg2: i1) -> (i1) {
    return %arg2 : i1
}

It looks like it is attempting to introduce a clone for the i1. With the changes submitted, I believe the error is corrected. Happy to improve the test. It would be nice if this is still merged (if deemed a fix) despite buffer-deallocation being deprecated as other people may encounter this.

Thanks, and happy to make changes to improve this fix.

llvmbot · 2025-01-03T16:52:54Z

@llvm/pr-subscribers-mlir

Author: None (erick-xanadu)

Changes

Hello,

I know that the buffer-deallocation pass is deprecated, but just in case anyone else is still relying on it, we found this issue with buffer-deallocation. The following test case segfaults on main:

func.func public @<!-- -->jit_main() {
    %c0 = arith.constant 0 : i1

    %4 = scf.while (%arg1 = %c0) : (i1) -&gt; (i1) {
      scf.condition(%c0) %arg1 : i1
    } do {
    ^bb0(%arg1: i1):
      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref&lt;i32&gt;

      %7 = func.call @<!-- -->jitted_fun_0(%alloc_1, %arg1) : (memref&lt;i32&gt;, i1) -&gt; (i1)

      scf.yield %7#<!-- -->0: i1
    }

    return
}

func.func private @<!-- -->jitted_fun_0(%arg1: memref&lt;i32&gt;, %arg2: i1) -&gt; (i1) {
    return %arg2 : i1
}

It looks like it is attempting to introduce a clone for the i1. With the changes submitted, I believe the error is corrected. Happy to improve the test. It would be nice if this is still merged (if deemed a fix) despite buffer-deallocation being deprecated as other people may encounter this.

Thanks, and happy to make changes to improve this fix.

Full diff: https://github.com/llvm/llvm-project/pull/121582.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp (+3)
(modified) mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir (+21)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
index a0a81d4add7121..36edbb3395eddf 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -308,6 +308,9 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
 
     // 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)))
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index 3fbe3913c6549e..a6e6f724a822ac 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1271,6 +1271,27 @@ func.func @while_two_arg(%arg0: index) {
 
 // -----
 
+// CHECK-LABEL: func @while_fun
+func.func @while_fun() {
+    %c0 = arith.constant 0 : i1
+    %4 = scf.while (%arg1 = %c0) : (i1) -> (i1) {
+      scf.condition(%c0) %arg1 : i1
+    } do {
+    ^bb0(%arg1: i1):
+      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<i32>
+      %7 = func.call @foo(%alloc_1, %arg1) : (memref<i32>, i1) -> (i1)
+      scf.yield %7#0: i1
+    }
+    return
+}
+
+func.func private @foo(%arg1: memref<i32>, %arg2: i1) -> (i1) {
+    return %arg2 : i1
+}
+
+// -----
+
+// CHECK-LABEL: func @while_three_arg
 func.func @while_three_arg(%arg0: index) {
 // CHECK: %[[ALLOC:.*]] = memref.alloc
   %a = memref.alloc(%arg0) : memref<?xf32>

llvmbot · 2025-01-03T16:52:55Z

@llvm/pr-subscribers-mlir-bufferization

Author: None (erick-xanadu)

Changes

Hello,

I know that the buffer-deallocation pass is deprecated, but just in case anyone else is still relying on it, we found this issue with buffer-deallocation. The following test case segfaults on main:

func.func public @<!-- -->jit_main() {
    %c0 = arith.constant 0 : i1

    %4 = scf.while (%arg1 = %c0) : (i1) -&gt; (i1) {
      scf.condition(%c0) %arg1 : i1
    } do {
    ^bb0(%arg1: i1):
      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref&lt;i32&gt;

      %7 = func.call @<!-- -->jitted_fun_0(%alloc_1, %arg1) : (memref&lt;i32&gt;, i1) -&gt; (i1)

      scf.yield %7#<!-- -->0: i1
    }

    return
}

func.func private @<!-- -->jitted_fun_0(%arg1: memref&lt;i32&gt;, %arg2: i1) -&gt; (i1) {
    return %arg2 : i1
}

It looks like it is attempting to introduce a clone for the i1. With the changes submitted, I believe the error is corrected. Happy to improve the test. It would be nice if this is still merged (if deemed a fix) despite buffer-deallocation being deprecated as other people may encounter this.

Thanks, and happy to make changes to improve this fix.

Full diff: https://github.com/llvm/llvm-project/pull/121582.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp (+3)
(modified) mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir (+21)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
index a0a81d4add7121..36edbb3395eddf 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -308,6 +308,9 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
 
     // 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)))
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index 3fbe3913c6549e..a6e6f724a822ac 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1271,6 +1271,27 @@ func.func @while_two_arg(%arg0: index) {
 
 // -----
 
+// CHECK-LABEL: func @while_fun
+func.func @while_fun() {
+    %c0 = arith.constant 0 : i1
+    %4 = scf.while (%arg1 = %c0) : (i1) -> (i1) {
+      scf.condition(%c0) %arg1 : i1
+    } do {
+    ^bb0(%arg1: i1):
+      %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<i32>
+      %7 = func.call @foo(%alloc_1, %arg1) : (memref<i32>, i1) -> (i1)
+      scf.yield %7#0: i1
+    }
+    return
+}
+
+func.func private @foo(%arg1: memref<i32>, %arg2: i1) -> (i1) {
+    return %arg2 : i1
+}
+
+// -----
+
+// CHECK-LABEL: func @while_three_arg
 func.func @while_three_arg(%arg0: index) {
 // CHECK: %[[ALLOC:.*]] = memref.alloc
   %a = memref.alloc(%arg0) : memref<?xf32>

github-actions · 2025-01-03T16:55:42Z

✅ With the latest revision this PR passed the C/C++ code formatter.

llvm/llvm-project#121582

matthias-springer

Are you using this pass a lot? I wanted to delete this pass from MLIR soon.

matthias-springer · 2025-01-07T08:22:01Z

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

@@ -308,6 +308,9 @@ class BufferDeallocation : public BufferPlacementTransformationBase {

    // Add new allocs and additional clone operations.
    for (Value value : valuesToFree) {
+      if (!isa<BaseMemRefType>(value.getType())) {


Can you put the check earlier, so that non-memref values are not even added to valuesToFree?

dime10 · 2025-01-07T18:40:07Z

Are you using this pass a lot? I wanted to delete this pass from MLIR soon.

Thanks for the review @matthias-springer ! We've been meaning to update to the new bufferization infrastructure although we haven't completed the migration yet (currently we are indeed relying on this pass heavily). Is the new deallocation pass useable without one-shot bufferization?

matthias-springer · 2025-01-08T08:34:43Z

Yes, the new deallocation pass (-buffer-deallocation-pipeline) can be used without One-Shot Bufferize. See documentation here.

…ation (#1408) **Context:** This patch fixes a bug in upstream MLIR that prevents buffer deallocation from working in certain cases. The patch was also submitted upstream: llvm/llvm-project#121582 **Description of the Change:** This PR adds a patch to the llvm-project source, and updates the relevant build recipes. The llvm step during wheel builds is now also handled through Make instead of replicating the CMake configuration in each wheel script. **Benefits:** Unblocks our friends over at Qrisp. **Possible Drawbacks:** Patching dependencies is not great, but hopefully temporary. **Related GitHub Issues:** [sc-81444]

erick-xanadu added 2 commits January 3, 2025 11:37

Introduce copies only for MemRef typed values.

5867e72

Add test

3da161b

erick-xanadu requested a review from matthias-springer as a code owner January 3, 2025 16:52

llvmbot added mlir mlir:bufferization Bufferization infrastructure labels Jan 3, 2025

style

df60731

dime10 added a commit to PennyLaneAI/catalyst that referenced this pull request Jan 3, 2025

Patch mlir buffer-deallocation bug

1677a25

llvm/llvm-project#121582

dime10 mentioned this pull request Jan 3, 2025

Allow quantum registers returned from functions during buffer deallocation PennyLaneAI/catalyst#1408

Merged

matthias-springer approved these changes Jan 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

erick-xanadu commented Jan 3, 2025

llvmbot commented Jan 3, 2025

llvmbot commented Jan 3, 2025

github-actions bot commented Jan 3, 2025 •

edited

Loading

matthias-springer left a comment

matthias-springer Jan 7, 2025

dime10 commented Jan 7, 2025 •

edited

Loading

matthias-springer commented Jan 8, 2025

[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

Are you sure you want to change the base?

[MLIR][buffer-deallocation] Introduce copies only for MemRef typed values. #121582

Conversation

erick-xanadu commented Jan 3, 2025

llvmbot commented Jan 3, 2025

llvmbot commented Jan 3, 2025

github-actions bot commented Jan 3, 2025 • edited Loading

matthias-springer left a comment

Choose a reason for hiding this comment

matthias-springer Jan 7, 2025

Choose a reason for hiding this comment

dime10 commented Jan 7, 2025 • edited Loading

matthias-springer commented Jan 8, 2025

github-actions bot commented Jan 3, 2025 •

edited

Loading

dime10 commented Jan 7, 2025 •

edited

Loading