gemini-code-assist[bot] commented on code in PR #19853: URL: https://github.com/apache/tvm/pull/19853#discussion_r3447803160
########## src/relax/ir/type_functor.cc: ########## @@ -0,0 +1,151 @@ +/* + * 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 type_functor.cc + * \brief Implementations of Relax type functors. + */ +#include <tvm/ffi/cast.h> +#include <tvm/relax/type_functor.h> + +namespace tvm { +namespace relax { + +void TypeVisitor::VisitType_(const ObjectTypeNode* op) {} + +void TypeVisitor::VisitType_(const PrimTypeNode* op) { + if (op->value.defined()) { + this->VisitTypeExprField(op->value.value()); + } +} + +void TypeVisitor::VisitType_(const ShapeTypeNode* op) { + if (op->values.defined()) { + for (PrimExpr value : op->values.value()) { + this->VisitTypeExprField(value); + } + } +} + +void TypeVisitor::VisitType_(const TensorTypeNode* op) { + if (op->shape.defined()) { + this->VisitTypeExprField(op->shape.value()); + } +} + +void TypeVisitor::VisitType_(const distributed::DTensorTypeNode* op) { + this->VisitType(op->tensor_ty); +} + +void TypeVisitor::VisitType_(const TupleTypeNode* op) { + for (Type field : op->fields) { + this->VisitType(field); + } +} + +void TypeVisitor::VisitType_(const FuncTypeNode* op) { + if (op->params.defined()) { + for (Type param : op->params.value()) { + this->VisitType(param); + } + } + this->VisitType(op->ret); +} + +Type TypeMutator::VisitType_(const ObjectTypeNode* op) { return ffi::GetRef<Type>(op); } + +Type TypeMutator::VisitType_(const PrimTypeNode* op) { + if (!op->value.defined()) { + return ffi::GetRef<Type>(op); + } + + auto new_expr = VisitTypeExprField(op->value.value()); + if (new_expr.same_as(op->value)) { + return ffi::GetRef<Type>(op); + } else { + return PrimType(new_expr); + } +} + +Type TypeMutator::VisitType_(const ShapeTypeNode* op) { + ffi::Optional<ffi::Array<PrimExpr>> values; + + if (op->values.defined()) { + // if no changes are made the original array will be returned. + values = op->values.value().Map( + [this](const PrimExpr& expr) { return this->VisitTypeExprField(expr); }); + } + + if (values.same_as(op->values)) { + return ffi::GetRef<Type>(op); + } else { + return ShapeType(values.value(), op->span); + } +} + +Type TypeMutator::VisitType_(const TensorTypeNode* op) { + ffi::Optional<Expr> shape; + + if (op->shape.defined()) { + shape = this->VisitTypeExprField(op->shape.value()); + } + + VDevice vdev = op->vdevice.value_or(VDevice()); + + if (shape.same_as(op->shape)) { + return ffi::GetRef<Type>(op); + } else { + return TensorType(shape.value(), op->dtype, vdev, op->span); + } +} + +Type TypeMutator::VisitType_(const distributed::DTensorTypeNode* op) { + TensorType tensor_ty = Downcast<TensorType>(this->VisitType(op->tensor_ty)); + return distributed::DTensorType(tensor_ty, op->device_mesh, op->placement); +} + +Type TypeMutator::VisitType_(const TupleTypeNode* op) { + ffi::Array<Type> fields = op->fields.Map([this](const Type& ty) { return this->VisitType(ty); }); + + if (fields.same_as(op->fields)) { + return ffi::GetRef<Type>(op); + } else { + return TupleType(fields, op->span); + } +} + +Type TypeMutator::VisitType_(const FuncTypeNode* op) { + ffi::Optional<ffi::Array<Type>> params; + + if (op->params.defined()) { + params = op->params.value().Map([this](const Type& ty) { return this->VisitType(ty); }); + } + + Type ret = this->VisitType(op->ret); + + if (params.same_as(op->params) && ret.same_as(op->ret)) { + return ffi::GetRef<Type>(op); + } else { + TVM_FFI_ICHECK(ret.defined()) << "FuncType that contains params must contain ret"; + return FuncType(params.value(), ret, op->purity, op->span); + } Review Comment:  **Critical Bug**: Potential empty optional dereference. If `op->params` is not defined (e.g., for an opaque function type), `params` remains undefined (`std::nullopt`). If the return type `ret` changes, the code enters the `else` branch and calls `params.value()`, which will crash. We should check if `params` is defined, and if not, construct an opaque function type using `FuncType::OpaqueFunc`. ```c } else { TVM_FFI_ICHECK(ret.defined()) << "FuncType must contain ret"; if (params.defined()) { return FuncType(params.value(), ret, op->purity, op->span); } else { return FuncType::OpaqueFunc(ret, op->purity, op->span); } } ``` ########## include/tvm/relax/type.h: ########## @@ -19,130 +19,438 @@ /*! * \file tvm/relax/type.h - * \brief Relax Types. + * \brief Relax types, including the richer dependent Relax type nodes. */ #ifndef TVM_RELAX_TYPE_H_ #define TVM_RELAX_TYPE_H_ #include <tvm/ffi/function.h> #include <tvm/ffi/reflection/registry.h> #include <tvm/ir/attrs.h> +#include <tvm/ir/cast.h> #include <tvm/ir/env_func.h> +#include <tvm/ir/global_info.h> #include <tvm/ir/type.h> #include <tvm/tirx/expr.h> #include <string> +#include <utility> namespace tvm { namespace relax { +using Expr = RelaxExpr; +using ExprNode = RelaxExprNode; + +class BlockBuilder; +class Call; + /*! \brief Indicates the number of dimensions of a tensor is unknown at compile time. */ static constexpr int kUnknownNDim = -1; +using tvm::TupleType; +using tvm::TupleTypeNode; + +class PackedFuncTypeNode : public TypeNode { + public: + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef<PackedFuncTypeNode>(); + } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeNode, TypeNode); +}; + +class PackedFuncType : public Type { + public: + TVM_DLL PackedFuncType(Span span = Span()); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeNode); +}; + +/*! + * \brief Base type of all Relax type information. + * + * Type stores possible type information deduced during compile-time. + * It encapsulates both static type and runtime information such as shape. + * + * Type of each non-primitive Expr can be deduced during compilation in a + * "best-effort" manner. + * + * When ty appears in function parameter and return signatures, it + * implies a runtime check that matches the type information with the value. + * + * When it appears in Expr, it follows "assume-semantics", which means the + * compiler will take the deduced information as it is and only do best effort + * proofs and checks. + * + * Each type can be uniquely erased to a static-type. The compiler will + * still compile the code, with less information, when we erase to the static + * type. + * + * If a Type contains an Expr field, then that field must already be + * normalized through NormalizeArg. This invariant is checked in constructors + * and simplifies assumptions during type deduction. + */ +/*! + * \brief Opaque object. + */ +class ObjectTypeNode : public TypeNode { + public: + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef<ObjectTypeNode>(); + } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ObjectType", ObjectTypeNode, TypeNode); +}; + +/*! + * \brief Managed reference to ObjectTypeNode. + * \sa ObjectTypeNode + */ +class ObjectType : public Type { + public: + TVM_DLL ObjectType(Span span = Span()); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ObjectType, Type, ObjectTypeNode); +}; + +/*! + * \brief Primitive value. + */ +class PrimTypeNode : public TypeNode { + public: + /*! \brief Underlying primitive value, if known */ + ffi::Optional<PrimExpr> value; + + /*! \brief Underlying data type of the primitive value */ + DataType dtype; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef<PrimTypeNode>() + .def_ro("value", &PrimTypeNode::value) + .def_ro("dtype", &PrimTypeNode::dtype); + } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PrimType", PrimTypeNode, TypeNode); +}; + +/*! + * \brief Managed reference to PrimTypeNode. + * \sa PrimTypeNode + */ +class PrimType : public Type { + public: + /* Construct a PrimType with a known dtype, but unknown value */ + TVM_DLL PrimType(DataType dtype, Span span = Span()); + + /* Construct a PrimType with a known value */ + TVM_DLL PrimType(PrimExpr value, Span span = Span()); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimType, Type, PrimTypeNode); +}; + +/*! + * \brief Type of shape value. + */ class ShapeTypeNode : public TypeNode { public: - /*! \brief size of the shape. */ + /*! \brief optionally stores the symbolic value patterns of the shape */ + ffi::Optional<ffi::Array<PrimExpr>> values; + /*! + * \brief The number of dimension of the shape, can be unknown. + * \sa kUnknownNDim + */ int ndim; + /*! \return Whether the type contains unknown ndim. */ + bool IsUnknownNdim() const { return ndim == kUnknownNDim; } + static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<ShapeTypeNode>().def_ro("ndim", &ShapeTypeNode::ndim); + refl::ObjectDef<ShapeTypeNode>() + .def_ro("values", &ShapeTypeNode::values) + .def_ro("ndim", &ShapeTypeNode::ndim); } TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeType", ShapeTypeNode, TypeNode); }; +/*! + * \brief Managed reference to ShapeTypeNode. + * \sa ShapeTypeNode + */ class ShapeType : public Type { public: + /*! + * \brief Construction with known symbolic shape patterns + * \param values The symbolic shape values + * \param span The span of the AST. + */ + TVM_DLL ShapeType(ffi::Array<PrimExpr> values, Span span = Span()); + /*! + * \brief Construction with known unknown symbolic shape patterns. + * \param ndim Number of dimensions -- can be kUnknownNDim + * \param span The span of the AST. + */ TVM_DLL ShapeType(int ndim, Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeType, Type, ShapeTypeNode); }; /*! - * \brief Dynamic version of TensorType - * - * Use relax::TensorStructInfo for more detailed (possibly dynamic) shape constrains + * \brief Type of Tensor. */ class TensorTypeNode : public TypeNode { public: /*! - * \brief The number of dimensions of the tensor, use -1 to denote tensor with unknown number of - * dimensions. + * \brief optionally store the shape expression of the tensor. + * \note shape must be normalized: it can only be std::nullopt or ShapeExpr or Var. */ - int ndim; + ffi::Optional<Expr> shape; + /*! \brief The virtual device, indicates where the tensor + * is expected to be executed. + */ + ffi::Optional<VDevice> vdevice; /*! \brief The content data type, use void to denote the dtype is unknown. */ DataType dtype; + /*! + * \brief The number of dimension of the tensor, can be unknown. + * \sa kUnknownNDim + */ + int ndim; + + /*! \return Whether the type contains unknown ndim. */ + bool IsUnknownNdim() const { return ndim == kUnknownNDim; } + + /*! \return Whether the type contains unknown dtype. */ + bool IsUnknownDtype() const { return dtype.is_void(); } + + /*! \return Shape if it is known. */ + ffi::Optional<ffi::Array<PrimExpr>> GetShape() const { + if (!shape.defined()) return {}; + ShapeType shape_ty = Downcast<ShapeType>(this->shape.value()->ty); + return shape_ty->values; Review Comment:  **Defensive Programming**: Potential null pointer dereference. If `this->shape.value()->ty` is not defined (null), `Downcast<ShapeType>` will fail or crash. We should defensively check if `ty` is defined and is of type `ShapeTypeNode` before accessing its values. ```c if (!shape.defined() || !shape.value()->ty.defined()) return {}; if (const auto* shape_ty = shape.value()->ty.as<ShapeTypeNode>()) { return shape_ty->values; } return {}; ``` ########## src/relax/distributed/type.cc: ########## @@ -118,29 +118,27 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // DTensor -DTensorStructInfo::DTensorStructInfo(TensorStructInfo tensor_sinfo, DeviceMesh device_mesh, - Placement placement, Span span) { +DTensorType::DTensorType(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement, + Span span) { TVM_FFI_CHECK_EQ(device_mesh->shape.size(), placement->dim_specs.size(), ValueError) << "The device mesh and placement must have the same dimension size"; Review Comment:  **Defensive Programming**: Potential null pointer dereference. If `device_mesh` or `placement` is null (not defined), dereferencing them will cause a segmentation fault. We should defensively check if they are defined first. ```c TVM_FFI_CHECK(device_mesh.defined(), ValueError) << "device_mesh must be defined"; TVM_FFI_CHECK(placement.defined(), ValueError) << "placement must be defined"; TVM_FFI_CHECK_EQ(device_mesh->shape.size(), placement->dim_specs.size(), ValueError) << "The device mesh and placement must have the same dimension size"; ``` ########## src/relax/ir/dependent_type.cc: ########## @@ -0,0 +1,219 @@ +/* + * 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 src/relax/ir/dependent_type.cc + * \brief Relax type nodes. + */ +#include <tvm/ffi/function.h> +#include <tvm/ffi/reflection/registry.h> +#include <tvm/relax/analysis.h> +#include <tvm/relax/type.h> +#include <tvm/relax/type_functor.h> + +namespace tvm { +namespace relax { + +TVM_FFI_STATIC_INIT_BLOCK() { + ObjectTypeNode::RegisterReflection(); + PrimTypeNode::RegisterReflection(); + ShapeTypeNode::RegisterReflection(); + TensorTypeNode::RegisterReflection(); + FuncTypeNode::RegisterReflection(); +} + +ObjectType::ObjectType(Span span) { + ffi::ObjectPtr<ObjectTypeNode> n = ffi::make_object<ObjectTypeNode>(); + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("relax.ObjectType", [](Span span) { return ObjectType(span); }); +} + +// Prim +PrimType::PrimType(PrimExpr value, Span span) { + ffi::ObjectPtr<PrimTypeNode> n = ffi::make_object<PrimTypeNode>(); + n->dtype = value->dtype; + n->value = std::move(value); + n->span = span; + data_ = std::move(n); +} + +PrimType::PrimType(DataType dtype, Span span) { + ffi::ObjectPtr<PrimTypeNode> n = ffi::make_object<PrimTypeNode>(); + n->dtype = dtype; + n->value = std::nullopt; + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef() + .def("relax.PrimTypeFromDtype", + [](DataType dtype, Span span) { return PrimType(dtype, span); }) + .def("relax.PrimTypeFromValue", + [](PrimExpr value, Span span) { return PrimType(value, span); }); +} + +// Shape +ShapeType::ShapeType(ffi::Array<PrimExpr> values, Span span) { + ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>(); + n->ndim = static_cast<int>(values.size()); + n->values = values.Map([](PrimExpr value) { + if (value->IsInstance<IntImmNode>()) { + return tvm::cast(DataType::Int(64), value); + } + TVM_FFI_ICHECK(value.dtype() == DataType::Int(64)) + << "the value in ShapeType can only have dtype of int64"; + return value; + }); + n->span = span; + data_ = std::move(n); +} + +ShapeType::ShapeType(int ndim, Span span) { + ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>(); + TVM_FFI_ICHECK_GE(ndim, -1) << "ndim of ShapeType must be >= -1, but got " << ndim; Review Comment:  **Compilation Risk**: `TVM_FFI_ICHECK_GE` might not be defined in the TVM FFI headers. It is safer and more consistent to use the standard `TVM_FFI_ICHECK` macro with a comparison operator. ```suggestion TVM_FFI_ICHECK(ndim >= -1) << "ndim of ShapeType must be >= -1, but got " << ndim; ``` ########## src/relax/ir/dependent_type.cc: ########## @@ -0,0 +1,219 @@ +/* + * 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 src/relax/ir/dependent_type.cc + * \brief Relax type nodes. + */ +#include <tvm/ffi/function.h> +#include <tvm/ffi/reflection/registry.h> +#include <tvm/relax/analysis.h> +#include <tvm/relax/type.h> +#include <tvm/relax/type_functor.h> + +namespace tvm { +namespace relax { + +TVM_FFI_STATIC_INIT_BLOCK() { + ObjectTypeNode::RegisterReflection(); + PrimTypeNode::RegisterReflection(); + ShapeTypeNode::RegisterReflection(); + TensorTypeNode::RegisterReflection(); + FuncTypeNode::RegisterReflection(); +} + +ObjectType::ObjectType(Span span) { + ffi::ObjectPtr<ObjectTypeNode> n = ffi::make_object<ObjectTypeNode>(); + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("relax.ObjectType", [](Span span) { return ObjectType(span); }); +} + +// Prim +PrimType::PrimType(PrimExpr value, Span span) { + ffi::ObjectPtr<PrimTypeNode> n = ffi::make_object<PrimTypeNode>(); + n->dtype = value->dtype; + n->value = std::move(value); + n->span = span; + data_ = std::move(n); +} + +PrimType::PrimType(DataType dtype, Span span) { + ffi::ObjectPtr<PrimTypeNode> n = ffi::make_object<PrimTypeNode>(); + n->dtype = dtype; + n->value = std::nullopt; + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef() + .def("relax.PrimTypeFromDtype", + [](DataType dtype, Span span) { return PrimType(dtype, span); }) + .def("relax.PrimTypeFromValue", + [](PrimExpr value, Span span) { return PrimType(value, span); }); +} + +// Shape +ShapeType::ShapeType(ffi::Array<PrimExpr> values, Span span) { + ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>(); + n->ndim = static_cast<int>(values.size()); + n->values = values.Map([](PrimExpr value) { + if (value->IsInstance<IntImmNode>()) { + return tvm::cast(DataType::Int(64), value); + } + TVM_FFI_ICHECK(value.dtype() == DataType::Int(64)) + << "the value in ShapeType can only have dtype of int64"; + return value; + }); + n->span = span; + data_ = std::move(n); +} + +ShapeType::ShapeType(int ndim, Span span) { + ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>(); + TVM_FFI_ICHECK_GE(ndim, -1) << "ndim of ShapeType must be >= -1, but got " << ndim; + n->ndim = ndim; + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def( + "relax.ShapeType", [](ffi::Optional<ffi::Array<PrimExpr>> values, int ndim, Span span) { + if (values.defined()) { + TVM_FFI_CHECK_EQ(ndim, kUnknownNDim, ValueError) << "Cannot both specify values and ndim"; + return ShapeType(values.value(), span); + } else { + return ShapeType(ndim, span); + } + }); +} + +// Tensor +TensorType::TensorType(Expr shape, DataType dtype, ffi::Optional<VDevice> vdevice, Span span) { + ffi::ObjectPtr<TensorTypeNode> n = ffi::make_object<TensorTypeNode>(); + // assign ndim before move + ffi::Optional<ShapeType> shape_ty = MatchType<ShapeType>(shape); + TVM_FFI_ICHECK(shape_ty) << "We expect shape to contain pre-set shape type"; + TVM_FFI_ICHECK(shape.defined()) << "Must provide a shape in this constructor"; + TVM_FFI_ICHECK(shape->IsInstance<ShapeExprNode>() || shape->IsInstance<VarNode>()) + << "We require shape to be normalized when constructing TensorType"; + n->ndim = shape_ty.value()->ndim; + // assign rest of the fields. + n->shape = std::move(shape); + n->dtype = dtype; + n->vdevice = vdevice; + n->span = span; + data_ = std::move(n); +} + +TensorType::TensorType(DataType dtype, int ndim, ffi::Optional<VDevice> vdevice, Span span) { + ffi::ObjectPtr<TensorTypeNode> n = ffi::make_object<TensorTypeNode>(); + TVM_FFI_ICHECK_GE(ndim, -1) << "ndim of TensorType must be >= -1, but got " << ndim; Review Comment:  **Compilation Risk**: `TVM_FFI_ICHECK_GE` might not be defined in the TVM FFI headers. It is safer and more consistent to use the standard `TVM_FFI_ICHECK` macro with a comparison operator. ```suggestion TVM_FFI_ICHECK(ndim >= -1) << "ndim of TensorType must be >= -1, but got " << ndim; ``` ########## src/relax/ir/dependent_type.cc: ########## @@ -0,0 +1,219 @@ +/* + * 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 src/relax/ir/dependent_type.cc + * \brief Relax type nodes. + */ +#include <tvm/ffi/function.h> +#include <tvm/ffi/reflection/registry.h> +#include <tvm/relax/analysis.h> +#include <tvm/relax/type.h> +#include <tvm/relax/type_functor.h> + +namespace tvm { +namespace relax { + +TVM_FFI_STATIC_INIT_BLOCK() { + ObjectTypeNode::RegisterReflection(); + PrimTypeNode::RegisterReflection(); + ShapeTypeNode::RegisterReflection(); + TensorTypeNode::RegisterReflection(); + FuncTypeNode::RegisterReflection(); +} + +ObjectType::ObjectType(Span span) { + ffi::ObjectPtr<ObjectTypeNode> n = ffi::make_object<ObjectTypeNode>(); + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("relax.ObjectType", [](Span span) { return ObjectType(span); }); +} + +// Prim +PrimType::PrimType(PrimExpr value, Span span) { + ffi::ObjectPtr<PrimTypeNode> n = ffi::make_object<PrimTypeNode>(); + n->dtype = value->dtype; + n->value = std::move(value); + n->span = span; + data_ = std::move(n); +} + +PrimType::PrimType(DataType dtype, Span span) { + ffi::ObjectPtr<PrimTypeNode> n = ffi::make_object<PrimTypeNode>(); + n->dtype = dtype; + n->value = std::nullopt; + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef() + .def("relax.PrimTypeFromDtype", + [](DataType dtype, Span span) { return PrimType(dtype, span); }) + .def("relax.PrimTypeFromValue", + [](PrimExpr value, Span span) { return PrimType(value, span); }); +} + +// Shape +ShapeType::ShapeType(ffi::Array<PrimExpr> values, Span span) { + ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>(); + n->ndim = static_cast<int>(values.size()); + n->values = values.Map([](PrimExpr value) { + if (value->IsInstance<IntImmNode>()) { + return tvm::cast(DataType::Int(64), value); + } + TVM_FFI_ICHECK(value.dtype() == DataType::Int(64)) + << "the value in ShapeType can only have dtype of int64"; + return value; + }); + n->span = span; + data_ = std::move(n); +} + +ShapeType::ShapeType(int ndim, Span span) { + ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>(); + TVM_FFI_ICHECK_GE(ndim, -1) << "ndim of ShapeType must be >= -1, but got " << ndim; + n->ndim = ndim; + n->span = span; + data_ = std::move(n); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def( + "relax.ShapeType", [](ffi::Optional<ffi::Array<PrimExpr>> values, int ndim, Span span) { + if (values.defined()) { + TVM_FFI_CHECK_EQ(ndim, kUnknownNDim, ValueError) << "Cannot both specify values and ndim"; + return ShapeType(values.value(), span); + } else { + return ShapeType(ndim, span); + } + }); +} + +// Tensor +TensorType::TensorType(Expr shape, DataType dtype, ffi::Optional<VDevice> vdevice, Span span) { + ffi::ObjectPtr<TensorTypeNode> n = ffi::make_object<TensorTypeNode>(); + // assign ndim before move + ffi::Optional<ShapeType> shape_ty = MatchType<ShapeType>(shape); + TVM_FFI_ICHECK(shape_ty) << "We expect shape to contain pre-set shape type"; + TVM_FFI_ICHECK(shape.defined()) << "Must provide a shape in this constructor"; Review Comment:  **Critical Bug**: Potential null pointer dereference. `MatchType<ShapeType>(shape)` is called on `shape` before checking if `shape` is defined on line 121. If `shape` is null (undefined), `MatchType` will dereference it and cause a segmentation fault. Please check `shape.defined()` first before calling `MatchType`. ```suggestion TVM_FFI_ICHECK(shape.defined()) << "Must provide a shape in this constructor"; ffi::Optional<ShapeType> shape_ty = MatchType<ShapeType>(shape); TVM_FFI_ICHECK(shape_ty) << "We expect shape to contain pre-set shape type"; ``` -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
