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-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
