Transformations mutating read-only memref blocks

MLIR
#1

Hello all! :slight_smile:

It seems there are two bufferization-related transformations in MLIR whose correctnesses are unclear or wrong.
If indeed wrong, I think there are low-cost solutions hopefully; I wrote my thoughts as well.

  1. -linalg-bufferize introduces linalg.fill to a read-only block:
func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %c0 = arith.constant 0.0 : f32 
  %0 = linalg.fill(%c0, %arg0) : f32, tensor<?xf32> -> tensor<?xf32>
  return %0 : tensor<?xf32>
}
=>
func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %0 = bufferization.to_memref %arg0 : memref<?xf32>
  // ^^^ According to BufferizationOps document, mutating %0 is UB.
  %cst = arith.constant 0.000000e+00 : f32
  linalg.fill(%cst, %0) : f32, memref<?xf32> 
  // ^^^ This is mutating %0 (hence UB)?
  %1 = bufferization.to_tensor %0 : memref<?xf32>
  return %1 : tensor<?xf32>
}

Perhaps a straightforward solution would be to add writable attribute to bufferization.to_memref.
If it is set to true, the returned buffer is writable. This will make -linalg-bufferize correct.

  1. -buffer-deallocation introduces dellocating read-only blocks, which affects the lifetime of the original block if the bufferization.clone is returning an aliased pointer:
func @dealloc_existing_clones(%arg0: memref<?x?xf64>, %arg1: memref<?x?xf64>) -> memref<?x?xf64> {
  %0 = bufferization.clone %arg0 : memref<?x?xf64> to memref<?x?xf64>
  %1 = bufferization.clone %arg1 : memref<?x?xf64> to memref<?x?xf64>
  return %0 : memref<?x?xf64>
}
=>
func @dealloc_existing_clones(%arg0: memref<?x?xf64>, %arg1: memref<?x?xf64>) -> memref<?x?xf64> {
    %0 = bufferization.clone %arg0 : memref<?x?xf64> to memref<?x?xf64>
    %1 = bufferization.clone %arg1 : memref<?x?xf64> to memref<?x?xf64>
    memref.dealloc %1 : memref<?x?xf64> // <- THIS ONE
    return %0 : memref<?x?xf64>
  }

To resolve this, we can add another operation bufferization.dealloc %p that
(1) is no-op if the %p is an aliased pointer
(2) is equivalent to memref.dealloc %p otherwise.

Thanks!

1 Like