<table border="1" cellspacing="0" cellpadding="8">
    <tr>
        <th>Issue</th>
        <td>
            <a href=https://github.com/llvm/llvm-project/issues/77096>77096</a>
        </td>
    </tr>

    <tr>
        <th>Summary</th>
        <td>
            [mlir][bufferization] Double free when using private-function-dynamic-ownership=true
        </td>
    </tr>

    <tr>
      <th>Labels</th>
      <td>
            mlir
      </td>
    </tr>

    <tr>
      <th>Assignees</th>
      <td>
      </td>
    </tr>

    <tr>
      <th>Reporter</th>
      <td>
          sabauma
      </td>
    </tr>
</table>

<pre>
    When using `ownership-based-buffer-deallocation=private-function-dynamic-ownership=true`, the generated deallocations produce a double-free for `memref` function arguments which are not subsequently returned.

## Reproducer

```mlir
// RUN:   mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership=true -test-lower-to-llvm \
// RUN:   | mlir-cpu-runner -entry-point-result=i32

func.func private @private_callee(%arg0: memref<f32>) -> memref<f32> {
  %alloc = memref.alloc() : memref<f32>
 return %alloc : memref<f32>
}

func.func @caller() -> (f32) {
 %alloc = memref.alloc() : memref<f32>
  %ret = call @private_callee(%alloc) : (memref<f32>) -> memref<f32>

  %val = memref.load %ret[] : memref<f32>
  return %val : f32
}

// Driver main function for mlir-cpu-runner
func.func @main() -> i32 {
  %res = func.call @caller() : () -> f32
  %val = arith.fptosi %res : f32 to i32
  return %val : i32
}
```

Executing this example crashes the `mlir-cpu-runner` with the error: `free(): double free detected in tcache 2`.

## What happens

`ownership-based-buffer-deallocation=private-function-dynamic-ownership=true` generates the following code for the caller and callee:

```mlir
func.func private @private_callee(%arg0: memref<f32>, %arg1: i1) -> (memref<f32>, i1) {
  %true = arith.constant true
  %alloc = memref.alloc() : memref<f32>
  %base_buffer, %offset = memref.extract_strided_metadata %arg0 : memref<f32> -> memref<f32>, index
 %base_buffer_0, %offset_1 = memref.extract_strided_metadata %alloc : memref<f32> -> memref<f32>, index
  %0 = bufferization.dealloc (%base_buffer, %base_buffer_0 : memref<f32>, memref<f32>) if (%arg1, %true) retain (%alloc : memref<f32>)
  return %alloc, %0 : memref<f32>, i1
}
func.func @caller() -> f32 {
  %true = arith.constant true
 %alloc = memref.alloc() : memref<f32>
  %0:2 = call @private_callee(%alloc, %true) : (memref<f32>, i1) -> (memref<f32>, i1)
 %1 = memref.load %0#0[] : memref<f32>
  %base_buffer, %offset = memref.extract_strided_metadata %alloc : memref<f32> -> memref<f32>, index
  %base_buffer_0, %offset_1 = memref.extract_strided_metadata %0#0 : memref<f32> -> memref<f32>, index
  bufferization.dealloc (%base_buffer, %base_buffer_0 : memref<f32>, memref<f32>) if (%true, %0#1)
  return %1 : f32
}
```

Two deallocations are inserted: one in `@private_callee`

```mlir
%0 = bufferization.dealloc (%base_buffer, %base_buffer_0 : memref<f32>, memref<f32>) if (%arg1, %true) retain (%alloc : memref<f32>)
```

and the other in `@caller`

```mlir
bufferization.dealloc (%base_buffer, %base_buffer_0 : memref<f32>, memref<f32>) if (%true, %0#1)
```

In both cases, `%base_buffer` refers to the allocation in `@caller` and the ownership indicator is `true`. At runtime, both `@private_callee` and `@caller` always free the memref value allocated within `@caller`.

## Other observations

Changing caller to the following:

```mlir
func.func @caller() -> (f32) {
  %alloc = memref.alloc() : memref<f32>
  %ret = call @private_callee(%alloc) : (memref<f32>) -> memref<f32>

  // Read from %alloc rather than %ret
  %val = memref.load %alloc[] : memref<f32>
  return %val : f32
}
```

Results in a read-after-free of `%alloc`, where `%alloc` is still deallocated within `@private_callee`, but has a use after the call.

