Definition-time `simplifyExpr` with concretization axioms

[jacobhinkle]

With dynamic resize, we have implemented a very basic system of trying to prove propositions at definition time (before concretization) if we have some constant inputs; for example, if we have constant pad values that sum to 2, we know that the padded extent will be >= 2 even if we don't know the input extent. This allows us to avoid using symbolic operations in some common cases. Other ops like cat will pad using complicated expressions, so the first point of this issue is that it might be advantageous to have more sophisticated theorem proving during definition.

Currently NVFuser possesses a powerful system for simplifying expressions: simplifyExpr. It implements iterative term rewriting, starting with a provided set of assumptions that defaults to a collection of "axioms" from IrContainer::axioms() which are a collection of Bools that are known to be True: namely this attests that parallelization dimensions are positive and thread and block IDs are non-negative.

There are more axioms that could empower this system. For example we know that any fusion input extents should be non-negative. During concretization, we analyze individual ops and look at scalars, extract their values and make a determination. These decision points are then predicates that can be considered constant for that concretization. For example, if we determine that the extent of a pad op is equal to 1, that will be true for any set of inputs sharing that concretization. The second point of this issue is in the other direction: that axioms accumulated during definition and concretization might be useful during lowering.

Approaches

I think one way to enable this flexibly is to add the ability to append axioms to IrContainer by adding an attestAxiom method. Creation of a TensorView would add an axiom that each extent is non-negative. We can then use simplifyExpr directly during definition or lowering. We could then also sprinkle in more axioms to aid the system. Alternatively, we could just override axioms() for Fusion in order to also return the input extent nonnegative axioms, but this is less flexible.

[jacobhinkle]

cc @zasdfgbnm

[zasdfgbnm]

Love this idea.

[zasdfgbnm]

I think I will need the assumptions that all extents are non-negative in order to simplify vectorization factor. Self assign.

[zasdfgbnm]

Instead of adding an axiom every time we create a tensor, I think it makes more sense to use a pre-concretization pass to fill in this information. And we should have another concretization pass to replace equal extent the same Val*, for example, if we have T1[I1{i1}] + T2[I2{i2}], then we should replace it as T1[I1{i1}] + T2[I2{i1}], this would help reduce the number of assumptions we have. I think I can reuse @jjsjann123's infrastructure for cast simplification to build this out.

[zasdfgbnm]

@jacobhinkle Am I understanding correct that expr simplifier might also potentially be useful for analysis in concretization? If this is the case, I think we should have a pre-conretization pass to replace all equal extents with the same Val*, then another pre-concretization pass to insert axioms like extent >= 0, and during concretization, we should find all zero extents and convert their extents to compile time 0, remove corresponding axioms, and convert all other axioms as extent > 0. If not, we should have a pre-conretization pass to replace all equal extents, and a post-concretization pass to insert extent > 0.

[jacobhinkle]

Yeah that's possible to do. The new pass seems very closely related to #382. In terms of new axioms, we could indeed do these after definition instead of during. I am thinking it may be useful to do a "partial concretization" pass, since one of the reasons I was thinking of this is for reshape: if we can prove that certain extents line up we can do the decomposition into Keep/Squeeze/Merge even before seeing inputs in some cases.

[zasdfgbnm]

Are you talking about detecting things like view(tv, {i1, i2}, {i1*i2}), if we can analyze it symbolically, then we don't need to know the actual value of i1 and i2 in order to know this is a merge. Generally, if you don't need the distributivity, you don't care about the sign of these values, see:

Fuser/csrc/expr_simplifier.h

Lines 444 to 460 in 8fbb798

	// Theorem 6.9: If a is a multiple of b, then a % b = 0
	// Proof: This can be proved directly from Euclid's division lemma
	//
	// Theorem 6.10: If b is a multiple of c, then we have: a*(b/c) = (a*b)/c
	// Proof: Same proof as 4.10
	//
	// Theorem 6.11: a/(b*c) = (a/b)/c
	// Proof: This is part of Theorem 6.1
	//
	// Theorem 6.12: a % (b * c) = a % b + ((a / b) % c) * b
	// Proof: Already proved in the proof of Theorem 6.1
	//
	// Theorem 6.13: If d divides a and b, then a % b = ((a / d) % (b / d)) * d
	// Proof: Same proof as 4.13
	//
	// Theorem 6.14: If b is a multiple of c, then a/(b/c) = (a*c)/b
	// Proof: Same proof as 4.14

So I don't think you need any assumption about the extents are non-negative. So probably in this case, we can just fill extent > 0 after concretization.

[jacobhinkle]

Sorry to not be so clear here: in reshape a value of -1 means to fill in with whatever value is needed to preserve the number of elements. It can appear in at most one position but could be any of them, so step one is convert all inputs to where statements checking for that condition. That may challenge the proofs

[zasdfgbnm]

Could you post an example here? I expect -1 should be replaced as things like size1 * size2 * size3 / size4 / size5? Where is the "where statement"?