Lunderberg opened a new pull request #9194: URL: https://github.com/apache/tvm/pull/9194
This PR replace the `AllocateNode::extents` variable, which represented the logical N-d shape of the tensor, with a 1-d `AllocateNode::extent` variable representing the amount of memory to be allocated. This is part of a larger push to support more flexible memory layouts across multiple device types, and to introduce a split between the physical layout in memory and the logical layout in a tensor or a compute definition. The commits of this PR are divided up into discrete sections to hopefully make it easier to review. 1. Change `AllocateNode::extents` (plural) to `AllocateNode::extent` (singular). This first commit is the motivation for the PR, and the remaining commits are the additional changes needed to be compatible with this first change. 2. C++ changes necessary to compile against the new `AllocateNode` definition. 3. A functionality addition to `IRBuilder` to output `BufferRealize`/`BufferLoad`/`BufferStore` nodes, in addition to `Allocate`/`Load`/`Store`. This minimizes the changes that will be needed in the unit tests, as they can continue using N-d indices and allocations with `BufferRealize`/`BufferLoad`/`BufferStore`, which are then flattened during the `StorageFlatten` pass. 4. Update topi schedules to use 1-d extent. 5. Update python unittests to use 1-d extent. 6. Simplify() call added to workspace size calculation, needed for the workspace calculation unit tests. 7. Updated UnrollLoop to count `BufferStore` in addition to `Store`. This resolves a breakage in the tests for unroll loop, and allows UnrollLoop to be used either before or after `StorageFlatten`. 8. Refactor of the tests for unroll loop. Done as part of the investigation for previous point, but kept as an independent commit for readability. -- 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]
