This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 8241385f59 [Unity] De-duplicate calls to TensorStructInfo constructor
(#16209)
8241385f59 is described below
commit 8241385f595fa17ba18c0d4b05d4704a21643f6c
Author: Eric Lunderberg <[email protected]>
AuthorDate: Mon Dec 11 17:42:33 2023 -0600
[Unity] De-duplicate calls to TensorStructInfo constructor (#16209)
* [Unity] Change TensorStructInfo argument to Optional<VDevice>
Prior to this commit, the `TensorStructInfo` constructor took as input
a `VDevice`, with a default value of `VDevice()`, and then assigned it
to a member of type `Optional<VDevice>`. This commit changes the
constructor signature to match the member type, which will remove
unnecessary type conversions when copying a `TensorStructInfo`.
* Remove duplication by passing Optional<Device>
---
include/tvm/relax/struct_info.h | 4 +-
src/relax/ir/struct_info.cc | 5 +-
src/relax/op/ccl/ccl.cc | 11 +-
src/relax/op/distributed/distributed.cc | 5 +-
src/relax/op/image/resize.cc | 10 +-
src/relax/op/nn/attention.cc | 5 +-
src/relax/op/nn/convolution.cc | 40 +----
src/relax/op/nn/nn.cc | 50 ++----
src/relax/op/nn/pooling.cc | 20 +--
src/relax/op/tensor/binary.cc | 20 +--
src/relax/op/tensor/create.cc | 5 +-
src/relax/op/tensor/index.cc | 49 ++----
src/relax/op/tensor/manipulate.cc | 290 ++++++++------------------------
src/relax/op/tensor/search.cc | 22 +--
src/relax/op/tensor/set.cc | 39 +----
src/relax/op/tensor/statistical.cc | 45 ++---
16 files changed, 141 insertions(+), 479 deletions(-)
diff --git a/include/tvm/relax/struct_info.h b/include/tvm/relax/struct_info.h
index 2e224f1830..7c13fcc531 100644
--- a/include/tvm/relax/struct_info.h
+++ b/include/tvm/relax/struct_info.h
@@ -231,7 +231,7 @@ class TensorStructInfo : public StructInfo {
*
* \note shape must already be normalized.
*/
- TVM_DLL TensorStructInfo(Expr shape, DataType dtype, VDevice vdevice =
VDevice(),
+ TVM_DLL TensorStructInfo(Expr shape, DataType dtype, Optional<VDevice>
vdevice = NullOpt,
Span span = Span());
/*!
@@ -241,7 +241,7 @@ class TensorStructInfo : public StructInfo {
* \param vdevice The virtual device.
* \param span The span of the AST.
*/
- TVM_DLL TensorStructInfo(DataType dtype, int ndim, VDevice vdevice =
VDevice(),
+ TVM_DLL TensorStructInfo(DataType dtype, int ndim, Optional<VDevice> vdevice
= NullOpt,
Span span = Span());
TVM_DEFINE_OBJECT_REF_METHODS(TensorStructInfo, StructInfo,
TensorStructInfoNode);
diff --git a/src/relax/ir/struct_info.cc b/src/relax/ir/struct_info.cc
index 9b635bb479..302534b414 100644
--- a/src/relax/ir/struct_info.cc
+++ b/src/relax/ir/struct_info.cc
@@ -105,7 +105,8 @@ TVM_REGISTER_GLOBAL("relax.ShapeStructInfo")
});
// Tensor
-TensorStructInfo::TensorStructInfo(Expr shape, DataType dtype, VDevice
vdevice, Span span) {
+TensorStructInfo::TensorStructInfo(Expr shape, DataType dtype,
Optional<VDevice> vdevice,
+ Span span) {
ObjectPtr<TensorStructInfoNode> n = make_object<TensorStructInfoNode>();
// assign ndim before move
Optional<ShapeStructInfo> sinfo = MatchStructInfo<ShapeStructInfo>(shape);
@@ -122,7 +123,7 @@ TensorStructInfo::TensorStructInfo(Expr shape, DataType
dtype, VDevice vdevice,
data_ = std::move(n);
}
-TensorStructInfo::TensorStructInfo(DataType dtype, int ndim, VDevice vdevice,
Span span) {
+TensorStructInfo::TensorStructInfo(DataType dtype, int ndim, Optional<VDevice>
vdevice, Span span) {
ObjectPtr<TensorStructInfoNode> n = make_object<TensorStructInfoNode>();
CHECK_GE(ndim, -1) << "ndim of TensorStructInfo must be >= -1, but got " <<
ndim;
n->ndim = ndim;
diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc
index 22ab22e940..c0fe6f4d88 100644
--- a/src/relax/op/ccl/ccl.cc
+++ b/src/relax/op/ccl/ccl.cc
@@ -72,11 +72,7 @@ StructInfo InferStructInfoAllGather(const Call& call, const
BlockBuilder& ctx) {
}
Array<PrimExpr> output_shape = input_shape.value();
output_shape.Set(0, floor(output_shape[0] * num_workers.value()));
- VDevice vdevice;
- if (input_sinfo->vdevice.defined()) {
- vdevice = input_sinfo->vdevice.value();
- }
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype, vdevice);
+ return TensorStructInfo(ShapeExpr(output_shape), output_dtype,
input_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.ccl.allgather")
@@ -141,10 +137,7 @@ StructInfo InferStructInfoScatter(const Call& call, const
BlockBuilder& ctx) {
Array<PrimExpr> output_shape = input_shape.value();
output_shape.Set(attrs->axis, div(output_shape[attrs->axis], num_workers));
- if (input_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype,
input_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), output_dtype,
input_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.ccl.scatter_from_worker0")
diff --git a/src/relax/op/distributed/distributed.cc
b/src/relax/op/distributed/distributed.cc
index 1159e0e830..67e11f1535 100644
--- a/src/relax/op/distributed/distributed.cc
+++ b/src/relax/op/distributed/distributed.cc
@@ -165,10 +165,7 @@ StructInfo InferStructInfoRtoS(const Call& call, const
BlockBuilder& ctx) {
Array<PrimExpr> output_shape = input_shape.value();
output_shape.Set(attrs->axis, div(output_shape[attrs->axis], num_workers));
- if (input_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype,
input_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), output_dtype,
input_sinfo->vdevice);
}
StructInfo InferDistStructInfoRtoS(const Call& call, const BlockBuilder& ctx) {
diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc
index 3a4cb26861..8b92f34edd 100644
--- a/src/relax/op/image/resize.cc
+++ b/src/relax/op/image/resize.cc
@@ -90,10 +90,7 @@ StructInfo InferStructInfoResize2D(const Call& call, const
BlockBuilder& ctx) {
Optional<ShapeExpr> data_shape =
CheckNdimPerLayoutAndGetShape(call, ctx,
GetRef<TensorStructInfo>(data_sinfo), data_layout);
if (!data_shape.defined() || size_value == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(out_dtype, data_layout.ndim(),
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(out_dtype, data_layout.ndim());
+ return TensorStructInfo(out_dtype, data_layout.ndim(),
data_sinfo->vdevice);
}
Array<PrimExpr> data_NCHW_shape =
data2NCHW.ForwardShape(data_shape.value()->values);
@@ -102,10 +99,7 @@ StructInfo InferStructInfoResize2D(const Call& call, const
BlockBuilder& ctx) {
out_NCHW_shape.Set(3, size_value->values[1]);
Array<PrimExpr> out_shape = data2NCHW.BackwardShape(out_NCHW_shape);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutResize2d(const Call& call,
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index c6aed941b6..ca3746ddad 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -133,10 +133,7 @@ StructInfo InferStructInfoAttention(const Call& call,
const BlockBuilder& ctx) {
}
Array<PrimExpr> output_shape = {num_batches, num_queries, num_heads,
head_dim_value};
- if (q_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), q_sinfo->dtype,
q_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), q_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), q_sinfo->dtype,
q_sinfo->vdevice);
}
Call InferMixedPrecisionAttention(const Call& call, const DataType& out_dtype)
{
diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc
index e8cb1916e8..cea234060e 100644
--- a/src/relax/op/nn/convolution.cc
+++ b/src/relax/op/nn/convolution.cc
@@ -79,10 +79,7 @@ StructInfo InferStructInfoConv1d(const Call& call, const
BlockBuilder& ctx) {
: attrs->out_dtype;
Optional<VDevice> vdevice = InferBinaryArithOpOutVDevice(call, ctx,
data_sinfo, weight_sinfo);
if (!data_shape.defined() || !weight_shape.defined()) {
- if (vdevice.defined()) {
- return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice.value());
- }
- return TensorStructInfo(out_dtype, out_layout.ndim());
+ return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice);
}
Array<PrimExpr> data_NCW_shape =
data2NCW.ForwardShape(data_shape.value()->values);
@@ -125,10 +122,7 @@ StructInfo InferStructInfoConv1d(const Call& call, const
BlockBuilder& ctx) {
out_NCW_shape[2] = analyzer->Simplify(floordiv(numerator_w,
attrs->strides[0]) + 1);
Array<PrimExpr> out_shape = out2NCW.BackwardShape(out_NCW_shape);
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice);
}
InferLayoutOutput InferLayoutConv1d(const Call& call,
@@ -248,10 +242,7 @@ StructInfo InferStructInfoConv2d(const Call& call, const
BlockBuilder& ctx) {
: attrs->out_dtype;
Optional<VDevice> vdevice = InferBinaryArithOpOutVDevice(call, ctx,
data_sinfo, weight_sinfo);
if (!data_shape.defined() || !weight_shape.defined()) {
- if (vdevice.defined()) {
- return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice.value());
- }
- return TensorStructInfo(out_dtype, out_layout.ndim());
+ return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice);
}
Array<PrimExpr> data_NCHW_shape =
data2NCHW.ForwardShape(data_shape.value()->values);
@@ -299,10 +290,7 @@ StructInfo InferStructInfoConv2d(const Call& call, const
BlockBuilder& ctx) {
out_NCHW_shape[3] = analyzer->Simplify(floordiv(numerator_w,
attrs->strides[1]) + 1);
Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice);
}
InferLayoutOutput InferLayoutConv2d(const Call& call,
@@ -427,10 +415,7 @@ StructInfo InferStructInfoConv1dTranspose(const Call&
call, const BlockBuilder&
: attrs->out_dtype;
Optional<VDevice> vdevice = InferBinaryArithOpOutVDevice(call, ctx,
data_sinfo, weight_sinfo);
if (!data_shape.defined() || !weight_shape.defined()) {
- if (vdevice.defined()) {
- return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice.value());
- }
- return TensorStructInfo(out_dtype, out_layout.ndim());
+ return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice);
}
Array<PrimExpr> data_NCW_shape =
data2NCW.ForwardShape(data_shape.value()->values);
@@ -483,10 +468,7 @@ StructInfo InferStructInfoConv1dTranspose(const Call&
call, const BlockBuilder&
out_NCW_shape[2] = analyzer->Simplify(out_w);
Array<PrimExpr> out_shape = out2NCW.BackwardShape(out_NCW_shape);
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice);
}
// TODO(relax-team): implement FInferMixedPrecision and FRelaxInferLayout for
conv1d_transpose
@@ -571,10 +553,7 @@ StructInfo InferStructInfoConv2dTranspose(const Call&
call, const BlockBuilder&
: attrs->out_dtype;
Optional<VDevice> vdevice = InferBinaryArithOpOutVDevice(call, ctx,
data_sinfo, weight_sinfo);
if (!data_shape.defined() || !weight_shape.defined()) {
- if (vdevice.defined()) {
- return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice.value());
- }
- return TensorStructInfo(out_dtype, out_layout.ndim());
+ return TensorStructInfo(out_dtype, out_layout.ndim(), vdevice);
}
Array<PrimExpr> data_NCHW_shape =
data2NCHW.ForwardShape(data_shape.value()->values);
@@ -635,10 +614,7 @@ StructInfo InferStructInfoConv2dTranspose(const Call&
call, const BlockBuilder&
out_NCHW_shape[3] = analyzer->Simplify(out_w);
Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice);
}
// TODO(relax-team): implement FInferMixedPrecision and FRelaxInferLayout for
conv2d_transpose
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
index f95cc9f4d6..6f8a90e3cb 100644
--- a/src/relax/op/nn/nn.cc
+++ b/src/relax/op/nn/nn.cc
@@ -266,15 +266,10 @@ StructInfo InferStructInfoBatchNorm(const Call& call,
const BlockBuilder& ctx) {
DataType dtype = input_sinfo[0]->dtype;
if (unknown_shape) {
- if (input_sinfo[0]->vdevice.defined()) {
- VDevice vdev = input_sinfo[0]->vdevice.value();
- return TupleStructInfo({TensorStructInfo(dtype, input_sinfo[0]->ndim,
vdev),
- TensorStructInfo(dtype, /*ndim=*/1, vdev),
- TensorStructInfo(dtype, /*ndim=*/1, vdev)});
- }
- return TupleStructInfo({TensorStructInfo(dtype, input_sinfo[0]->ndim),
- TensorStructInfo(dtype, /*ndim=*/1),
- TensorStructInfo(dtype, /*ndim=*/1)});
+ auto vdev = input_sinfo[0]->vdevice;
+ return TupleStructInfo({TensorStructInfo(dtype, input_sinfo[0]->ndim,
vdev),
+ TensorStructInfo(dtype, /*ndim=*/1, vdev),
+ TensorStructInfo(dtype, /*ndim=*/1, vdev)});
} else {
return TupleStructInfo({input_sinfo[0], input_sinfo[3], input_sinfo[4]});
}
@@ -337,12 +332,8 @@ StructInfo InferStructInfoLayerNorm(const Call& call,
const BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<LayerNormAttrs>();
bool unknown_shape = NormCheckDtypeAndShape(call, ctx, input_sinfo,
attrs->axes);
- if (input_sinfo[0]->vdevice.defined()) {
- return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype,
input_sinfo[0]->ndim,
- input_sinfo[0]->vdevice.value())
- : input_sinfo[0];
- }
- return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype,
input_sinfo[0]->ndim)
+ return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype,
input_sinfo[0]->ndim,
+ input_sinfo[0]->vdevice)
: input_sinfo[0];
}
@@ -514,12 +505,8 @@ StructInfo InferStructInfoRMSNorm(const Call& call, const
BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<RMSNormAttrs>();
bool unknown_shape = NormCheckDtypeAndShape(call, ctx, input_sinfo,
attrs->axes);
- if (input_sinfo[0]->vdevice.defined()) {
- return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype,
input_sinfo[0]->ndim,
- input_sinfo[0]->vdevice.value())
- : input_sinfo[0];
- }
- return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype,
input_sinfo[0]->ndim)
+ return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype,
input_sinfo[0]->ndim,
+ input_sinfo[0]->vdevice)
: input_sinfo[0];
}
@@ -629,10 +616,7 @@ StructInfo InferStructInfoCrossEntropy(const Call& call,
const BlockBuilder& ctx
}
}
}
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(Array<PrimExpr>()), dtype,
vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(Array<PrimExpr>()), dtype);
+ return TensorStructInfo(ShapeExpr(Array<PrimExpr>()), dtype, vdevice);
}
Expr cross_entropy_with_logits(Expr predictions, Expr labels) {
@@ -860,24 +844,14 @@ StructInfo InferStructInfoNLLLoss(const Call& call, const
BlockBuilder& ctx) {
if (reduction == "none") {
// () or (N,) or (N, d1, d2, ..., dk)
if (pred_sinfo->shape.as<ShapeExprNode>()) {
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype,
vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), output_dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), output_dtype, vdevice);
} else {
int output_ndim = pred_sinfo->ndim == kUnknownNDim ? kUnknownNDim :
pred_sinfo->ndim - 1;
- if (vdevice.defined()) {
- return TensorStructInfo(output_dtype, /*ndim=*/output_ndim,
vdevice.value());
- }
- return TensorStructInfo(output_dtype, /*ndim=*/output_ndim);
+ return TensorStructInfo(output_dtype, /*ndim=*/output_ndim, vdevice);
}
} else {
// sum or mean. output is scalar
- if (vdevice.defined()) {
- return TensorStructInfo(/*shape=*/ShapeExpr(Array<PrimExpr>()),
output_dtype,
- vdevice.value());
- }
- return TensorStructInfo(/*shape=*/ShapeExpr(Array<PrimExpr>()),
output_dtype);
+ return TensorStructInfo(/*shape=*/ShapeExpr(Array<PrimExpr>()),
output_dtype, vdevice);
}
}
diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc
index c26fae08c2..6c81f5310a 100644
--- a/src/relax/op/nn/pooling.cc
+++ b/src/relax/op/nn/pooling.cc
@@ -86,10 +86,7 @@ StructInfo InferStructInfoPool2D(const Call& call, const
BlockBuilder& ctx) {
Optional<ShapeExpr> data_shape =
CheckNdimPerLayoutAndGetShape(call, ctx, data_sinfo, data_layout);
if (!data_shape.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, out_layout.ndim(),
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, out_layout.ndim());
+ return TensorStructInfo(data_sinfo->dtype, out_layout.ndim(),
data_sinfo->vdevice);
}
Array<PrimExpr> data_NCHW_shape =
data2NCHW.ForwardShape(data_shape.value()->values);
@@ -117,10 +114,7 @@ StructInfo InferStructInfoPool2D(const Call& call, const
BlockBuilder& ctx) {
out_NCHW_shape[3] = analyzer->Simplify(floordiv(numerator_w,
attrs->strides[1]) + 1);
Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutPool2d(const Call& call,
@@ -210,10 +204,7 @@ StructInfo InferStructInfoAdaptiveAvgPool2D(const Call&
call, const BlockBuilder
!attrs->output_size.defined()) {
return data_sinfo;
} else {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, out_layout.ndim(),
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, out_layout.ndim());
+ return TensorStructInfo(data_sinfo->dtype, out_layout.ndim(),
data_sinfo->vdevice);
}
}
@@ -225,10 +216,7 @@ StructInfo InferStructInfoAdaptiveAvgPool2D(const Call&
call, const BlockBuilder
}
Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutAdaptiveAvgPool2D(const Call& call,
diff --git a/src/relax/op/tensor/binary.cc b/src/relax/op/tensor/binary.cc
index 87afc24397..f1427156e0 100644
--- a/src/relax/op/tensor/binary.cc
+++ b/src/relax/op/tensor/binary.cc
@@ -58,28 +58,16 @@ StructInfo InferStructInfoBroadcast(const Call& call, const
BlockBuilder& ctx,
Optional<Array<PrimExpr>> output_shape =
InferBinaryBroadcastShape(call, ctx, x1_shape->values,
x2_shape->values);
if (!output_shape.defined()) {
- if (vdevice.defined()) {
- return TensorStructInfo(output_dtype, /*ndim=*/output_ndim,
vdevice.value());
- }
- return TensorStructInfo(output_dtype, /*ndim=*/output_ndim);
+ return TensorStructInfo(output_dtype, /*ndim=*/output_ndim, vdevice);
} else {
ICHECK_EQ(static_cast<int>(output_shape.value().size()), output_ndim);
- if (vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape.value()), output_dtype,
vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape.value()), output_dtype);
+ return TensorStructInfo(ShapeExpr(output_shape.value()), output_dtype,
vdevice);
}
} else if (x1_sinfo->shape.defined() &&
x1_sinfo->shape.same_as(x2_sinfo->shape)) {
- if (vdevice.defined()) {
- return TensorStructInfo(x1_sinfo->shape.value(), output_dtype,
vdevice.value());
- }
- return TensorStructInfo(x1_sinfo->shape.value(), output_dtype);
+ return TensorStructInfo(x1_sinfo->shape.value(), output_dtype, vdevice);
} else {
- if (vdevice.defined()) {
- return TensorStructInfo(output_dtype, /*ndim=*/output_ndim,
vdevice.value());
- }
- return TensorStructInfo(output_dtype, /*ndim=*/output_ndim);
+ return TensorStructInfo(output_dtype, /*ndim=*/output_ndim, vdevice);
}
}
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index f5893d64b1..fd6fea6e70 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -77,10 +77,7 @@ StructInfo InferStructInfoFull(const Call& call, const
BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<InitAttrs>();
DataType out_dtype = attrs->dtype.is_void() ? fill_value_sinfo->dtype :
attrs->dtype;
- if (fill_value_sinfo->vdevice.defined()) {
- return TensorStructInfo(/*shape=*/call->args[0], out_dtype,
fill_value_sinfo->vdevice.value());
- }
- return TensorStructInfo(/*shape=*/call->args[0], out_dtype);
+ return TensorStructInfo(/*shape=*/call->args[0], out_dtype,
fill_value_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.full")
diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc
index 6d9dfc86ba..7ab98e9468 100644
--- a/src/relax/op/tensor/index.cc
+++ b/src/relax/op/tensor/index.cc
@@ -66,10 +66,7 @@ StructInfo InferStructInfoTake(const Call& call, const
BlockBuilder& ctx) {
<< data_sinfo->ndim);
}
if (data_sinfo->IsUnknownNdim() || indices_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
int axis = attrs->axis.defined()
@@ -78,11 +75,8 @@ StructInfo InferStructInfoTake(const Call& call, const
BlockBuilder& ctx) {
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
const auto* indices_shape = indices_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr || indices_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, indices_sinfo->ndim +
data_sinfo->ndim - 1,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, indices_sinfo->ndim +
data_sinfo->ndim - 1);
+ return TensorStructInfo(data_sinfo->dtype, indices_sinfo->ndim +
data_sinfo->ndim - 1,
+ data_sinfo->vdevice);
}
Array<PrimExpr> output_shape;
@@ -94,11 +88,7 @@ StructInfo InferStructInfoTake(const Call& call, const
BlockBuilder& ctx) {
output_shape.push_back(data_shape->values[i]);
}
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.take")
@@ -191,19 +181,13 @@ StructInfo InferStructInfoStridedSlice(const Call& call,
const BlockBuilder& ctx
}
if (data_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
std::vector<int> axes = NormalizeAxes(call, ctx, data_sinfo->ndim,
attrs->axes);
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
}
int n_axis = axes.size();
@@ -216,10 +200,7 @@ StructInfo InferStructInfoStridedSlice(const Call& call,
const BlockBuilder& ctx
for (int i = 0; i < n_axis; ++i) {
const auto* int_stride = strides[i].as<IntImmNode>();
if (!int_stride) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
}
int_strides.push_back(int_stride->value);
}
@@ -231,11 +212,7 @@ StructInfo InferStructInfoStridedSlice(const Call& call,
const BlockBuilder& ctx
output_shape.Set(axes[i], GetLength(attrs->begin[i], attrs->end[i],
int_strides[i],
data_shape->values[axes[i]],
attrs->assume_inbound));
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutStridedSlice(const Call& call,
@@ -289,10 +266,7 @@ StructInfo InferStructInfoDynStridedSlice(const Call&
call, const BlockBuilder&
LOG(WARNING) << "When data rank is unknown, dynamic strided slice assumes
begin/end/strides "
"tensors are well-formed. It could produce runtime error
when this assumption "
"turns out to be wrong.";
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
if (data_sinfo->IsUnknownDtype()) {
LOG(WARNING) << "When data type is unknown, dynamic strided slice assumes
to have a valid "
@@ -332,10 +306,7 @@ StructInfo InferStructInfoDynStridedSlice(const Call&
call, const BlockBuilder&
// The output shape will depend on the runtime value in begin/end/strides
tensors.
// TODO(tvm-team): Currently, it is unable to express partially-static
shape. Revisit when
// PrimValue lands.
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, n_axis,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, n_axis);
+ return TensorStructInfo(data_sinfo->dtype, n_axis, data_sinfo->vdevice);
} // namespace relax
// TODO(tvm-team): Register FRelaxInferLayout, TMixedPrecisionPolicy
diff --git a/src/relax/op/tensor/manipulate.cc
b/src/relax/op/tensor/manipulate.cc
index 38b761d04f..12342aecf2 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -71,19 +71,11 @@ StructInfo InferStructInfoBroadcastTo(const Call& call,
const BlockBuilder& ctx)
// Trust the input target shape when there is no possibility to do any
compile-time check.
if (!data_sinfo->shape.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype);
+ return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype,
data_sinfo->vdevice);
}
ShapeStructInfo shape_sinfo =
Downcast<ShapeStructInfo>(data_sinfo->shape.value()->struct_info_);
if (!shape_sinfo->values.defined() || !tgt_shape_sinfo->values.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype);
+ return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype,
data_sinfo->vdevice);
}
arith::Analyzer* analyzer = ctx->GetAnalyzer();
@@ -108,11 +100,7 @@ StructInfo InferStructInfoBroadcastTo(const Call& call,
const BlockBuilder& ctx)
// Todo(relax-team): revisit here for better check on if the tensor length
// is consistent with the length in the given shape.
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype);
+ return TensorStructInfo(/*shape=*/call->args[1], data_sinfo->dtype,
data_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.broadcast_to")
@@ -202,7 +190,7 @@ StructInfo InferStructInfoConcat(const Call& call, const
BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<ConcatAttrs>();
int output_ndim = attrs->axis.defined() ? kUnknownNDim : 1;
DataType output_dtype = DataType::Void();
- VDevice vdev = VDevice();
+ Optional<VDevice> vdev = NullOpt;
bool shape_unknown = false;
bool is_void_dtype = false;
bool vdevice_unknown = false;
@@ -269,12 +257,13 @@ StructInfo InferStructInfoConcat(const Call& call, const
BlockBuilder& ctx) {
if (is_void_dtype) {
output_dtype = DataType::Void();
}
+ if (vdevice_unknown) {
+ vdev = NullOpt;
+ }
+
if (output_ndim == kUnknownNDim) {
- if (!vdevice_unknown) {
- return tensor_sinfo.size() == 1 ? tensor_sinfo[0]
- : TensorStructInfo(output_dtype,
output_ndim, vdev);
- }
- return tensor_sinfo.size() == 1 ? tensor_sinfo[0] :
TensorStructInfo(output_dtype, output_ndim);
+ return tensor_sinfo.size() == 1 ? tensor_sinfo[0]
+ : TensorStructInfo(output_dtype,
output_ndim, vdev);
}
int axis =
@@ -359,10 +348,7 @@ StructInfo InferStructInfoExpandDims(const Call& call,
const BlockBuilder& ctx)
}
if (data_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
int n_new_dim = attrs->axis.size();
@@ -371,10 +357,7 @@ StructInfo InferStructInfoExpandDims(const Call& call,
const BlockBuilder& ctx)
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, output_ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, output_ndim);
+ return TensorStructInfo(data_sinfo->dtype, output_ndim,
data_sinfo->vdevice);
}
std::vector<PrimExpr> output_shape;
@@ -393,11 +376,7 @@ StructInfo InferStructInfoExpandDims(const Call& call,
const BlockBuilder& ctx)
++i_data_shape;
}
ICHECK_EQ(i_data_shape, data_sinfo->ndim);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutExpandDims(const Call& call,
@@ -466,32 +445,20 @@
TVM_REGISTER_GLOBAL("relax.op.flatten").set_body_typed(flatten);
StructInfo InferStructInfoFlatten(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
if (data_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, /*ndim=*/1,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, /*ndim=*/1);
+ return TensorStructInfo(data_sinfo->dtype, /*ndim=*/1,
data_sinfo->vdevice);
} else if (data_sinfo->ndim == 0) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr({1}), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr({1}), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr({1}), data_sinfo->dtype,
data_sinfo->vdevice);
} else if (data_sinfo->ndim == 1) {
return data_sinfo;
}
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, /*ndim=*/1,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, /*ndim=*/1);
+ return TensorStructInfo(data_sinfo->dtype, /*ndim=*/1,
data_sinfo->vdevice);
}
PrimExpr shape_prod = ComputeShapeProduct(data_shape->values);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr({std::move(shape_prod)}),
data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr({std::move(shape_prod)}),
data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr({std::move(shape_prod)}),
data_sinfo->dtype,
+ data_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.flatten")
@@ -535,11 +502,8 @@ StructInfo InferStructInfoLayoutTransform(const Call&
call, const BlockBuilder&
if (data_sinfo->IsUnknownNdim()) {
// Todo(relax-team): revisit here for better check on if the input tensor
has desired ndim.
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size(),
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size());
+ return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size(),
+ data_sinfo->vdevice);
}
// If rank is known, check that it is compatible with the index_map, i.e.,
#dims match.
@@ -551,29 +515,19 @@ StructInfo InferStructInfoLayoutTransform(const Call&
call, const BlockBuilder&
}
if (!data_sinfo->shape.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size(),
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size());
+ return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size(),
+ data_sinfo->vdevice);
}
ShapeStructInfo shape_sinfo =
Downcast<ShapeStructInfo>(data_sinfo->shape.value()->struct_info_);
if (!shape_sinfo->values.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size(),
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size());
+ return TensorStructInfo(data_sinfo->dtype,
/*ndim=*/index_map->final_indices.size(),
+ data_sinfo->vdevice);
}
arith::Analyzer analyzer;
Array<PrimExpr> output_shape =
index_map->MapShape(shape_sinfo->values.value(), &analyzer);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.layout_transform")
@@ -614,10 +568,7 @@ StructInfo InferStructInfoPermuteDims(const Call& call,
const BlockBuilder& ctx)
// Todo(relax-team): revisit here for better check on if the input tensor has
// ndim same as the number of input axes.
if (!attrs->axes.defined() && data_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
if (attrs->axes.defined()) {
@@ -644,20 +595,14 @@ StructInfo InferStructInfoPermuteDims(const Call& call,
const BlockBuilder& ctx)
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
}
std::vector<PrimExpr> new_shape;
new_shape.reserve(data_sinfo->ndim);
for (int i = 0; i < data_sinfo->ndim; ++i) {
new_shape.push_back(data_shape->values[axes[i]]);
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(new_shape), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(new_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(new_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutPermuteDims(const Call& call,
@@ -848,16 +793,10 @@ StructInfo InferStructInfoReshape(const Call& call, const
BlockBuilder& ctx) {
Expr target_shape = call->args[1];
// If shape values are defined, use them
if (target_shape->IsInstance<VarNode>() &&
new_shape_sinfo->values.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(new_shape_sinfo->values.value()),
data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(new_shape_sinfo->values.value()),
data_sinfo->dtype);
- }
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(target_shape, data_sinfo->dtype,
data_sinfo->vdevice.value());
+ return TensorStructInfo(ShapeExpr(new_shape_sinfo->values.value()),
data_sinfo->dtype,
+ data_sinfo->vdevice);
}
- return TensorStructInfo(target_shape, data_sinfo->dtype);
+ return TensorStructInfo(target_shape, data_sinfo->dtype,
data_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.reshape")
@@ -914,26 +853,18 @@ StructInfo InferStructInfoSplit(const Call& call, const
BlockBuilder& ctx) {
}
// Fall back to unknown shape when the input tensor doesn't have ShapeExpr
as shape.
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TupleStructInfo(Array<StructInfo>(
- p_indices->size() + 1,
- TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value())));
- }
return TupleStructInfo(Array<StructInfo>(
- p_indices->size() + 1, TensorStructInfo(data_sinfo->dtype,
data_sinfo->ndim)));
+ p_indices->size() + 1,
+ TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice)));
}
ICHECK_NE(axis, -1);
const auto* axis_length = data_shape->values[axis].as<IntImmNode>();
// Fall back to unknown shape when the input tensor shape at the given
axis is symbolic.
if (axis_length == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TupleStructInfo(Array<StructInfo>(
- p_indices->size() + 1,
- TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value())));
- }
return TupleStructInfo(Array<StructInfo>(
- p_indices->size() + 1, TensorStructInfo(data_sinfo->dtype,
data_sinfo->ndim)));
+ p_indices->size() + 1,
+ TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice)));
}
// Only do output shape inference when all the indices and the total
length are integers.
@@ -950,12 +881,8 @@ StructInfo InferStructInfoSplit(const Call& call, const
BlockBuilder& ctx) {
Array<PrimExpr> shape = data_shape->values;
shape.Set(axis, tvm::max(zero, r - l));
- if (data_sinfo->vdevice.defined()) {
- output_sinfo.push_back(
- TensorStructInfo(ShapeExpr(shape), data_sinfo->dtype,
data_sinfo->vdevice.value()));
- } else {
- output_sinfo.push_back(TensorStructInfo(ShapeExpr(shape),
data_sinfo->dtype));
- }
+ output_sinfo.push_back(
+ TensorStructInfo(ShapeExpr(shape), data_sinfo->dtype,
data_sinfo->vdevice));
}
return TupleStructInfo(output_sinfo);
} else if (const auto* p_n_section =
attrs->indices_or_sections.as<IntImmNode>()) {
@@ -967,13 +894,8 @@ StructInfo InferStructInfoSplit(const Call& call, const
BlockBuilder& ctx) {
}
// Fall back to unknown shape when the input tensor doesn't have ShapeExpr
as shape.
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TupleStructInfo(Array<StructInfo>(
- n_section,
- TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value())));
- }
- return TupleStructInfo(
- Array<StructInfo>(n_section, TensorStructInfo(data_sinfo->dtype,
data_sinfo->ndim)));
+ return TupleStructInfo(Array<StructInfo>(
+ n_section, TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice)));
}
ICHECK_NE(axis, -1);
PrimExpr split_len = ceildiv(data_shape->values[axis], n_section);
@@ -981,22 +903,13 @@ StructInfo InferStructInfoSplit(const Call& call, const
BlockBuilder& ctx) {
// Construct struct info for tensors except the last one.
Array<PrimExpr> shape = data_shape->values;
shape.Set(axis, split_len);
- if (data_sinfo->vdevice.defined()) {
- std::vector<StructInfo> output_sinfo(
- n_section - 1,
- TensorStructInfo(ShapeExpr(shape), data_sinfo->dtype,
data_sinfo->vdevice.value()));
- }
- std::vector<StructInfo> output_sinfo(n_section - 1,
- TensorStructInfo(ShapeExpr(shape),
data_sinfo->dtype));
+ std::vector<StructInfo> output_sinfo(
+ n_section - 1, TensorStructInfo(ShapeExpr(shape), data_sinfo->dtype,
data_sinfo->vdevice));
// Construct struct info for the last tensor.
shape.Set(axis, data_shape->values[axis] - split_len * (n_section - 1));
- if (data_sinfo->vdevice.defined()) {
- output_sinfo.push_back(
- TensorStructInfo(ShapeExpr(shape), data_sinfo->dtype,
data_sinfo->vdevice.value()));
- } else {
- output_sinfo.push_back(TensorStructInfo(ShapeExpr(shape),
data_sinfo->dtype));
- }
+ output_sinfo.push_back(
+ TensorStructInfo(ShapeExpr(shape), data_sinfo->dtype,
data_sinfo->vdevice));
return TupleStructInfo(output_sinfo);
}
ICHECK(false) << "Cannot reach here.";
@@ -1054,10 +967,7 @@ StructInfo InferStructInfoSqueeze(const Call& call, const
BlockBuilder& ctx) {
}
if (data_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
Optional<Array<PrimExpr>> shape_value;
@@ -1072,11 +982,8 @@ StructInfo InferStructInfoSqueeze(const Call& call, const
BlockBuilder& ctx) {
std::vector<int> axes = NormalizeAxes(call, ctx, data_sinfo->ndim,
attrs->axis.value());
if (!shape_value.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim -
axes.size(),
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim -
axes.size());
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim -
axes.size(),
+ data_sinfo->vdevice);
}
for (int i = 0; i < static_cast<int>(axes.size()); ++i) {
// Todo(relax-team): revisit here for better check on if the axis being
squeezed has length 1.
@@ -1098,19 +1005,13 @@ StructInfo InferStructInfoSqueeze(const Call& call,
const BlockBuilder& ctx) {
//
(https://data-apis.org/array-api/latest/API_specification/generated/array_api.squeeze.html).
// Consider discourage usage later.
if (!shape_value.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
for (int i = 0; i < data_sinfo->ndim; ++i) {
// Whenever a dimension length is symbolic, fall back to unknown ndim.
const auto* int_len = shape_value.value()[i].as<IntImmNode>();
if (int_len == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
if (int_len->value == 1) {
axis_removal_mask[i] = true;
@@ -1130,23 +1031,12 @@ StructInfo InferStructInfoSqueeze(const Call& call,
const BlockBuilder& ctx) {
if (static_cast<int>(output_shape.size()) == data_sinfo->ndim) {
return data_sinfo;
} else if (attrs->axis.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, output_shape.size(),
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, output_shape.size());
+ return TensorStructInfo(data_sinfo->dtype, output_shape.size(),
data_sinfo->vdevice);
} else {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
} else {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
}
@@ -1282,17 +1172,11 @@ StructInfo InferStructInfoCollapseSumLike(const Call&
call, const BlockBuilder&
}
if (collapse_target_sinfo->shape.defined()) {
- if (collapse_target_sinfo->vdevice.defined()) {
- return TensorStructInfo(collapse_target_sinfo->shape.value(),
output_dtype,
- collapse_target_sinfo->vdevice.value());
- }
- return TensorStructInfo(collapse_target_sinfo->shape.value(),
output_dtype);
+ return TensorStructInfo(collapse_target_sinfo->shape.value(), output_dtype,
+ collapse_target_sinfo->vdevice);
} else {
- if (collapse_target_sinfo->vdevice.defined()) {
- return TensorStructInfo(output_dtype, collapse_target_sinfo->ndim,
- collapse_target_sinfo->vdevice.value());
- }
- return TensorStructInfo(output_dtype, collapse_target_sinfo->ndim);
+ return TensorStructInfo(output_dtype, collapse_target_sinfo->ndim,
+ collapse_target_sinfo->vdevice);
}
}
@@ -1343,10 +1227,7 @@ StructInfo InferStructInfoCollapseSumTo(const Call&
call, const BlockBuilder& ct
if (data_shape_value.defined() && shape_sinfo->values.defined()) {
CheckCollapseShape(call, ctx, data_shape_value.value(),
shape_sinfo->values.value());
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(/*shape=*/call->args[1], output_dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(/*shape=*/call->args[1], output_dtype);
+ return TensorStructInfo(/*shape=*/call->args[1], output_dtype,
data_sinfo->vdevice);
}
TVM_REGISTER_OP("relax.collapse_sum_to")
@@ -1394,36 +1275,24 @@ StructInfo InferStructInfoRepeat(const Call& call,
const BlockBuilder& ctx) {
// the shape does not changes
return data_sinfo;
} else {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
}
} else {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, 1,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, 1);
+ return TensorStructInfo(data_sinfo->dtype, 1, data_sinfo->vdevice);
}
}
if (!attrs->axis.defined()) {
PrimExpr new_shape =
analyzer->Simplify(ComputeShapeProduct(data_shape->values) *
attrs->repeats);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(Array<PrimExpr>({new_shape})),
data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(Array<PrimExpr>({new_shape})),
data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(Array<PrimExpr>({new_shape})),
data_sinfo->dtype,
+ data_sinfo->vdevice);
}
int axis = NormalizeAxis(call, ctx, data_sinfo->ndim,
attrs->axis.value()->value);
auto shape_array = data_shape->values;
shape_array.Set(axis, analyzer->Simplify(shape_array[axis] *
attrs->repeats));
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(shape_array), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(shape_array), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(shape_array), data_sinfo->dtype,
data_sinfo->vdevice);
}
// TODO(relax-team): implement FRelaxInferLayout for repeat
@@ -1457,24 +1326,14 @@ StructInfo InferStructInfoTile(const Call& call, const
BlockBuilder& ctx) {
if (data_shape == nullptr) {
if (data_sinfo->IsUnknownNdim()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
if (l > ndim) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, l,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, l);
+ return TensorStructInfo(data_sinfo->dtype, l, data_sinfo->vdevice);
} else {
for (auto i : attrs->repeats) {
if (!analyzer->CanProveEqual(i, 1)) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
}
}
// if control reaches here, the shape should not be changed
@@ -1497,10 +1356,7 @@ StructInfo InferStructInfoTile(const Call& call, const
BlockBuilder& ctx) {
}
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
// TODO(relax-team): implement FRelaxInferLayout for tile
@@ -1582,10 +1438,7 @@ StructInfo InferStructInfoScatterElements(const Call&
call, const BlockBuilder&
if (data_sinfo->IsUnknownNdim()) {
// When `data` has unknown rank, assume rest of arguments are correct and
proceed.
// If the assumption turns out to be wrong, runtime error will be
triggered.
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
}
if (!indices_sinfo->IsUnknownNdim() && !updates_sinfo->IsUnknownNdim()) {
@@ -1651,16 +1504,9 @@ StructInfo InferStructInfoScatterElements(const Call&
call, const BlockBuilder&
}
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(data_shape->values), data_sinfo->dtype,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(data_shape->values), data_sinfo->dtype);
- }
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value());
+ return TensorStructInfo(ShapeExpr(data_shape->values), data_sinfo->dtype,
data_sinfo->vdevice);
}
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
}
// TODO(relax-team): implement FRelaxInferLayout for scatter_elements
diff --git a/src/relax/op/tensor/search.cc b/src/relax/op/tensor/search.cc
index 14fa287494..81cb6f87e2 100644
--- a/src/relax/op/tensor/search.cc
+++ b/src/relax/op/tensor/search.cc
@@ -161,21 +161,12 @@ StructInfo InferStructInfoArgmaxArgmin(const Call& call,
const BlockBuilder& ctx
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
if (!attrs->axis.defined() && attrs->keepdims && out_ndim != kUnknownNDim)
{
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(
- ShapeExpr(Array<PrimExpr>(out_ndim, IntImm(out_dtype,
/*value=*/1))), out_dtype,
- data_sinfo->vdevice.value());
- }
return TensorStructInfo(ShapeExpr(Array<PrimExpr>(out_ndim,
IntImm(out_dtype, /*value=*/1))),
- out_dtype);
+ out_dtype, data_sinfo->vdevice);
} else {
- if (data_sinfo->vdevice.defined()) {
- return out_ndim == 0 ? TensorStructInfo(ShapeExpr(Array<PrimExpr>()),
out_dtype,
- data_sinfo->vdevice.value())
- : TensorStructInfo(out_dtype, out_ndim,
data_sinfo->vdevice.value());
- }
- return out_ndim == 0 ? TensorStructInfo(ShapeExpr(Array<PrimExpr>()),
out_dtype)
- : TensorStructInfo(out_dtype, out_ndim);
+ return out_ndim == 0
+ ? TensorStructInfo(ShapeExpr(Array<PrimExpr>()), out_dtype,
data_sinfo->vdevice)
+ : TensorStructInfo(out_dtype, out_ndim, data_sinfo->vdevice);
}
}
@@ -193,10 +184,7 @@ StructInfo InferStructInfoArgmaxArgmin(const Call& call,
const BlockBuilder& ctx
}
}
ICHECK_EQ(static_cast<int>(out_shape.size()), out_ndim);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype,
data_sinfo->vdevice);
}
#define RELAX_REGISTER_ARGMAX_ARGMIN_OP(OpName)
\
diff --git a/src/relax/op/tensor/set.cc b/src/relax/op/tensor/set.cc
index 3920cccadd..29d9d52c60 100644
--- a/src/relax/op/tensor/set.cc
+++ b/src/relax/op/tensor/set.cc
@@ -86,45 +86,22 @@ StructInfo InferStructInfoUnique(const Call& call, const
BlockBuilder& ctx) {
// unique values
if (data_sinfo->ndim == 0) {
- if (data_sinfo->vdevice.defined()) {
-
output_sinfo.push_back(TensorStructInfo(ShapeExpr({IntImm(DataType::Int(64),
/*value=*/1)}),
- data_sinfo->dtype,
data_sinfo->vdevice.value()));
- } else {
- output_sinfo.push_back(
- TensorStructInfo(ShapeExpr({IntImm(DataType::Int(64),
/*value=*/1)}), data_sinfo->dtype));
- }
+
output_sinfo.push_back(TensorStructInfo(ShapeExpr({IntImm(DataType::Int(64),
/*value=*/1)}),
+ data_sinfo->dtype,
data_sinfo->vdevice));
} else if (axis.defined()) {
- if (data_sinfo->vdevice.defined()) {
- output_sinfo.push_back(
- TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice.value()));
- } else {
- output_sinfo.push_back(TensorStructInfo(data_sinfo->dtype,
data_sinfo->ndim));
- }
+ output_sinfo.push_back(
+ TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice));
} else {
- if (data_sinfo->vdevice.defined()) {
- output_sinfo.push_back(
- TensorStructInfo(data_sinfo->dtype, /*ndim=*/1,
data_sinfo->vdevice.value()));
- } else {
- output_sinfo.push_back(TensorStructInfo(data_sinfo->dtype, /*ndim=*/1));
- }
+ output_sinfo.push_back(TensorStructInfo(data_sinfo->dtype, /*ndim=*/1,
data_sinfo->vdevice));
}
// index, reverse and counts
TensorStructInfo int_return{nullptr};
if (data_sinfo->ndim == 0) {
- if (data_sinfo->vdevice.defined()) {
- int_return = TensorStructInfo(ShapeExpr({IntImm(DataType::Int(64),
/*value=*/1)}),
- DataType::Int(64),
data_sinfo->vdevice.value());
- } else {
- int_return =
- TensorStructInfo(ShapeExpr({IntImm(DataType::Int(64),
/*value=*/1)}), DataType::Int(64));
- }
+ int_return = TensorStructInfo(ShapeExpr({IntImm(DataType::Int(64),
/*value=*/1)}),
+ DataType::Int(64), data_sinfo->vdevice);
} else {
- if (data_sinfo->vdevice.defined()) {
- int_return = TensorStructInfo(DataType::Int(64), /*ndim=*/1,
data_sinfo->vdevice.value());
- } else {
- int_return = TensorStructInfo(DataType::Int(64), /*ndim=*/1);
- }
+ int_return = TensorStructInfo(DataType::Int(64), /*ndim=*/1,
data_sinfo->vdevice);
}
for (int i = 0; i < n_int_return; ++i) {
output_sinfo.push_back(int_return);
diff --git a/src/relax/op/tensor/statistical.cc
b/src/relax/op/tensor/statistical.cc
index c450738a1d..b861aafe21 100644
--- a/src/relax/op/tensor/statistical.cc
+++ b/src/relax/op/tensor/statistical.cc
@@ -61,23 +61,13 @@ StructInfo InferStructInfoStatistical(const Call& call,
const BlockBuilder& ctx)
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
if (!attrs->axis.defined() && attrs->keepdims && out_ndim != kUnknownNDim)
{
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(
- ShapeExpr(Array<PrimExpr>(out_ndim, IntImm(DataType::Int(64),
/*value=*/1))),
- data_sinfo->dtype, data_sinfo->vdevice.value());
- }
return TensorStructInfo(
ShapeExpr(Array<PrimExpr>(out_ndim, IntImm(DataType::Int(64),
/*value=*/1))),
- data_sinfo->dtype);
+ data_sinfo->dtype, data_sinfo->vdevice);
} else {
- if (data_sinfo->vdevice.defined()) {
- return out_ndim == 0
- ? TensorStructInfo(ShapeExpr(Array<PrimExpr>()),
data_sinfo->dtype,
- data_sinfo->vdevice.value())
- : TensorStructInfo(data_sinfo->dtype, out_ndim,
data_sinfo->vdevice.value());
- }
- return out_ndim == 0 ? TensorStructInfo(ShapeExpr(Array<PrimExpr>()),
data_sinfo->dtype)
- : TensorStructInfo(data_sinfo->dtype, out_ndim);
+ return out_ndim == 0 ? TensorStructInfo(ShapeExpr(Array<PrimExpr>()),
data_sinfo->dtype,
+ data_sinfo->vdevice)
+ : TensorStructInfo(data_sinfo->dtype, out_ndim,
data_sinfo->vdevice);
}
}
@@ -91,10 +81,7 @@ StructInfo InferStructInfoStatistical(const Call& call,
const BlockBuilder& ctx)
}
}
ICHECK_EQ(static_cast<int>(out_shape.size()), out_ndim);
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype);
+ return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutStatistical(const Call& call,
@@ -172,33 +159,21 @@ StructInfo InferStructInfoCumsum(const Call& call, const
BlockBuilder& ctx) {
// flattened
const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(out_type, data_sinfo->ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(out_type, data_sinfo->ndim);
+ return TensorStructInfo(out_type, data_sinfo->ndim, data_sinfo->vdevice);
} else {
PrimExpr flattened_d = 1;
for (const auto v : data_shape->values) {
flattened_d *= v;
}
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(ShapeExpr(Array<PrimExpr>({flattened_d})),
out_type,
- data_sinfo->vdevice.value());
- }
- return TensorStructInfo(ShapeExpr(Array<PrimExpr>({flattened_d})),
out_type);
+ return TensorStructInfo(ShapeExpr(Array<PrimExpr>({flattened_d})),
out_type,
+ data_sinfo->vdevice);
}
}
if (data_sinfo->shape.defined()) {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(data_sinfo->shape.value(), out_type,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(data_sinfo->shape.value(), out_type);
+ return TensorStructInfo(data_sinfo->shape.value(), out_type,
data_sinfo->vdevice);
} else {
- if (data_sinfo->vdevice.defined()) {
- return TensorStructInfo(out_type, data_sinfo->ndim,
data_sinfo->vdevice.value());
- }
- return TensorStructInfo(out_type, data_sinfo->ndim);
+ return TensorStructInfo(out_type, data_sinfo->ndim, data_sinfo->vdevice);
}
}
