You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "sunggg (via GitHub)" <gi...@apache.org> on 2023/03/31 22:05:46 UTC
[GitHub] [tvm] sunggg opened a new pull request, #14447: [Unity][Op] introduce `shape_to_tensor` op
sunggg opened a new pull request, #14447:
URL: https://github.com/apache/tvm/pull/14447
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)
```
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 enable `PackedFunc` evaluation, we should be able to remove these unnecessary special handling.
cc. @jwfromm @yongwww @psrivas2 @slyubomirsky
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 merged pull request #14447: [Unity][Op] introduce `shape_to_tensor` op
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 merged PR #14447:
URL: https://github.com/apache/tvm/pull/14447
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tvm-bot commented on pull request #14447: [Unity][Op] introduce `shape_to_tensor` op
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #14447:
URL: https://github.com/apache/tvm/pull/14447#issuecomment-1492665421
<!---bot-comment-->
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](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
<!--bot-comment-ccs-start-->
* cc @quic-sanirudh <sub>See [#10317](https://github.com/apache/tvm/issues/10317) for details</sub><!--bot-comment-ccs-end-->
<sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] sunggg commented on a diff in pull request #14447: [Unity][Op] introduce `shape_to_tensor` op
Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #14447:
URL: https://github.com/apache/tvm/pull/14447#discussion_r1154950195
##########
src/relax/op/op_common.cc:
##########
@@ -38,7 +38,7 @@ Array<TensorStructInfo> GetInputTensorStructInfo(const Call& call, const BlockBu
if (sinfo == nullptr) {
ctx->ReportFatal(Diagnostic::Error(call)
<< op << " requires the input " << op->arguments[i]->name
- << " to be Tensor. However, the given one is "
+ << " to be Tensor. However, the given one has a "
Review Comment:
This is unrelated to this PR, but found the current message confusing during the debugging.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14447: [Unity][Op] introduce `shape_to_tensor` op
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14447:
URL: https://github.com/apache/tvm/pull/14447#discussion_r1154967695
##########
python/tvm/relax/op/base.py:
##########
@@ -253,6 +254,19 @@ def render_object(val: tvm.Object) -> str:
return str(val)
+@tvm.register_func("relax.run.shape_to_tensor")
+def relax_shape_to_tensor(shape_tuple: tvm.runtime.ShapeTuple) -> tvm.nd.NDArray:
+ """
+ Takes a ShapeTuple and convert it to NDArray.
+
+ Parameters
+ ----------
+ shape_tuple: tvm.runtime.ShapeTuple
+ Shape tuple that we want to convert to NDArray at runtime
+ """
+ return tvm.nd.array([int(v) for v in shape_tuple])
Review Comment:
do we assume it's always on cpu?
##########
src/relax/transform/fold_constant.cc:
##########
@@ -279,6 +279,25 @@ class ConstantFolder : public ExprMutator {
}
return ShapeExpr(shape_values);
}
+ } else if (op->name == "relax.shape_to_tensor") {
+ // Special handling for "relax.shape_to_tensor" since it is implemented in PackedFunc.
+ // TODO(sunggg): revisit this when we extend ConstantFolding to fold PackedFunc.
+ Expr arg = post_call->args[0];
+ ShapeExpr shape = Downcast<ShapeExpr>(arg);
+ Array<PrimExpr> values = shape->values;
+ Array<Integer> arr;
+ bool isKnown = true;
Review Comment:
nit
```suggestion
bool is_known = true;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] sunggg commented on a diff in pull request #14447: [Unity][Op] introduce `shape_to_tensor` op
Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #14447:
URL: https://github.com/apache/tvm/pull/14447#discussion_r1154986001
##########
python/tvm/relax/op/base.py:
##########
@@ -253,6 +254,19 @@ def render_object(val: tvm.Object) -> str:
return str(val)
+@tvm.register_func("relax.run.shape_to_tensor")
+def relax_shape_to_tensor(shape_tuple: tvm.runtime.ShapeTuple) -> tvm.nd.NDArray:
+ """
+ Takes a ShapeTuple and convert it to NDArray.
+
+ Parameters
+ ----------
+ shape_tuple: tvm.runtime.ShapeTuple
+ Shape tuple that we want to convert to NDArray at runtime
+ """
+ return tvm.nd.array([int(v) for v in shape_tuple])
Review Comment:
Yes. shape tuple and shape computation happen on the CPU side.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org