Refactor updates functionality and Scan inner-graph compilation
#848
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
brandonwillard
added
enhancement
brandonwillard
force-pushed
the
more-inner-graph-updates
branch
5 times, most recently
from
028787a to
18c1b10
Compare
brandonwillard
force-pushed
the
more-inner-graph-updates
branch
14 times, most recently
from
bda1cd2 to
a412af3
Compare
Codecov Report
@@ Coverage Diff @@
## main #848 +/- ##
==========================================
+ Coverage 78.99% 79.02% +0.02%
==========================================
Files 152 152
Lines 47809 47947 +138
Branches 10893 10900 +7
==========================================
+ Hits 37768 37889 +121
- Misses 7541 7547 +6
- Partials 2500 2511 +11
|
brandonwillard
changed the title
Scan inner-graph compilation
brandonwillard
force-pushed
the
more-inner-graph-updates
branch
from
a412af3 to
c5f4093
Compare
These changes allow one to pass an `fgraph` argument to all key functions in the `Function` compilation pipeline. The given `FunctionGraph` will be directly compiled without cloning. Unlike the previous `FunctionMaker.__init__`, this one's `fgraph` argument *will* be optimized according to the given `mode` unless the keyword argument `no_fgraph_prep` is `True`.
A few core `Feature` implementations were using `assert`s and generic `Exception` instead of `AlreadyThere`, which unnecessarily produces errors when already-attached `Feature` requirements are added by `Rewriter`s.
This also prevents `SeqOptimizer.apply` from calling `Rewriter.add_requirements` on each iteration through its list of sub-`Rewriter`s.
This changes the `update_storage` parameter from a list containing the input indices that are to be updated with the last N-many outputs to a tuple of tuples specifying input/output indices. Now, arbitrary output-to-input update pairings are possible, instead of forcing graphs and code to compensate for this unnecessary restriction.
In order to use shared updates to pre-allocate the storage for mit-mot input and output loops, `Scan` would need to remove the corresponding mit-mot outputs from its inner-`FunctionGraph` before compilation and it would expect the `Function` compilation pipeline to add them back at the end of the remaining outputs. Now, `Scan`'s inner-`FunctionGraph`s maintain the same form at every point, and no special logic is needed to compensate for post-compilation changes in the order/location of inputs and outputs.
All VM implementations now perform variable updates themselves. This leaves some now redundant update code in `Function`, but it also removes some from `Scan.perform`.
This new method provides an interface for stateful `Feature`s to be easily cloned and attached to other `FunctionGraph`s.
This method adds the ability to clone a `HasInnerGraph` `Op` and its inner-`FunctionGraph`.
Since it is no longer necessary to clone `Constant`s, and they add extra work for the merge rewrites, `clone_get_equiv`'s `copy_orphans` option has been prevented from cloning `Constant`s. The option is mostly applicable to constants, but it has been retained for other non-`Constant` cases (just in case) and backward compatibility.
This adds the missing `update_mapping`s step during cloning, as well as a new `Feature` cloning step that prevents issues when features are copied to their clones.
brandonwillard
force-pushed
the
more-inner-graph-updates
branch
from
36e29a4 to
fba4184
Compare
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.
This PR consists of independent
FunctionGraph,Scan, andFunctionconstruction changes that have been split off from #824.One of the main results of these changes is that
Scan's inner-graph is handled almost exclusively as a single coherentFunctionGraph, and the inner-graph compilation andScan.performlogic is much more concise.For instance, when pre-allocation is applied,
Scancurrently removes/reorders some of its inner-graph outputs just before compiling. This is necessary because theFunctioncompilation pipeline is somewhat rigid with respect to the way it handles updates (e.g. it assumes that all update steps are at the end of the outputs list of aFunctionGraph). Those issues have been fixed in this PR and, now, aScan's inner-graph is consistent for the lifetime of theScanand requires no manipulation before compilation.This PR also adds some long-overdue refactoring of the
VMinterface and implementations. Previously, theCVMwas the onlyVMthat would perform updates, andFunctionwould perform the updates for all otherVMs. This meant that the updating logic needed to be reproduced inScan.perform(see below)—mostly becauseScan.performwas (perhaps unnecessarily) designed to use aFunction'sVMand not theFunctionitself. Now, everyVMimplementation performs updates, so there's no need for that logic/scope leak inScan(orFunctionfor that matter, but I've left it in for now).Scan.perform(and its Cython implementation).This would be another big step toward simplifying the
Op.performlogic inScan.