[TIR] CSE-TIR Pass - More deterministic behavior

[AndrewZhaoLuo]

Running the CSE-TIR pass led to some variation in lowering to TIR from the same schedule multiple times. I tracked this down to iterating over an unordered map. I think having consistency in lowering to TIR is a desirable property to have so I made the offending code a little better in this regard.

Example test script:
https://github.com/AndrewZhaoLuo/TVM-Sandbox/blob/19284ddbd6bb28af61c0c2aa8bb334c5c53731a7/tir/test_inconsistent_tir_lowering.py#L1

^--- Running the above off main produces different hashes every time, running after this diff it produces the same hash.

There may be more examples of non-determinism in CSE-TIR but this is what I found for now 👀

[AndrewZhaoLuo]

cc some folks from #9482: @masahi @FranckQC @mbs-octoml @jroesch

[masahi]

@AndrewZhaoLuo maybe we can also solve this by changing the comparison functor in

tvm/src/tir/transforms/common_subexpr_elim_tools.h

Line 54 in 09f7be2

using ComputationTable = std::unordered_map<PrimExpr, size_t, StructuralHash, ExprDeepEqual>;

[AndrewZhaoLuo]

Hmm, so the unordered map doesn't have an ordering so all it needs is some way of determining equality of elements.

An alternative would be std::map which does have a notion of ordering (and you can define a comparator). However, it's performance characteristics are different (usually self-balancing binary tree so O(log n) lookup instead of O(~1)).

In general though yeah, we need to be on the lookout for iterating over unordered collections cause it will be non-determinstic so might be good for maintainability to use std::map

I'll try refactoring unordered_map --> map sometime later, though need to read more about performance (it's already kind of slow when processing a statement with 100 common subexpressions).

[masahi]

Hmm, so the unordered map doesn't have an ordering

The actual order doesn't matter, does it? unordered_map doesn't store elements in a sorted order, but it does have some ordering internally determined by the comparison functor. And what we need is just a fixed order, not necessarily the sorted order, so I think unordered map should do the job.

[masahi]

ah sorry, I was confused with something, "but it does have some ordering internally determined by the comparison functor" is not correct.

[masahi]

@AndrewZhaoLuo I think we can merge this as is for now. The CI failed on the mac os, probably need to kick another job.

[FranckQC]

Hi,

Sorry I couldn't discuss this with you last Thursday/Friday. Thanks for improving the CSE!

Yes, indeed, iterating through the hash table will not necessary lead to accessing its elements always in the same order. That's because the iteration order depends on the hashes (and we use structural_hash for that), which themselves are based on pointers, so different run will give different addresses, which will give different hashes, and thus a different order in which elements are accessed.

I guess it is not an issue as far as the CSE is concerned, because after producing this hashtable, we use it to produce a vector of semantic entities (which are pairs <expr, count>).
The main loop of the CSE will later iterate through the elements of this vector.
But before that, this vector gets sorted by the size (called "complexity") of the expression (the first component).

So the order of the end result (the vector) is non-deterministic only for elements with the same size / complexity, for which it does not matter (again, only for the CSE) in which order we common them. (i.e. it should not add or remove opportunities for commoning more later). Said differently, this non-determinism should appear for orthogonal choices.
I guess as in:

Mem[i1] = (a+b)+ (c+d);
Mem[i2] = a+b;               // We can common this one before or after
Mem[i3] = c+d;               // this one, it does not matter for the CSE
Mem[i4] = (a+b) + (c+d) + e;

for which we could produce:

let cse_var_2 = a+b in
let cse_var_3 = c+d in
let cse_var_1 = cse_var_2 + cse_var_3 in
Mem[i1] = cse_var_1;
Mem[i2] = cse_var_2;
Mem[i3] = cse_var_3;
Mem[i4] = cse_var_1 + e;

Or the alternative version were (c+d) gets introduced before (a+b). In practice, what is chosen currently depends on the addresses of the free variables 'a', 'b', 'c' and 'd' because they are used for the hashes.

I did not think that having this level of non-determinism would be an issue, but you're clearly right that its better to always produce exactly the same resulting TIR code, that just makes testing easier. Sorry I did not think about it when I wrote it!
Your way of making it deterministic (by first copying the hashtable into an array, which you then sort using the actual syntax that the expressions represent) is perfectly fine to me.

According to the doc about structural_hash ( here : https://tvm.apache.org/docs/reference/api/python/ir.html ), it seems an alternative could have been to set its map_free_vars parameter to true, as that's the reason the hashes use pointers. If that was the only place were addresses are used in the hashes, it should make everything deterministic too. I'll see if that works too, out of curiosity.

Many thanks for having spotted that and for having fixed it!

Kind regards,

Franck