[WIP-DO NOT REVIEW] Equivalence tracking in Fusion#382
Closed
jacobhinkle wants to merge 23 commits intoNVIDIA:mainfrom
Closed
[WIP-DO NOT REVIEW] Equivalence tracking in Fusion#382jacobhinkle wants to merge 23 commits intoNVIDIA:mainfrom
jacobhinkle wants to merge 23 commits intoNVIDIA:mainfrom
Conversation
Still needs converting copy ctors and assignment operators that check bounds when converting.
This will ultimately handle simplifying expressions, but for now just wraps a UnionFind to track equivalence between vals of a single ValType.
This allows Vals to track their position in the originating IrContainer's vals() vector.
This is important so that we can merge Scalars with NamedScalars, which is not possible if we restrict by ValType.
This includes a simple test. The implementation is probably prematurely optimized and not very well.
Not working yet of course
This is just the start of an approach to equality saturation in NVFuser
This was referenced May 25, 2023
This was referenced Jun 8, 2023
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
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.
There are multiple stages in which we use the notion of equivalence of
Vals. For example, both scheduling and lowering use root domain maps which map rfactorIterDomains of producers to rootIterDomains of consumers. Later in lowering (but increasingly in other stages as well see #320),simplifyExprdoes term rewriting which is based on exploring spaces of equivalent expression trees and doing replacement based on complexity heuristics.Meanwhile, we also need to prove things, such as whether a predicate is always true unconditionally. This is also increasingly common since the introduction of dynamic fusions. For example, as mentioned in #380 it is difficult currently to check that the expanded extent of a concretized
IterDomainis equal to 1, since we must do so without using anExpressionEvaluatorin that case.In order to increase the power of our simplification and equivalence checking ability, this PR makes the following key changes:
.name()which numbers each val by itsValType, eachValnow has a.number()which is independent of val type, andFusioncan look up anyVal*from its number.UnionFindclass is introduced which is a very efficient class for tracking equivalence classes. https://en.wikipedia.org/wiki/Disjoint-set_data_structure. This is similar toDisjointSetsand may some day be a valid replacement for it, but the focus ofUnionFindis on quick merging and pairwise comparison, with slower extraction of entire equivalence classes.ValEquivalenceclass is introduced which plumbs the.number()of each val into aUnionFindand implements some methods for simplification/selection ofVals. An object of this class can quickly attest or test equivalence of any twoVals even if they are of different types, such asIntandNamedScalar. This forms the basis of a term rewriting system. TheFusionholds a singleValEquivalenceobject representing the weakest form of equivalence (TODO).ExprEquivalenceclass is introduced which is based onValEquivalence. It allows us to propose "sketches" of newExprs and check whether anExprwhose inputs and attributes are all equivalent already exists. This will be useful functionality for implementing simplification/optimization via equality saturation as it helps to limit the number of createdVals andExprs when exploring rewrites. This could be useful forsimplifyExpr, for passes like reduce-broadcast fusions in lowering, or for other graph optimization passes such as in Cast opt pass #355.Vals at any point. NOT YET STARTEDAll of these are currently an early work in progress.