LLVM Discussion Forums

Towards a generic Buffer Assignment (BA) Phase

Hi all. We are currently working on a generic BufferAssignment (BA) pass (this is also related to the post Generic Alias Analysis). During our initial integrations into the HLO->LHLO and HLO->Linalg (with buffers) legalizations we stumbled upon the general question whether a buffer is read only or write enabled. If we “move” a value into a buffer, we have to ensure that it will not be changed unexpectedly by a dialect-specific operation. Consider the following code snippet and its lowered version below:

module {
  func ....(%arg0: tensor<f32>) -> ... {
    %0 = "my_amazing_op" (%arg0) : tensor<f32> -> tensor<f32>
    ...
  }
}

Lowered version:

module {
  func ....(%arg0: memref<tensor<f32>>, %arg1: memref<tensor<f32>>) -> ... {
    "my_amazing_write_to_all_buffers_op" (%arg0, %arg1) :
        (memref<tensor<f32>>, memref<tensor<f32>>)
    ...
  }
}

Originally, the operation “my_amazing_op” received a tensor and returned an updated version. However, its lowered version (comparable to an LHLO operation that corresponds to an HLO operation) has “different” semantics and writes into all buffers. This will basically invalidate the contents of %arg0 in this example and the lowering will be invalid. Therefore, it is important to know which buffers are read only and which of them are write enabled. In this example, the following code snippet would fix the issue:

module {
  func ....(%arg0: memref<tensor<f32>>, %arg1: memref<tensor<f32>>) -> ... {
    %0 = "alloc" (...) : ...
    "copy_op"(%arg0, %0) : ...
    "my_amazing_write_to_all_buffers_op" (%0, %arg1) :
        (memref<tensor<f32>>, memref<tensor<f32>>)
    ...
  }
}

You could argue that the BA phase should insert the copies automatically. As outlined above, this is not trivial without knowing details about the semantics of the operations. We already had a discussion with @herhut and thought about this challenge. We came the conclusion that it might not be beneficial to annotate all operations with some “special flags”. Instead, a dialect expert that lowers a certain dialect into another dialect with buffers should take care about this problem “somehow”.

The general question is whether to provide an interface (again similar to “addIllegalOp”) to give dialect-specific information about read only and write-enabled arguments to a generic BA converter. Alternatively, you can argue that it is required to insert copies “manually” during the legalization phase.

We added the SideEffect interface specifically for this purpose: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Interfaces/SideEffects.td ; I’m interested to understand better what is missing here.

I had a closer look at the interface SideEffects. Looks pretty complete to me. You have introduced Alloc, Free, Read and Write effects. For certain programs and problems it might also be beneficial to think about Copy semantics in relation to buffers. However, this is usually achieved by introducing dialect-specific mydialect.copy operations, which can be easily analyzed without having a particular effect.