[TIR] 1-d AllocateNode::extent #9194
Closed
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
Lunderberg
requested review from
Huyuwei,
Hzfengsy,
Laurawly,
MarisaKirisame,
ZihengJiang,
anijain2305,
areusch,
comaniac,
jcf94,
jroesch,
junrushao,
jwfromm,
kevinthesun,
kparzysz-quic,
masahi,
mbrookhart,
merrymercy,
slyubomirsky,
tmoreau89,
tqchen,
vegaluisjose,
vinx13,
wweic,
yzhliu and
zhiics
as code owners
Contributor
Author
Contributor
Lunderberg
force-pushed
the
allocate_node_1d
branch
from
b4f2b37 to
b162288
Compare
Contributor
Author
|
I've started a PR for an RFC that describes the motivation for this change in more detail. This change to |
Lunderberg
force-pushed
the
allocate_node_1d
branch
from
8771fc8 to
623374c
Compare
Contributor
Author
|
Updated to resolve conflicts with main, change |
Lunderberg
force-pushed
the
allocate_node_1d
branch
2 times, most recently
from
756ce1d to
73a854d
Compare
This commit has the primary goal of the PR, replacing the AllocateNode::extents variable, which represented the 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.
After updating to a 1-d index in AllocateNode::extent, this commit adds all the secondary changes needed elsewhere in the codebase to be compatible with the 1-d extent.
Allow the IRBuilder to output BufferRealize/BufferLoad/BufferStore nodes, as well as Load/Store. This minimizes the changes that will be needed in the unit tests, as they can continue using N-d indices and allocations, which are then flattened during the `StorageFlatten` pass.
Lunderberg
force-pushed
the
allocate_node_1d
branch
2 times, most recently
from
0e294bf to
b12082d
Compare
Some of the updated unit tests express allocation size as a product rather than as a single integer, for readability, but those products should still be treated as constant integers when determining the workspace size needed.
If UnrollLoop is applied before StorageFlatten, there may be BufferLoad/BufferStore nodes that haven't yet been lowered to Load/Store nodes. These should be counted equivalently to Store nodes.
These changes were not necessary for the UnrollLoop change in the previous commit to function, but were made as part of that investigation. They are kept as a separate commit in the PR for ease of review.
Previous test of CopyOnWrite semantics during StmtMutator relied on AllocateNode having an array parameter. Rewrote these tests to check the same behavior, but implemented over SeqStmtNode instead.
Lunderberg
force-pushed
the
allocate_node_1d
branch
from
b12082d to
9b1d6e0
Compare
Contributor
Author
|
Updated to resolve breaks this caused in ethos unit tests. |
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Oct 22, 2021
Currently in PR at apache#9194, waiting on RFC apache/tvm-rfcs#40. Locally, this is a merge commit for the 'allocate_node_1d', to help me keep track of which commits belong in which PRs.
Contributor
Author
|
Closing this PR, as changes resulting from discussion on apache/tvm-rfcs#39 has made this change no longer desirable. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
This PR replace the
AllocateNode::extentsvariable, which represented the logical N-d shape of the tensor, with a 1-dAllocateNode::extentvariable 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.
Change
AllocateNode::extents(plural) toAllocateNode::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.C++ changes necessary to compile against the new
AllocateNodedefinition.A functionality addition to
IRBuilderto outputBufferRealize/BufferLoad/BufferStorenodes, in addition toAllocate/Load/Store. This minimizes the changes that will be needed in the unit tests, as they can continue using N-d indices and allocations withBufferRealize/BufferLoad/BufferStore, which are then flattened during theStorageFlattenpass.Update topi schedules to use 1-d extent.
Update python unittests to use 1-d extent.
Simplify() call added to workspace size calculation, needed for the workspace calculation unit tests.
Updated UnrollLoop to count
BufferStorein addition toStore. This resolves a breakage in the tests for unroll loop, and allows UnrollLoop to be used either before or afterStorageFlatten.Refactor of the tests for unroll loop. Done as part of the investigation for previous point, but kept as an independent commit for readability.