[Unity][Op] introduce shape_to_tensor op

[sunggg]

In Unity, we have a clear distinction between tensor and shape: we have ShapeExpr and ShapeStructInfo in AST, ShapeTuple in runtime container. Meanwhile, most of operators and their TOPI implementations are defined with tensors. For example, relax.take is defined as follows:

// Operator definition
TVM_REGISTER_OP("relax.take")
    .set_attrs_type<TakeAttrs>()
    .set_num_inputs(2)
    .add_argument("x", "Tensor", "The source tensor.")
    .add_argument("indices", "Tensor", "The indices of the values to extract.")
    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoTake);
       
StructInfo InferStructInfoTake(const Call& call, const BlockBuilder& ctx) {
  Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
   // Assume inputs are `Tensor` 
  TensorStructInfo data_sinfo = input_sinfo[0];
  TensorStructInfo indices_sinfo = input_sinfo[1];
  ...
}

// TOPI
inline Tensor take(const Tensor& a, const Tensor& indices, int batch_dims, int axis, ...)

To allow the shape computation, this PR introduces a shape_to_tensor op that converts ShapeTuple to NDArray at runtime. This enables the common shape computation patterns like following:

shape_var: R.Shape(ndim=4)
lv: R.Tensor((4,), dtype="int64") = R.shape_to_tensor(shape_var)
lv1: R.Tensor((1,), dtype="int64") = R.take(lv, indices, axis=0)
lv2: R.Tensor((1, 1), dtype="int64") = R.expand_dims(lv1, axis=[0])
gv: R.Tensor((1, 1), dtype="int64") = R.concat((lv2,), axis=0)

It's worth noting that tensor_to_shape op is already introduced in #14282, so roundtrip between shape and tensor would be now possible.

Currently, this op requires special handling in FoldConstant pass since this pass is only able to evaluate TIR primfunc, not PackedFunc. Once we extend FoldConstant to support PackedFunc evaluation, we should be able to remove these unnecessary special handling.

cc. @jwfromm @yongwww @psrivas2 @slyubomirsky @tqchen

[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}

[sunggg]

This is unrelated to this PR, but found the current message confusing during the debugging.

[vinx13]

do we assume it's always on cpu?

[sunggg]

Yes. shape tuple and shape computation happen on the CPU side.

[yongwww]

LGTM!