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