As a follow-up to the previous discussion around extending one-shot bufferization support to user-specified types, I would like to continue along this path as I am in the process of using the newly added TensorLike and BufferLike type interfaces inside of the implementation.
In particular, my proposal is this: Bufferization’s operations (to_tensor, to_memref, clone but also alloc_tensor / dealloc_tensor?) need to be extended to work with TensorLike / BufferLike.
I can justify this for to_tensor and to_memref - from my perspective, these are the “unrealized_conversion_casts” with more semantics.
I can not really justify it for other ops however: tensor / memref cloning kind of depends on what is a tensor / memref - so user-specific; allocation / deallocation is also something that is user-specific?
In general, the pitfall to me is this:
if bufferization ops remain “as is”, users have to provide a complete set of similar ops themselves → this suggests introducing op interfaces and bufferization starts to rely on interfaces instead of actual operations
downside: every new op has to have an associated interface
this also ultimately renders the new type interfaces useless (TensorLike / BufferLike)
if bufferization ops change, it is not really clear what is “tensor/memref allocation”, “tensor/memref copy”, etc.
perhaps this is fine? - i.e. downstream projects could lower from bufferization-specific ops to “proper” ops separately, after one-shot bufferization is done
I do not think that “more customization points” (i.e. new callbacks in bufferization options) would help here as the underlying issue is still the same: “roll your own ops” vs “support custom types in existing bufferization ops”.
This RFC is to discuss whether changing operations (effectively, their semantics) is viable and hopefully become aware of potential issues.
Tensor copy can be just to_tensor(to_memref(x)), so IMO there is no difference between allowing these two ops support custom types and not allowing copies. All of them are similarly “aware” of the underlying type structure.
alloc_tensor is inserted by the bufferization framework. See TensorCopyInsertion.cpp. That’s an internal pass that brings the IR into a form where no further analysis is needed and all tensors can be directly replaced with memrefs.
So I think you have to extend alloc_tensor (and for consistency also dealloc_tensor) in the same way. The implementation of AllocTensorOp::bufferize calls lambdas from BufferizationOptions to insert the buffer alloc/copy ops. These lambdas must create the correct buffer ops. That’s where the user-specific code is.
Is there a problem with this approach?
if bufferization ops change, it is not really clear what is “tensor/memref allocation”, “tensor/memref copy”, etc.
The naming of alloc_tensor may not be ideal. It returns a tensor that is guaranteed to bufferize to a new buffer. I.e., the future buffer of the alloc_tensor result is guaranteed to be distinct from the future buffer of the copy operand (or any other “previous” buffer). In essence, this op gives you a tensor that will turn into buffer that doesn’t alias with anything (up to that point).
Which additional callbacks would be needed? I think we already have everything that we need.
I haven’t looked into this function in detail, but you probably assume that shapedValue has a type that implements either TensorLike or BufferLike. That should make it possible to implement this function in a generic way. (By querying type interface methods to convert from tensor-like → buffer-like types.)
Not really. I am just not sure upstream “allocation” makes sense for user types. In any case, thanks for the hints, I think it does indeed look “expected” to have alloc_tensor / dealloc_tensor extended.
What about dealloc and clone (both operate on memrefs)? They seem pretty memref-related to me (i also see that options.getMemCpy ends up producing memref::CopyOp by default so i’m not even sure why clone is special).
This way we hide the TensorType / BaseMemRefType details behind options. But I’ll need to dwell on this as I’m not sure if it’s good in general (there are multiple places where the getMemRefType is being used).
So far, I can ignore the problem in one-shot-bufferize (at least) by assuming shapedValue is a builtin tensor/memref. Generally though, you’re correct.
Do you propose to introduce TensorLikeType::getBufferType and perhaps a handful of other interface methods? Doing it this way might end up reducing the need for certain option-level hooks (and maybe eliminate op-level interfaces methods e.g. BufferizableOpInterface::getBufferType) and also actually allow one to use one-shot-bufferize pass directly. For instance, the current issue is that one needs to set options manually in order to use: function boundary bufferization, allocateTensorForShapedValue and likely other primitives that rely on the “not BufferizableOpInterface then call getMemRefType” kind of model.
Edit: for “tensor-like → buffer-like types” conversion, I guess we also need to agree on whether this should be external or internal w.r.t. the types. External: smth.getBufferType(TensorLike) → BufferLike; internal: TensorLike.getBufferType() → BufferLike.
I am still thinking about TensorLike / BufferLike to be eventually hoisted to builtins. If that ever to be considered, having “internal conversion” model is a hard blocker.
bufferization.clone is used by the buffer deallocation pass. I don’t think we use it anywhere else. I’m also not sure why we need the op at all. We could probably remove it and replace it with memref.alloc + memref.copy.
It’s been too long since I looked at this… unknownTypeConverterFn does not take a layout, so I’m not sure if it can we wired like that. I would look into the places that are calling unknownTypeConverterFn. What if you put a failed assertion in the lambda? What tests are failing? I think this lambda is kind of a “fallback” when we don’t know which type to use. But when is that actually the case? I don’t remember…
getMemRefType produces a memref type from a tensor type. Now that we have multiple buffer types, this sounds like a candidate for a type interface method to me.
yes
I think you’re still going to need that. This function is predicting the future buffer type for a tensor result. E.g., for tensor.extract_slice, the bufferized result must have a certain layout map. We need to predict this type when bufferizing a loop op: if I remember correctly, the loop op is bufferized before the loop body.
Inside of BufferizableOpInterface::getBufferType, the builtin MemRefType can be hard-coded. This is an interface method on a specific operation. E.g., memref.subview (bufferized operation of tensor.extract_slice) supports only MemRefType and no custom buffer types, so it’s OK to hardcode MemRefType.
What’s the difference? The first one is a member function, the second one is a static function? What is smth?
I see. So the BufferizableOpInterface::getBufferType is a building block of BufferizableOpInterface::bufferize kind of and simultaneously the method to preliminary query the output buffer without actually bufferizing. I guess the renewed logic is then something like:
Try BufferizableOpInterface::getBufferType if op supports BufferizableOpInterface
Go to “new getBufferType API” otherwise
(the cases could likely be enclosed into the free-standing getBufferType function or something along these lines)
Do you mean at API level or in implementation? I guess the latter does make sense indeed. I’d still prefer BufferLikeType BufferizableOpInterface::getBufferType(TensorLikeType, ...) at the signature level though.
Yes. Member function - clear enough (again, main concern is that we couple “types” and “operations” on these types). Non-member function:
converter.getBufferType(TensorLike) → BufferLike (smth - some converter object that is a DialectInterface that one just creates for user-specified tensor/memref - they likely live in own dialect anyway)
this has to live inside the context forever though (so slightly higher memory usage)
From the standpoint of custom tensors / memrefs support, options is probably most cumbersome because they require one to basically reimplement one-shot-bufferization pass itself (not the underlying mlir::runOneShotModuleBufferize() function with the actual logic though) - because that’s the only way to “seed” options object? For instance, this is what we do in our downstream and it works fine, I am not against this in general.
One more thing I just remembered, our downstream also extends builtin tensor (via encoding) and memref (via layout interface).
In general, this would require a customization point to overwrite builtin.tensor → builtin.memref conversion as well [1] - we want to be able to convert encoding to layout in a particular fashion ourselves. Which is why having TensorLikeType::getBufferType might be problematic - the TensorLike interface is not “promised” for builtins but hard-attached instead (see details) - so supplying a custom implementation is hard-ish.
[1]: I think so far we manage to avoid solving this issue on our end by a combination of: options.copyBeforeWrite = false; options.testAnalysisOnly = true; (these two luckily “turn off” any tensor copies and allocations during one-shot bufferization). Any Op::bufferize implementation also uses our own “get buffer type” method that also handles builtins.