TensorView::computeAt() recursive traversal is non-deterministic

[tlemo]

TensorView::computeAt() is propagating to all the uses of the same producer:

pytorch/torch/csrc/jit/codegen/cuda/tensor_view.cpp

Line 176 in 172bcdb

for (Expr* other_use : FusionGuard::getCurFusion()->uses(this)) {

The problem is that Fusion::uses() returns a std::set<Expr*>:

pytorch/torch/csrc/jit/codegen/cuda/fusion.cpp

Line 289 in 172bcdb

std::set<Expr*> Fusion::uses(Val* val) const {

std::set normally offers a predictable (sorted) order of values, but in this case the keys are Expr*, which are non-deterministic. This is an issue since the order of computeAt() traversal changes the end results.

This issue manifested as a failure of a unit test triggered by an unrelated change in a different unit test (but which impacted the heap layout as a side effect).

[tlemo]

Here's an example where the current

  Fusion fusion;
  FusionGuard fg(&fusion);

  TensorView* tv0 = makeDummyTensor(2);
  fusion.addInput(tv0);

  TensorView* tv1 = static_cast<TensorView*>(mul(tv0, new Float(-1.0)));
  TensorView* tv2 = static_cast<TensorView*>(add(tv0, new Float(3.0)));
  TensorView* tv3 = static_cast<TensorView*>(mul(tv0, new Float(2.0)));
  TensorView* tv4 = static_cast<TensorView*>(add(tv2, tv1));

  TensorView* tv5 = static_cast<TensorView*>(add(tv4, tv3));
  TensorView* tv6 = static_cast<TensorView*>(add(tv0, tv3));

  fusion.addOutput(tv5);
  fusion.addOutput(tv6);

  tv6->merge(0);
  tv6->split(0, 128);
  tv6->split(0, 4);

  tv6->axis(0)->parallelize(ParallelType::BIDx);

  tv5->reorder({{-1, 0}});

  tv0->computeAt(tv3, 1); // <---------------------------- the order of traversal matters here

---

One potential order - tv5 will end up with rank=3
   0->computeAt(3, 1)
     3->computeAt(6, 1)
       5->computeAt(6, 1)
     2->computeAt(3, 1)
       4->computeAt(3, 1)
         3->computeAt(5, 1)
           6->computeAt(5, 1)
     1->computeAt(3, 1)
       4->computeAt(3, 1)
         3->computeAt(5, 1)
           6->computeAt(5, 1)
   0->computeAt(6, 1)
     3->computeAt(6, 1)
       5->computeAt(6, 1)
     2->computeAt(6, 1)
       4->computeAt(6, 1)
         5->computeAt(6, 1)
     1->computeAt(6, 1)
       4->computeAt(6, 1)
         5->computeAt(6, 1)

Second order, tv5 will have rank=2
   0->computeAt(3, 1)
     2->computeAt(3, 1)
       4->computeAt(3, 1)
         3->computeAt(5, 1)
           6->computeAt(5, 1)
     1->computeAt(3, 1)
       4->computeAt(3, 1)
         3->computeAt(5, 1)
           6->computeAt(5, 1)
     3->computeAt(6, 1)
       5->computeAt(6, 1)
   0->computeAt(6, 1)
     3->computeAt(6, 1)
       5->computeAt(6, 1)
     2->computeAt(6, 1)
       4->computeAt(6, 1)
         5->computeAt(6, 1)
     1->computeAt(6, 1)
       4->computeAt(6, 1)
         5->computeAt(6, 1)

[csarofeen]

This has become deterministic a while back, issues is from desynced work. @tlemo do we want to close this or do you want me to pull up the reference to the fix and put it here first?

[tlemo]

I think it's safe to close it.