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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
