You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "Lunderberg (via GitHub)" <gi...@apache.org> on 2024/02/13 18:45:47 UTC
[PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Lunderberg opened a new pull request, #16563:
URL: https://github.com/apache/tvm/pull/16563
Relax is capable of expressing tensors whose element type is unknown. However, these must typically be replaced with a known dtype prior to compilation, as most operators require known data types prior to legalization. This can be done by using a `relax::MatchCast` node, such as accepting a parameter `arg: R.Tensor([16,16])`, then defining the dtype using `R.match_cast(arg, R.Tensor([16,16],'float16'))`.
However, using a `R.match_cast` node requires knowing which data type should be used in the new `R.Tensor`, and raises an error for an incorrect data type. If an argument may be one of two distinct data types, `R.match_cast` cannot be used to check which data type is in use.
This commit adds Relax operators to read the runtime values of a `DLTensor*` argument. These can be be used to normalize arguments prior to a compute step. For example, pre-processing a model weight that may be provided in either `float16` or `bfloat16` format.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on PR #16563:
URL: https://github.com/apache/tvm/pull/16563#issuecomment-1947309703
> The implementation seems solid, a very good change. I like the parameterized test cases too. The use of the _DLTensorShapeProxy resulted in an elegant UI in tvmscript.
Thank you! I really like polishing up the interface to be as clean as possible. (For my own sake if nothing else, as I am liable to forget a builtin-function name, but am much less likely to forget `obj.dtype`.)
> Idle musing: I wonder if there's any way the PrimFuncs can be unrolled in cases where the parameters are known at compile time.
Prior to `LegalizeOps`, this is implicitly done by the `FNormalize` implementation for each operator. After `LegalizeOps`, not so much. While the `FoldConstant` pass can compile/run TIR functions if their arguments are known, there isn't a good way to indicate that a TIR function only requires the DLTensor struct, and not the data itself.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491708611
##########
src/relax/op/tensor/unpack.cc:
##########
@@ -0,0 +1,349 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file unpack.cc
Review Comment:
I like `relax.inspect` quite a bit, and will update the tensors to that namespace.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg merged PR #16563:
URL: https://github.com/apache/tvm/pull/16563
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on PR #16563:
URL: https://github.com/apache/tvm/pull/16563#issuecomment-1955086643
All CI tests passing, and thank you for the review @slyubomirsky ! I'll follow up with another PR to add inspection of the remainder of the `DLTensor*` fields.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491708293
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
+ to find the invalid usage.
+ """
+ if self.struct_info_ is None:
+ return
+
+ if not isinstance(self.struct_info_, tvm.relax.TensorStructInfo):
+ raise TypeError(
+ f"Runtime unpacking of DLDataType is only implemented for tensors, "
+ f"but was applied to object {self} of type {type(self)}."
+ )
+
+ @property
+ def dtype(self) -> "_DLTensorDTypeProxy":
+ """Returns a proxy object for accessing DLTensor::dtype"""
+ self._check_for_tensor_struct_info()
+ return _DLTensorDTypeProxy(self)
+
+ @property
+ def ndim(self) -> "Expr":
+ """Returns the runtime value of DLTensor::ndim"""
+ self._check_for_tensor_struct_info()
+ op = tvm.ir.Op.get("relax.tensor_ndim")
+ return tvm.relax.Call(op, [self])
+
+ @property
+ def shape(self) -> "_DLTensorShapeProxy":
+ """Returns a proxy object for accessing DLTensor::shape"""
+ self._check_for_tensor_struct_info()
+ return _DLTensorShapeProxy(self)
+
+
+class _DLTensorDTypeProxy(tvm.runtime.ObjectGeneric):
+ """A proxy object for unpacking DLDatatype from DLTensor
+
+ Exposes accessors for `DLDataType` fields `type_code`, `lanes`,
+ and `bits` within a `DLTensor::dtype`. Accessing these fields
Review Comment:
Good point. At the moment, I stuck with the values that have a direct presence elsewhere in Relax, but I agree that it would be good to be able to extract any of the `DLTensor*` fields.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491893261
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
Review Comment:
Thank you, and fixed
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491520155
##########
src/relax/op/tensor/unpack.cc:
##########
@@ -0,0 +1,349 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file unpack.cc
Review Comment:
how about relax.inspect(as a sub namespace for the ops)
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1490323766
##########
src/relax/op/tensor/unpack.cc:
##########
@@ -0,0 +1,349 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file unpack.cc
Review Comment:
I wonder if "unpack" is the best name. Perhaps "accessors" or "runtime_accessors" could be a little more descriptive? I'm not sure.
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
Review Comment:
```suggestion
raised during shape inference, an earlier check makes it easier
```
typo
##########
src/relax/transform/legalize_ops.cc:
##########
@@ -157,17 +167,72 @@ class LegalizeMutator : public ExprMutator {
if (op_node == nullptr) {
return visited_call;
}
-
auto op = GetRef<Op>(op_node);
- std::string op_name(op->name);
- bool is_data_dependent_op = (op_name.find("dynamic") != std::string::npos);
- // Not all shape values are known
- // Data-dependent ops are exception since their output shape will be identified at runtime.
- // Legalizer will insert their shape functions, which are manually registered, and match cast
- // to define symbolic output shape at compile time.
- if (!std::all_of(visited_call->args.begin(), visited_call->args.end(),
- [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); }) ||
- (!is_data_dependent_op && !KnowAllShapeValues(GetStructInfo(visited_call)))) {
+
+ bool can_legalize = [&]() -> bool {
+ bool requires_arg_shapes = requires_arg_shapes_map.get(op, Bool(true))->value;
+ if (!requires_arg_shapes) {
+ // This operator does not require its arguments to have a
+ // known shape/dtype. For example, the "relax.tensor_ndim"
+ // operator can output the dimensionality of a tensor at
+ // runtime, and does not require the dimensionality to be
+ // known at compile-time.
+ return true;
+ }
+
+ bool arg_shapes_defined =
+ std::all_of(visited_call->args.begin(), visited_call->args.end(),
+ [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); });
+ if (!arg_shapes_defined) {
+ // This operator cannot be legalized, because legalization
+ // requires the argument shapes to be known.
+ //
+ // TODO(Lunderberg):
+ //
+ // Improve this fallback case, as failure to legalize can
+ // produce unexpected errors during CodeGenVM. This could
+ // be done by having `R.Tensor(ndim=2)` be syntactic sugar
+ // for `R.Tensor(shape=[m, n])`, where `m` and `n` are new
+ // shape variables. This would allow legalization into
+ // dynamic TIR PrimFuncs.
+ //
+ // This fallback would only be applicable for cases where
+ // both the dtype and the dimensionality are known. While
+ // Relax can express a tensor with unknown dtype and
+ // dimensionality as `TensorStructInfo(DataType::Void(),
+ // kUnknownNDim)`, TIR cannot express unknown dtype or
+ // unknown dimensionality.
Review Comment:
Interesting idea. This could be done by inserting a `MatchCast` that introduces the new vars. Perhaps this should be filed as an issue rather than made a long comment.
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
+ to find the invalid usage.
+ """
+ if self.struct_info_ is None:
+ return
+
+ if not isinstance(self.struct_info_, tvm.relax.TensorStructInfo):
+ raise TypeError(
+ f"Runtime unpacking of DLDataType is only implemented for tensors, "
+ f"but was applied to object {self} of type {type(self)}."
+ )
+
+ @property
+ def dtype(self) -> "_DLTensorDTypeProxy":
+ """Returns a proxy object for accessing DLTensor::dtype"""
+ self._check_for_tensor_struct_info()
+ return _DLTensorDTypeProxy(self)
+
+ @property
+ def ndim(self) -> "Expr":
+ """Returns the runtime value of DLTensor::ndim"""
+ self._check_for_tensor_struct_info()
+ op = tvm.ir.Op.get("relax.tensor_ndim")
+ return tvm.relax.Call(op, [self])
+
+ @property
+ def shape(self) -> "_DLTensorShapeProxy":
+ """Returns a proxy object for accessing DLTensor::shape"""
+ self._check_for_tensor_struct_info()
+ return _DLTensorShapeProxy(self)
+
+
+class _DLTensorDTypeProxy(tvm.runtime.ObjectGeneric):
+ """A proxy object for unpacking DLDatatype from DLTensor
+
+ Exposes accessors for `DLDataType` fields `type_code`, `lanes`,
+ and `bits` within a `DLTensor::dtype`. Accessing these fields
Review Comment:
These are good to have. Offset might also be useful to add, as it might help for memory reuse.
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
+ to find the invalid usage.
+ """
+ if self.struct_info_ is None:
+ return
+
+ if not isinstance(self.struct_info_, tvm.relax.TensorStructInfo):
+ raise TypeError(
+ f"Runtime unpacking of DLDataType is only implemented for tensors, "
+ f"but was applied to object {self} of type {type(self)}."
+ )
+
+ @property
+ def dtype(self) -> "_DLTensorDTypeProxy":
Review Comment:
Why does the return type have to be in quotes? I assume it has to do with the property decorator.
##########
src/relax/transform/legalize_ops.cc:
##########
@@ -157,17 +167,72 @@ class LegalizeMutator : public ExprMutator {
if (op_node == nullptr) {
return visited_call;
}
-
auto op = GetRef<Op>(op_node);
- std::string op_name(op->name);
- bool is_data_dependent_op = (op_name.find("dynamic") != std::string::npos);
- // Not all shape values are known
- // Data-dependent ops are exception since their output shape will be identified at runtime.
- // Legalizer will insert their shape functions, which are manually registered, and match cast
- // to define symbolic output shape at compile time.
- if (!std::all_of(visited_call->args.begin(), visited_call->args.end(),
- [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); }) ||
- (!is_data_dependent_op && !KnowAllShapeValues(GetStructInfo(visited_call)))) {
+
+ bool can_legalize = [&]() -> bool {
+ bool requires_arg_shapes = requires_arg_shapes_map.get(op, Bool(true))->value;
+ if (!requires_arg_shapes) {
+ // This operator does not require its arguments to have a
+ // known shape/dtype. For example, the "relax.tensor_ndim"
+ // operator can output the dimensionality of a tensor at
+ // runtime, and does not require the dimensionality to be
+ // known at compile-time.
+ return true;
+ }
+
+ bool arg_shapes_defined =
+ std::all_of(visited_call->args.begin(), visited_call->args.end(),
+ [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); });
+ if (!arg_shapes_defined) {
+ // This operator cannot be legalized, because legalization
+ // requires the argument shapes to be known.
+ //
+ // TODO(Lunderberg):
+ //
+ // Improve this fallback case, as failure to legalize can
+ // produce unexpected errors during CodeGenVM. This could
+ // be done by having `R.Tensor(ndim=2)` be syntactic sugar
+ // for `R.Tensor(shape=[m, n])`, where `m` and `n` are new
+ // shape variables. This would allow legalization into
+ // dynamic TIR PrimFuncs.
+ //
+ // This fallback would only be applicable for cases where
+ // both the dtype and the dimensionality are known. While
+ // Relax can express a tensor with unknown dtype and
+ // dimensionality as `TensorStructInfo(DataType::Void(),
+ // kUnknownNDim)`, TIR cannot express unknown dtype or
+ // unknown dimensionality.
+ return false;
+ }
+
+ std::string op_name(op->name);
+ bool is_data_dependent_op = (op_name.find("dynamic") != std::string::npos);
+ bool ret_shape_defined = KnowAllShapeValues(GetStructInfo(visited_call));
+ if (!is_data_dependent_op && !ret_shape_defined) {
+ // This operator cannot be legalized, because legalization by
+ // default requires the output shape. The exception is
+ // data-dependent operators (e.g. `R.dynamic_strided_slice`),
+ // where the shape of the output depends on the runtime values
+ // stored in a tensor.
+ //
+ // For data-dependent ops, the output shape will be identified
+ // at runtime. The Legalizer will insert their shape
+ // functions, which are manually registered for each
+ // data-dependent op, and match cast to define symbolic output
+ // shapes. These symbolic output shapes at compile time can
+ // be by later operations to refer to the runtime shape.
+ //
+ // TODO(Lunderberg): Make a new operator attribute
+ // `.set_attr<Bool>("DataDependent")`, rather than relying on
+ // the name of the operator.
Review Comment:
I agree with this, probably should be a separate PR.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491707010
##########
src/relax/transform/legalize_ops.cc:
##########
@@ -157,17 +167,72 @@ class LegalizeMutator : public ExprMutator {
if (op_node == nullptr) {
return visited_call;
}
-
auto op = GetRef<Op>(op_node);
- std::string op_name(op->name);
- bool is_data_dependent_op = (op_name.find("dynamic") != std::string::npos);
- // Not all shape values are known
- // Data-dependent ops are exception since their output shape will be identified at runtime.
- // Legalizer will insert their shape functions, which are manually registered, and match cast
- // to define symbolic output shape at compile time.
- if (!std::all_of(visited_call->args.begin(), visited_call->args.end(),
- [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); }) ||
- (!is_data_dependent_op && !KnowAllShapeValues(GetStructInfo(visited_call)))) {
+
+ bool can_legalize = [&]() -> bool {
+ bool requires_arg_shapes = requires_arg_shapes_map.get(op, Bool(true))->value;
+ if (!requires_arg_shapes) {
+ // This operator does not require its arguments to have a
+ // known shape/dtype. For example, the "relax.tensor_ndim"
+ // operator can output the dimensionality of a tensor at
+ // runtime, and does not require the dimensionality to be
+ // known at compile-time.
+ return true;
+ }
+
+ bool arg_shapes_defined =
+ std::all_of(visited_call->args.begin(), visited_call->args.end(),
+ [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); });
+ if (!arg_shapes_defined) {
+ // This operator cannot be legalized, because legalization
+ // requires the argument shapes to be known.
+ //
+ // TODO(Lunderberg):
+ //
+ // Improve this fallback case, as failure to legalize can
+ // produce unexpected errors during CodeGenVM. This could
+ // be done by having `R.Tensor(ndim=2)` be syntactic sugar
+ // for `R.Tensor(shape=[m, n])`, where `m` and `n` are new
+ // shape variables. This would allow legalization into
+ // dynamic TIR PrimFuncs.
+ //
+ // This fallback would only be applicable for cases where
+ // both the dtype and the dimensionality are known. While
+ // Relax can express a tensor with unknown dtype and
+ // dimensionality as `TensorStructInfo(DataType::Void(),
+ // kUnknownNDim)`, TIR cannot express unknown dtype or
+ // unknown dimensionality.
+ return false;
+ }
+
+ std::string op_name(op->name);
+ bool is_data_dependent_op = (op_name.find("dynamic") != std::string::npos);
+ bool ret_shape_defined = KnowAllShapeValues(GetStructInfo(visited_call));
+ if (!is_data_dependent_op && !ret_shape_defined) {
+ // This operator cannot be legalized, because legalization by
+ // default requires the output shape. The exception is
+ // data-dependent operators (e.g. `R.dynamic_strided_slice`),
+ // where the shape of the output depends on the runtime values
+ // stored in a tensor.
+ //
+ // For data-dependent ops, the output shape will be identified
+ // at runtime. The Legalizer will insert their shape
+ // functions, which are manually registered for each
+ // data-dependent op, and match cast to define symbolic output
+ // shapes. These symbolic output shapes at compile time can
+ // be by later operations to refer to the runtime shape.
+ //
+ // TODO(Lunderberg): Make a new operator attribute
+ // `.set_attr<Bool>("DataDependent")`, rather than relying on
+ // the name of the operator.
Review Comment:
Agreed. The note is more so that it has a record somewhere, but it would be enough of a change that it should be part of a separate PR. As it was, I wanted to make as few changes to `LegalizeOps` as possible in this PR.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491670908
##########
src/relax/transform/legalize_ops.cc:
##########
@@ -157,17 +167,72 @@ class LegalizeMutator : public ExprMutator {
if (op_node == nullptr) {
return visited_call;
}
-
auto op = GetRef<Op>(op_node);
- std::string op_name(op->name);
- bool is_data_dependent_op = (op_name.find("dynamic") != std::string::npos);
- // Not all shape values are known
- // Data-dependent ops are exception since their output shape will be identified at runtime.
- // Legalizer will insert their shape functions, which are manually registered, and match cast
- // to define symbolic output shape at compile time.
- if (!std::all_of(visited_call->args.begin(), visited_call->args.end(),
- [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); }) ||
- (!is_data_dependent_op && !KnowAllShapeValues(GetStructInfo(visited_call)))) {
+
+ bool can_legalize = [&]() -> bool {
+ bool requires_arg_shapes = requires_arg_shapes_map.get(op, Bool(true))->value;
+ if (!requires_arg_shapes) {
+ // This operator does not require its arguments to have a
+ // known shape/dtype. For example, the "relax.tensor_ndim"
+ // operator can output the dimensionality of a tensor at
+ // runtime, and does not require the dimensionality to be
+ // known at compile-time.
+ return true;
+ }
+
+ bool arg_shapes_defined =
+ std::all_of(visited_call->args.begin(), visited_call->args.end(),
+ [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); });
+ if (!arg_shapes_defined) {
+ // This operator cannot be legalized, because legalization
+ // requires the argument shapes to be known.
+ //
+ // TODO(Lunderberg):
+ //
+ // Improve this fallback case, as failure to legalize can
+ // produce unexpected errors during CodeGenVM. This could
+ // be done by having `R.Tensor(ndim=2)` be syntactic sugar
+ // for `R.Tensor(shape=[m, n])`, where `m` and `n` are new
+ // shape variables. This would allow legalization into
+ // dynamic TIR PrimFuncs.
+ //
+ // This fallback would only be applicable for cases where
+ // both the dtype and the dimensionality are known. While
+ // Relax can express a tensor with unknown dtype and
+ // dimensionality as `TensorStructInfo(DataType::Void(),
+ // kUnknownNDim)`, TIR cannot express unknown dtype or
+ // unknown dimensionality.
Review Comment:
Good call, though I haven't had time to flesh out the idea yet. This would be part of a general cleanup I'd propose for the `StructInfo` interactions:
* Remove the `Optional<Expr> shape` and `int ndim` in `TensorStructInfo`. Instead, have `Optional<ShapeStructInfo>`.
* Remove the `ndim` field from `ShapeStructInfo`. Instead, the `ndim` is unknown when `Optional<Array<PrimExpr>>` is `NullOpt`.
* If the dimensionality of a `ShapeStructInfo` is known, every dimension must have an associated `PrimExpr`. The constructor that accepts `ndim` initializes fresh TIR variables to represent the unknown size.
--
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
Re: [PR] [Relax] Implement operators to read runtime DLTensor* information [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491661307
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
+ to find the invalid usage.
+ """
+ if self.struct_info_ is None:
+ return
+
+ if not isinstance(self.struct_info_, tvm.relax.TensorStructInfo):
+ raise TypeError(
+ f"Runtime unpacking of DLDataType is only implemented for tensors, "
+ f"but was applied to object {self} of type {type(self)}."
+ )
+
+ @property
+ def dtype(self) -> "_DLTensorDTypeProxy":
Review Comment:
More the order of definitions in the file. Function annotations are evaluated when the class is being defined. Since the `_DLTensorDTypeProxy` class is defined lower in the file, the type annotations are provided as a string, rather than as a class object.
--
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