Lunderberg opened a new pull request #9091: URL: https://github.com/apache/tvm/pull/9091
This started after noticing that StorageFlatten incorrectly handled BufferLoad/BufferStore nodes that pointed to a buffer defined in an `attr::buffer_bind_scope` annotation. Rather than adding more logic into the existing `StorageFlattener` mutator, I split up the existing behavior into multiple independent mutators. This PR includes a series of commits, each of which refactors one of the behaviors out of the `StorageFlattener` class and into a separate class. While all of the transforms are called sequentially in the `tir.transform.StorageFlatten` to maintain the same overall behavior, each transform results in a valid TIR tree. * BufferShapeLegalize, which rewrites Buffer nodes to have sizes that match the BufferRealize node in which they are defined. * BufferStrideLegalize, which rewrites the strides of Buffer nodes that are annotated with `attr::dim_align`. * ThreadScopePropagate, which defines the allocation scope of Buffer nodes based on the thread iter in which they are declared, if no allocation scope was already defined. * BufferBindUnwrapper, which rewrites access into Buffer objects that are defined by `attr::buffer_bind_scope`. Refactoring this behavior into a separate mutator was my original goal, in order to resolve the issue of BufferLoad/BufferStore nodes that point to bound buffers, but doing so required the previous three behaviors to also be refactored into separate mutators. * StorageFlattener, which contains all remaining behavior from the original StorageFlattener, and outputs the final Allocate/Store/Load nodes. This refactor will also help in the future, when introducing layout transformations. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
