[Unity] Allow eliminating only call nodes in CSE pass

[masahi]

The CSE pass eliminates expressions very eagerly. This can result in an undesirable result, for example relax.arange(...) op which requires all of its inputs to be PrimValue complains when its inputs are CSE variables.

https://github.com/apache/tvm/blob/unity/src/relax/op/tensor/create.cc#L248-L253

  File "/Users/masa/projects/dev/tvm/src/relax/ir/block_builder.cc", line 138
TVMError: Arange expects the `start` to be a PrimValue, but got relax.expr.DataflowVar

That error can be easily fixed, but I've also hit similar issues with other ops like strided_slice as well. Since I only need common CallNodes to be eliminated, I'm adding an option to avoid the default fine-grained elimination.

cc @slyubomirsky @sunggg

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[slyubomirsky]

The proposed changes are implemented correctly and I don't object to them per se. But I wonder if this is the right solution. Perhaps we should exclude PrimValues from CSE in general instead? There isn't really a good reason to CSE them in the first place. Maybe there are other types of expressions we shouldn't CSE either. (One other candidate I could think of: scalar constants. Shape expressions might be another.)

[slyubomirsky]

I've decided to approve because, while we might consider a setting to leave only scalar constants, shape constants, and PrimValues (we could call them atomic values), calls-only is indeed a reasonable simplification that will account for almost all expensive computations. (Other situations besides calls, though, where CSE is useful include large tensor constants, repeated inner functions, and large tuples.)