Skip to content

Comments

[WIP-DO NOT REVIEW] Equivalence tracking in Fusion#382

Closed
jacobhinkle wants to merge 23 commits intoNVIDIA:mainfrom
jacobhinkle:graph_indexing
Closed

[WIP-DO NOT REVIEW] Equivalence tracking in Fusion#382
jacobhinkle wants to merge 23 commits intoNVIDIA:mainfrom
jacobhinkle:graph_indexing

Conversation

@jacobhinkle
Copy link
Collaborator

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 rfactor IterDomains of producers to root IterDomains of consumers. Later in lowering (but increasingly in other stages as well see #320), simplifyExpr does 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 IterDomain is equal to 1, since we must do so without using an ExpressionEvaluator in that case.

In order to increase the power of our simplification and equivalence checking ability, this PR makes the following key changes:

  1. In addition to .name() which numbers each val by its ValType, each Val now has a .number() which is independent of val type, and Fusion can look up any Val* from its number.
  2. The UnionFind class is introduced which is a very efficient class for tracking equivalence classes. https://en.wikipedia.org/wiki/Disjoint-set_data_structure. This is similar to DisjointSets and may some day be a valid replacement for it, but the focus of UnionFind is on quick merging and pairwise comparison, with slower extraction of entire equivalence classes.
  3. The ValEquivalence class is introduced which plumbs the .number() of each val into a UnionFind and implements some methods for simplification/selection of Vals. An object of this class can quickly attest or test equivalence of any two Vals even if they are of different types, such as Int and NamedScalar. This forms the basis of a term rewriting system. The Fusion holds a single ValEquivalence object representing the weakest form of equivalence (TODO).
  4. The ExprEquivalence class is introduced which is based on ValEquivalence. It allows us to propose "sketches" of new Exprs and check whether an Expr whose 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 created Vals and Exprs when exploring rewrites. This could be useful for simplifyExpr, for passes like reduce-broadcast fusions in lowering, or for other graph optimization passes such as in Cast opt pass #355.
  5. Adds the ability to test or attest equivalences of Vals at any point. NOT YET STARTED
  6. Organizes integer constants such that int constant expressions have a single representative, so that individual pairwise attestations propagate properly. This makes it trivial to check int equality proofs at any point in the program (TODO). Note that inequalities are more difficult to handle, but that quick equality checking might aid in proving inequalities in some cases.

All of these are currently an early work in progress.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant