ValReplacementMutator fails to mutate an expression even when its output is mutated

[naoyam]

pytorch/third_party/nvfuser/csrc/ir_utils.cpp

Line 435 in 8ed9540

class ValReplacementMutator : private OptOutMutator {

I don't have a failing case with devel, but encountered a problem with #2480.

ValReplacementMutator traverses all fusion vals and exprs from inputs to outputs. For an expression, mutate is called for its inputs, and then mutate is called for the expression itself, and then finally it's called for the outputs. This means that when mutating an expression, its outputs are not yet mutated. However, in OptOutMutator::mutate(Expr*), it seems it assumes the outputs are already mutated:

std::vector<Val*> mutated_outputs;
  mutated_outputs.reserve(op->outputs().size());
  for (auto output : op->outputs()) {
    mutated_outputs.emplace_back(maybeMutated(output));
  }

https://github.com/csarofeen/pytorch/blob/devel/third_party/nvfuser/csrc/mutator.cpp#L131

Since they are not yet mutated, maybeMutated should always just return the original (unmutated) output, even when they may eventually be mutated when mutate is called for them. This could obviously result in inconsistent states. For example, if an output is indeed mutated, its defining expression may not be created since when mutating the expression it may determine nothing is mutated.

I encountered some strange error due to this when ValReplacementMutator was used by replaceSymbolicSizes.

I think ValReplacementMutator should first visit all vals and mutate them as needed, and then visit all exprs.

[zasdfgbnm]

I think ValReplacementMutator should first visit all vals and mutate them as needed, and then visit all exprs.

This makes sense to me.