Switch to one-shot bufferization from MLIR upstream #169
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nikita-kud
approved these changes
Nov 6, 2025
hrotuna
approved these changes
Nov 7, 2025
Contributor
hrotuna
left a comment
hrotuna
left a comment
There was a problem hiding this comment.
Thank you for handling this! Do you know if all of these patches have already landed in LLVM 21?
Contributor
Author
Not all. I think a couple of them are in LLVM 21, the rest are in LLVM 22. |
andrey-golubev
force-pushed
the
agolubev/new_bufferization
branch
from
37ba4b2 to
9c1649d
Compare
This reverts commit 51f2bc8.
This reverts commit aef8177.
…4220) Current one-shot bufferization infrastructure operates on top of TensorType and BaseMemRefType. These are non-extensible base classes of the respective builtins: tensor and memref. Thus, the infrastructure is bound to work only with builtin tensor/memref types. At the same time, there are customization points that allow one to provide custom logic to control the bufferization behavior. This patch introduces new type interfaces: tensor-like and buffer-like that aim to supersede TensorType/BaseMemRefType within the bufferization dialect and allow custom tensors / memrefs to be used. Additionally, these new type interfaces are attached to the respective builtin types so that the switch is seamless. Note that this patch does very minimal initial work, it does NOT refactor bufferization infrastructure. See https://discourse.llvm.org/t/rfc-changing-base-types-for-tensors-and-memrefs-from-c-base-classes-to-type-interfaces/85509
As part of the work on transitioning bufferization dialect, ops, and associated logic to operate on newly added type interfaces (see 00eaff3), rename the bufferization.to_memref to highlight the generic nature of the op. Bufferization process produces buffers while memref is a builtin type rather than a generic term. Preserve the current API (to_buffer still produces a memref), however, as the new type interfaces are not used yet.
… (#144658) Generally, bufferization should be able to create a memref from a tensor without needing to know more than just a mlir::Type. Thus, change BufferizationOptions::UnknownTypeConverterFn to accept just a type (mlir::TensorType for now) instead of mlir::Value. Additionally, apply the same rationale to getMemRefType() helper function. Both changes are prerequisites to enable custom types support in one-shot bufferization.
Following the addition of TensorLike and BufferLike type interfaces (see 00eaff3), introduce minimal changes required to bufferize a custom tensor operation into a custom buffer operation. To achieve this, new interface methods are added to TensorLike type interface that abstract away the differences between existing (tensor -> memref) and custom conversions. The scope of the changes is intentionally limited (for example, BufferizableOpInterface is untouched) in order to first understand the basics and reach consensus design-wise. --- Notable changes: * mlir::bufferization::getBufferType() returns BufferLikeType (instead of BaseMemRefType) * ToTensorOp / ToBufferOp operate on TensorLikeType / BufferLikeType. Operation argument "memref" renamed to "buffer" * ToTensorOp's tensor type inferring builder is dropped (users now need to provide the tensor type explicitly)
…e (#144867) Support custom types (2/N): allow value-owning operations (e.g. allocation ops) to bufferize custom tensors into custom buffers. This requires BufferizableOpInterface::getBufferType() to return BufferLikeType instead of BaseMemRefType. Affected implementors of the interface are updated accordingly. Relates to ee070d0.
…59766) Support custom types (3/N): allow custom tensor and buffer types in function signatures and at call-sites. This is one of the major building blocks to move in the direction of module-level one-shot-bufferization support. To achieve this, `BufferizationOptions::FunctionArgTypeConverterFn` callback is converted to work with tensor-like and buffer-like types, instead of the builtin counterparts. The default behavior for builtins remains unchanged, while custom types by default go through `TensorLikeType::getBufferType()` which is a general conversion interface.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Over the course of the year, upstream MLIR was gradually updated to allow custom types in one-shot bufferization. The result of it is a collection of patches that, from my perspective, should be sufficient to allow us to use "upstream" MLIR directly, without needing to maintain our own patches. This PR summarizes all of this work.
Notable changes:
JIRA ticket
Related PR in NPU Compiler and/or OpenVINO repository with sub-module update
Other related tickets