[Mlir-commits] [mlir] [mlir][bufferization] Make alloc_tensor's bufferization customizable (PR #215590)

[Victor Perez]

victor-eds wrote:

> What's the problem here? Can we generalize the "bufferize" implementation?

This can be done already via the `allocationFn` override I mentioned. The problem is, if we want to change the type of the resulting allocation, this may lead to problems. For instance, let's say we want the resulting buffer to have a different layout than the identity (something we can do without this PR):

```mlir
%0 = bufferization.alloc_tensor() : tensor<8x16xf32>  # bufferized to memref<8x16xf32, #layout>
```

Operations using the result of `alloc_tensor` can now query the type of the allocation _before actually bufferizing_ using `getBufferType`. However, as that cannot be customized, and we haven't called `alloc_tensor`'s `bufferize` method yet, the analysis will infer that the type of the buffer is `memref<8x16xf32>`. This way, if I have something like:

```mlir
%0 = bufferization.alloc_tensor() : tensor<8x16xf32> # bufferized to memref<8x16xf32, #layout>
%1 = tensor.collapse_shape %0 [[0, 1]] : tensor<8x16xf32> into tensor<128xf32>
```

We qill create an illegal shape transformation operation when bufferizing `collapse_shape`, as the analysis has told us that the type carries the identity layout:

```mlir
%alloc = memref.alloc() : memref<8x16xf32, #layout>
%1 = memref.collapse_shape %alloc [[0, 1]] : memref<8x16xf32, #layout> into memref<128xf32>  # may be illegal
```

> External models were not designed to be "swappable". Personally, I wouldn't mind merging this because it make the infrastructure a bit more uniform (all `BufferizableOpInterface` implementations are external models, with the same file name in each dialect). But I'm wondering if something could break when there are multiple interface implementations. E.g. one transformation registers a certain impl, another one registers another impl.

This is an interesting topic. Yes, there is no way to swap these implementations, as, once registered, there is no way to override them for now. However, different projects can simply not take the upstream ones and use their own versions. As of now, for this operation, that is not a possibility. However, this aims to enable it :)

https://github.com/llvm/llvm-project/pull/215590