apache · Lunderberg · Jun 14, 2024
diff --git a/include/tvm/runtime/container/variant.h b/include/tvm/runtime/container/variant.h
@@ -82,7 +82,7 @@ class Variant : public ObjectRef {
  public:
   /* \brief Helper utility to check if the type is part of the variant */
   template <typename T>
-  static constexpr bool is_variant = (std::is_same_v<T, V> || ...);
+  static constexpr bool is_variant = (std::is_base_of_v<V, T> || ...);
 
   /* \brief Helper utility for SFINAE if the type is part of the variant */
   template <typename T>

diff --git a/include/tvm/tir/function.h b/include/tvm/tir/function.h
@@ -264,7 +264,7 @@ class TensorIntrin : public ObjectRef {
  *              B[vi, vj] = A[vi, vj]
  * \endcode
  */
-PrimFunc Specialize(PrimFunc func, const Map<Var, ObjectRef>& param_map);
+PrimFunc Specialize(PrimFunc func, const Map<Var, Variant<Buffer, PrimExpr>>& param_map);
 
 /*!
  * \brief PrimFunc specific attribute names.

diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
@@ -36,7 +36,7 @@ namespace relax {
 TVM_REGISTER_NODE_TYPE(InitAttrs);
 
 /* relax.full */
-Expr full(ObjectRef shape, Expr fill_value, DataType dtype) {
+Expr full(Variant<Expr, Array<PrimExpr>> shape, Expr fill_value, DataType dtype) {
   Expr shape_in_expr{nullptr};
   if (const auto* expr = shape.as<ExprNode>()) {
     shape_in_expr = GetRef<Expr>(expr);

diff --git a/src/relax/op/tensor/create.h b/src/relax/op/tensor/create.h
@@ -39,7 +39,7 @@ namespace relax {
  * If dtype is not given, it will by default use the dtype of fill_value.
  * \return The result tensor.
  */
-Expr full(ObjectRef shape, Expr fill_value, DataType dtype);
+Expr full(Variant<Expr, Array<PrimExpr>> shape, Expr fill_value, DataType dtype);
 
 /*!
  * \brief Construct a tensor such that

diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc
@@ -652,7 +652,7 @@ TVM_REGISTER_OP("relax.permute_dims")
     .set_attr<Bool>("FPurity", Bool(true));
 
 /* relax.reshape */
-Expr ConvertNewShapeToExpr(const Expr& data, const ObjectRef& shape) {
+Expr ConvertNewShapeToExpr(const Expr& data, const Variant<Expr, Array<PrimExpr>>& shape) {
   const ArrayNode* array;
   // Treat shape expressions as constant arrays to handle special values.
   if (const auto* e = shape.as<ShapeExprNode>()) {
@@ -745,7 +745,7 @@ Expr ConvertNewShapeToExpr(const Expr& data, const ObjectRef& shape) {
   return ShapeExpr(array_ref);
 }
 
-Expr reshape(Expr x, ObjectRef shape) {
+Expr reshape(Expr x, Variant<Expr, Array<PrimExpr>> shape) {
   Expr shape_in_expr = ConvertNewShapeToExpr(x, shape);
   static const Op& op = Op::Get("relax.reshape");
   return Call(op, {std::move(x), std::move(shape_in_expr)}, Attrs(), {});
@@ -810,7 +810,7 @@ TVM_REGISTER_OP("relax.reshape")
 /* relax.split */
 TVM_REGISTER_NODE_TYPE(SplitAttrs);
 
-Expr split(Expr x, ObjectRef indices_or_sections, int axis) {
+Expr split(Expr x, Variant<IntImm, Array<IntImm>> indices_or_sections, int axis) {
   ObjectPtr<SplitAttrs> attrs = make_object<SplitAttrs>();
   if (const auto* indices = indices_or_sections.as<ArrayNode>()) {
     for (int i = 0; i < static_cast<int>(indices->size()); ++i) {

diff --git a/src/relax/op/tensor/manipulate.h b/src/relax/op/tensor/manipulate.h
@@ -88,7 +88,7 @@ Expr permute_dims(Expr x, Optional<Array<Integer>> axes);
  * It is required to be either an Array of PrimExpr, or a Shape in Relax
  * \return The reshaped result.
  */
-Expr reshape(Expr x, ObjectRef shape);
+Expr reshape(Expr x, Variant<Expr, Array<PrimExpr>> shape);
 
 /*!
  * \brief Split input tensor along axis by sections or indices.
@@ -103,7 +103,7 @@ Expr reshape(Expr x, ObjectRef shape);
  * \param axis The axis over which to split.
  * \return The computed result.
  */
-Expr split(Expr x, ObjectRef indices_or_sections, int axis);
+Expr split(Expr x, Variant<IntImm, Array<IntImm>> indices_or_sections, int axis);
 
 /*!
  * \brief Squeeze axes in the array.

diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
@@ -66,7 +66,7 @@ using CachedCastNodes = std::unordered_map<std::pair<const ExprNode*, DataType>,
 // Return array is of type : [MixedTypeConversionCategory (int), String, String]
 // The fields are          : [ConversionCategory, accumulation_datatype, output_datatype]
 // Call is a call node, DataType is the mixed precision type
-using FTVMMixedPrecisionConversionType = runtime::TypedPackedFunc<Array<ObjectRef>(
+using FTVMMixedPrecisionConversionType = runtime::TypedPackedFunc<Array<Variant<Integer, String>>(
     const Call& call_node, const std::string& target_dtype_str)>;
 
 /*! \brief This class transforms the given relay module into a version where
@@ -372,7 +372,7 @@ class MixedPrecisionPass : public MixedModeMutator {
       if (attr_map.count(op)) {
         // Calculate the conversion category and dtypes from registered attribute.
         FTVMMixedPrecisionConversionType func = attr_map[op];
-        Array<ObjectRef> op_descriptor =
+        Array<Variant<Integer, String>> op_descriptor =
             func(GetRef<Call>(pre_call_node), DLDataType2String(mixed_precision_type_));
         ICHECK(op_descriptor.size() == 3)
             << "got the wrong number of returned arguments (expected 3 got " << op_descriptor.size()

diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc
@@ -546,7 +546,9 @@ Call::Call(DataType dtype, RelayExpr op, Array<PrimExpr> args, Span span) {
 }
 
 TVM_REGISTER_GLOBAL("tir.Call")
-    .set_body_typed([](DataType type, RelayExpr op, Array<ObjectRef> args, Span span) {
+    .set_body_typed([](DataType type, RelayExpr op,
+                       Array<Variant<runtime::String, IterVar, BufferRegion, PrimExpr>> args,
+                       Span span) {
       Array<PrimExpr> prim_expr_args;
       for (const auto& it : args) {
         ICHECK(it->IsInstance<runtime::StringObj>() || it->IsInstance<PrimExprNode>() ||

diff --git a/src/tir/ir/specialize.cc b/src/tir/ir/specialize.cc
@@ -414,7 +414,7 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const PrimEx
 
 /**************** Implementation ****************/
 
-PrimFunc Specialize(PrimFunc func, const Map<Var, ObjectRef>& param_map) {
+PrimFunc Specialize(PrimFunc func, const Map<Var, Variant<Buffer, PrimExpr>>& param_map) {
   VarMap var_map;
   for (const auto& kv : param_map) {
     const Var& param = kv.first;

diff --git a/src/tir/transforms/inline_private_functions.cc b/src/tir/transforms/inline_private_functions.cc
@@ -231,7 +231,7 @@ class PrimFuncInliner : StmtExprMutator {
         << "Inlining of PrimFuncs with buffer arguments is not yet supported, "
         << "but callee " << gvar << " has non-empty buffer map " << callee->buffer_map;
 
-    Map<Var, ObjectRef> param_map;
+    Map<Var, Variant<tir::Buffer, tvm::PrimExpr>> param_map;
     for (size_t i = 0; i < callee->params.size(); i++) {
       param_map.Set(callee->params[i], args[i]);
     }