[TIR][Utility] More flexible tir::Substitute arguments

[Lunderberg]

Previously, the tir::Substitute method had overloads that supported a few ways of providing the variable map (e.g. const Map<Var,PrimExpr>&, std::unordered_map<const VarNode*, PrimExpr>&, etc.), delegating out to the overload that uses std::function<Optional<PrimExpr>(const Var&)>. However, the types supported for the variable map depended on the type being substituted (e.g. only supporting const Map<Var,PrimExpr>& with substituting into a Array<Range>), which would be unexpected to new developers.

This PR makes the tir::Substitute utility more uniform in the arguments that it accepts.

• For any type that is supported by tir::Substitute, Array<T> is also supported.

• Any variable mapping type can be used with any substitution type. All variable mapping types are normalized to std::function<Optional<PrimExpr>(const Var&)>.

• For Map and std::unordered_map arguments, the value type may be any subclass of PrimExpr (e.g. Map<Var, Var> instead of Map<Var, PrimExpr>). Previously, the calling scope needed to either construct a temporary map that returned PrimExpr, or to use a broader value type in the map than otherwise required.

The initial and primary goal was to allow a Map<Var, Var> to be used as an argument to tir::Substitute, rather than a Map<Var, PrimExpr>, and making the utility more general was more straightforward than adding multiple overloads specificall for Map<Var, Var>.

[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 @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[kparzysz-quic]

LGTM, pending some comment fixes.

[Lunderberg]

Thank you, and corrected the copy/paste errors in the docstring.