[Cleanup] Accept Variant<...> instead of ObjectRef when possible

Prior to the implementation of `Variant<...>` in
#15672, functions that were
polymorphic over an argument type would typically accept an
`ObjectRef` argument, then downcast to an allowed type.  This delays
the catching of an error, and can accidentally omit automatic
conversions applied by the FFI.

This commit updates several locations using this pattern to instead
accept a `Variant`, templated over the allowed types.  This enables
C++ type checking for C++ callers, standardizes the type-checking in
the FFI for non-C++ callers, and ensures that FFI type conversions are
uniformly applied.

Author: Lunderberg

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.

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);

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

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) {

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.

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()

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>() ||

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;

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]);
     }