</pre>
<img width="1px" height="1px" alt="" src="http://email.email.llvm.org/o/eJzUWF1v27gS_TX0y0AGTVmy8-CH1G6A-3IvUNxFHwNKGllcUKSWpOxkf_2CpPytbNwmWGyBwm0lcnjm8MzhUNxasVWIK5J9IdlmwnvXaLOyvOB9yyeFrl5X3xtU0FuhtkByqvcKjW1ElxTcYpUUfV2jSSrkUuqSO6EVSTedETvuMKl7VfpHSfWqeCvK5DidpBtneiQ5JWwNrkHYokLDHVZwHsxCZ3TVlwgcKt0XEpPaIEKtjYfTYmuwJjmFw1LAzbZvUTkL-0aUDXCDoLQD2xcW_-hROfkKBl1vFFZTQjeEPg6_LCUshW84LGkuXuY0_mmlOLxgT4Q9wbff_kvSRwDwbxLdOSAsszBGTdKJDqVQ-CMcQeLQukTqPZrE6UTKXQskW4-DIIt1BFJ2fWJ6pdBAgsqZ16TTQrnEoO2lI-lGpOw8QY9j6n9gQAZkTod_PpdcSkTCloRl3GypX2vgPl3XKSPpV8IeICHp1-vnQBZf4grgiQlkAEk3w7hpeBAiP8BY2Dg1bth5gDdGksVmPCsypyELM6wVsBK29JP90keUPw_STzXowkS_1psMxmgxFGHLO5k8zyustePyHKTUvBogxHr-O6gnQmOUR6iPerikcFDYxogdGmi5UKdi82V4pbYR3v2cc9ZFyq5UYdCGTMK0A3UX2zVQdYxxRHvBBDfCNdO6c9qKU9yQGzgN4jTpNn9xm_-h5s_p-PqCZe-8HbpGWMAX3nYSoTTcNmiDlQ0ucU5KTmEvXBNeozHahHxy6s0spuUfRIeD4HAVOiy9HQoFruRlg8BITscc63vDHTS861DZK8_6fL8-GnXMtdZS6r2no9RV9GX_OO4dcFXBoPz08R07_bgBrSG-nYXtnJ1X-cjYOOBSh8FvT0IqtbKOKwch9w-bmJ_qt-E57sIAWNe1HUxjiIYvzvDSPVtnRIXVc4uOV9zxIT06tsS4Y_g0VYUvJ2s7W_-ZXiB4nt2N4S0LvguEj0DDUhGG-DPIcDpoMlb5CE8XyEdZZuuxQ0nUcNTNbIgV9pM9eBfwfnZmy-OBH25tY9jy9SGdUThidmUp75xH9Y0z3qHIjwnSlxK798i6IO-t02t9Z-2d4M_GjjFKWErfP8g-o6Y-rOfPqKqQ78-j-IeLKYpgfQA-G62R2VuNxdjB-v-9vmr9fesulEXjsPKBtPL_98fmrUovY41267-q64yy5U9Wf9Bq16A5sTJYyrts_FvUMprbfxQU2jVQcos2TPPXxAtcOQWDNRrr2zrPw0k3I2TAka1DP-NrR5TcaQPC-uFDdzOFRwemV060AXDA8YbgQtSbleSev9rYwvkVIyuw47I_gsQqNIO3OMeau_-FDdaFRbOLdXE-aN1wtQ3dV2y4BjaOXdmPtF1335B-nStSvBsjr6A2uj3hNjyw6hquDjemd69VEdLnXKzGVP8tXMytly8Hg7xKeO3QxO8duh6KIKKIn032DRq8eu7lbJ2Q8uSl13K7dU4v9N5fISxw6C1CWPnYyA-ynFSrtHpIH_gEV7MFnc_zlObzSbPK8xLzJRZ8XtVIF3xZLGqWpsuywDpbpulErBhlczqj2SzNFvP5lJXLLFsu8grzIqs4J3OKLRdyKuWunWqznQhre1wtFvQhn0heoLThIxVjQbmMkWwzMSs_PCn6rSVzKoV19hTACSfDh60wIduQ7MuF6flt3Jzdt_an71x3X4cmvZGrxrnO-joLatsK1_TFtNQtYU8ezPBX0hn9O5aOsKeQmiXsKWT3VwAAAP__20XLhA">