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

[Andrei Golubev]

andrey-golubev wrote:

> > 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 :)

Funnily enough, before talking to Matthias on this exact thing some time ago, I thought the same as Victor: "one can replace an upstream external model with their own". As a matter of fact, this is something we sometimes do in our downstream. Generally, it's a bad idea (e.g. synchronizing with upstream changes could be painful) but it can work. 

> 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.

Last time I checked the actual interface addition logic, the new interface would gracefully replace the old one (this is just how "add interface" machinery works behind the scenes), which is quite convenient.


> > 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>
> ```

Our project is also quite a heavy user of custom layout information and we've suffered with it for quite a bit. I wonder if the problem you're describing is something we've also seen (asking @Devjiu to take a look). Overall, I agree with Matthias. I treat external model replacement as a "last resort" option and it is definitely not something that would work solidly in the long run imho. I don't get the full problem statement but perhaps there could be a good addition to the bufferization infrastructure to make your case work without requiring to rewrite the bufferization model completely? We've had success with SCF recently by introducing new bufferization hooks for example (https://github.com/llvm/llvm-project/pull/202667 and https://github.com/llvm/llvm-project/pull/205080), perhaps this case here could also be treated similarly.

> 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).

+1 to this.

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