Very WIP: Refactor core inference loops to use less memory #43999
Conversation
This file contains 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
aviatesk
force-pushed
the
kf/inferencerefactor
branch
from
9ae2a51 to
36e3abe
Compare
|
Rebased against the latest master, and dropped changes from #43899 (I believe these two PRs are orthogonal and can be merged separately). Still there seems to be some problems in this PR. Will try to dig into them. |
aviatesk
force-pushed
the
kf/inferencerefactor
branch
4 times, most recently
from
2ec47da to
60a82a6
Compare
Currently inference uses `O(<number of statements>*<number of slots>)` state in the core inference loop. This is usually fine, because users don't tend to write functions that are particularly long. However, MTK does generate functions that are excessively long and we've observed MTK models that spend 99% of their inference time just allocating and copying this state. It is possible to get away with significantly smaller state, and this PR is a first step in that direction, reducing the state to `O(<number of basic blocks>*<number of slots>)`. Further improvements are possible by making use of slot liveness information and only storing those slots that are live across a particular basic block. The core change here is to keep a full set of `slottypes` only at basic block boundaries rather than at each statement. For statements in between, the full variable state can be fully recovered by linearly scanning throughout the basic block, taking note of slot assignments (together with the SSA type) and NewVarNodes. The current status of this branch is that the changes appear correct (no known functional regressions) and significantly improve the MTK test cases in question (no exact benchmarks here for now, since the branch still needs a number of fixes before final numbers make sense), but somewhat regress optimizer quality (which is expected and just a missing TODO) and bootstrap time (which is not expected and something I need to dig into).
aviatesk
force-pushed
the
kf/inferencerefactor
branch
from
60a82a6 to
dfeb2e9
Compare
aviatesk
reviewed
Apr 26, 2022
Comment on lines
+1816
to
+1818
| if s.undef | ||
| sv.src.slotflags[sn] |= SLOT_USEDUNDEF | SLOT_STATICUNDEF | ||
| end |
There was a problem hiding this comment.
After digging into the regressions in the optimizer, I found this change could be problematic, and especially causes some regressions in the SROA pass.
I think these changes are orthogonal to the main purpose of this PR, so I am trying to fix regressions by separating related changes into another branch.
Still there seems to be another cause of the regressions though.
aviatesk
added a commit
that referenced
this pull request
Apr 26, 2022
Separated from #43999. xref: <#43999 (comment)>
aviatesk
added a commit
that referenced
this pull request
May 4, 2022
Separated from #43999. xref: <#43999 (comment)>
aviatesk
added a commit
that referenced
this pull request
May 4, 2022
Separated from #43999. xref: <#43999 (comment)>
aviatesk
added a commit
that referenced
this pull request
May 4, 2022
Separated from #43999. xref: <#43999 (comment)>
aviatesk
added a commit
that referenced
this pull request
May 4, 2022
Separated from #43999. xref: <#43999 (comment)>
aviatesk
added a commit
to aviatesk/julia
that referenced
this pull request
May 10, 2022
Separated from JuliaLang#43999. xref: <JuliaLang#43999 (comment)>
|
Replaced by #45276 . |
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.
Currently inference uses
O(<number of statements>*<number of slots>)statein the core inference loop. This is usually fine, because users don't tend
to write functions that are particularly long. However, MTK does generate
functions that are excessively long and we've observed MTK models that spend
99% of their inference time just allocating and copying this state.
It is possible to get away with significantly smaller state, and this PR is
a first step in that direction, reducing the state to
O(<number of basic blocks>*<number of slots>).Further improvements are possible by making use of slot liveness information
and only storing those slots that are live across a particular basic block.
The core change here is to keep a full set of slottypes only at
basic block boundaries rather than at each statement. For statements
in between, the full variable state can be fully recovered by
linearly scanning throught the basic block, taking note of
slot assignments (together with the SSA type) and NewVarNodes.
The current status of this branch is that the changes appear correct
(no known functional regressions) and significantly improve the MTK
test cases in question (no exact benchmarks here for now, since
the branch still needs a number of fixes before final numbers make
sense), but somewhat regress optimizer quality (which is expected
and just a missing TODO) and bootstrap time (which is not expected
and something I need to dig into).