[FFI][REFACTOR] Move Downcast out of ffi for now

ffi/include/tvm/ffi/any.h

@@ -635,7 +635,6 @@ struct AnyEqual {
     }
   }
 };
-
 }  // namespace ffi
 
 // Expose to the tvm namespace for usability

ffi/include/tvm/ffi/cast.h

@@ -18,21 +18,18 @@
  */
 /*!
  * \file tvm/ffi/cast.h
- * \brief Value casting support
+ * \brief Extra value casting helpers
  */
 #ifndef TVM_FFI_CAST_H_
 #define TVM_FFI_CAST_H_
 
 #include <tvm/ffi/any.h>
-#include <tvm/ffi/dtype.h>
-#include <tvm/ffi/error.h>
 #include <tvm/ffi/object.h>
 #include <tvm/ffi/optional.h>
 
-#include <utility>
-
 namespace tvm {
 namespace ffi {
+
 /*!
  * \brief Get a reference type from a raw object ptr type
  *
@@ -46,7 +43,7 @@ namespace ffi {
  * \return The corresponding RefType
  */
 template <typename RefType, typename ObjectType>
-TVM_FFI_INLINE RefType GetRef(const ObjectType* ptr) {
+inline RefType GetRef(const ObjectType* ptr) {
   static_assert(std::is_base_of_v<typename RefType::ContainerType, ObjectType>,
                 "Can only cast to the ref of same container type");
 
@@ -75,92 +72,9 @@ inline ObjectPtr<BaseType> GetObjectPtr(ObjectType* ptr) {
                 "Can only cast to the ref of same container type");
   return details::ObjectUnsafe::ObjectPtrFromUnowned<BaseType>(ptr);
 }
-
-/*!
- * \brief Downcast a base reference type to a more specific type.
- *
- * \param ref The input reference
- * \return The corresponding SubRef.
- * \tparam SubRef The target specific reference type.
- * \tparam BaseRef the current reference type.
- */
-template <typename SubRef, typename BaseRef,
-          typename = std::enable_if_t<std::is_base_of_v<ObjectRef, BaseRef>>>
-inline SubRef Downcast(BaseRef ref) {
-  if (ref.defined()) {
-    if (!ref->template IsInstance<typename SubRef::ContainerType>()) {
-      TVM_FFI_THROW(TypeError) << "Downcast from " << ref->GetTypeKey() << " to "
-                               << SubRef::ContainerType::_type_key << " failed.";
-    }
-    return SubRef(details::ObjectUnsafe::ObjectPtrFromObjectRef<Object>(std::move(ref)));
-  } else {
-    if constexpr (is_optional_type_v<SubRef> || SubRef::_type_is_nullable) {
-      return SubRef(ObjectPtr<Object>(nullptr));
-    }
-    TVM_FFI_THROW(TypeError) << "Downcast from undefined(nullptr) to `"
-                             << SubRef::ContainerType::_type_key
-                             << "` is not allowed. Use Downcast<Optional<T>> instead.";
-    TVM_FFI_UNREACHABLE();
-  }
-}
-
-/*!
- * \brief Downcast any to a specific type
- *
- * \param ref The input reference
- * \return The corresponding SubRef.
- * \tparam T The target specific reference type.
- */
-template <typename T>
-inline T Downcast(const Any& ref) {
-  if constexpr (std::is_same_v<T, Any>) {
-    return ref;
-  } else {
-    return ref.cast<T>();
-  }
-}
-
-/*!
- * \brief Downcast any to a specific type
- *
- * \param ref The input reference
- * \return The corresponding SubRef.
- * \tparam T The target specific reference type.
- */
-template <typename T>
-inline T Downcast(Any&& ref) {
-  if constexpr (std::is_same_v<T, Any>) {
-    return std::move(ref);
-  } else {
-    return std::move(ref).cast<T>();
-  }
-}
-
-/*!
- * \brief Downcast std::optional<Any> to std::optional<T>
- *
- * \param ref The input reference
- * \return The corresponding SubRef.
- * \tparam OptionalType The target optional type
- */
-template <typename OptionalType, typename = std::enable_if_t<is_optional_type_v<OptionalType>>>
-inline OptionalType Downcast(const std::optional<Any>& ref) {
-  if (ref.has_value()) {
-    if constexpr (std::is_same_v<OptionalType, Any>) {
-      return *ref;
-    } else {
-      return (*ref).cast<OptionalType>();
-    }
-  } else {
-    return OptionalType(std::nullopt);
-  }
-}
-
 }  // namespace ffi
 
-// Expose to the tvm namespace
-// Rationale: convinience and no ambiguity
-using ffi::Downcast;
+using ffi::GetObjectPtr;
 using ffi::GetRef;
 }  // namespace tvm
 #endif  // TVM_FFI_CAST_H_

ffi/tests/cpp/test_string.cc

@@ -18,7 +18,6 @@
  */
 #include <gtest/gtest.h>
 #include <tvm/ffi/any.h>
-#include <tvm/ffi/cast.h>
 #include <tvm/ffi/string.h>
 
 namespace {
@@ -266,7 +265,7 @@ TEST(String, Cast) {
   string source = "this is a string";
   String s{source};
   Any r = s;
-  String s2 = Downcast<String>(r);
+  String s2 = r.cast<String>();
 }
 
 TEST(String, Concat) {

include/tvm/node/cast.h

@@ -0,0 +1,117 @@
+/*
+ * 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 tvm/node/cast.h
+ * \brief Value casting helpers
+ */
+#ifndef TVM_NODE_CAST_H_
+#define TVM_NODE_CAST_H_
+
+#include <tvm/ffi/any.h>
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/object.h>
+#include <tvm/ffi/optional.h>
+
+#include <utility>
+
+namespace tvm {
+
+/*!
+ * \brief Downcast a base reference type to a more specific type.
+ *
+ * \param ref The input reference
+ * \return The corresponding SubRef.
+ * \tparam SubRef The target specific reference type.
+ * \tparam BaseRef the current reference type.
+ */
+template <typename SubRef, typename BaseRef,
+          typename = std::enable_if_t<std::is_base_of_v<ffi::ObjectRef, BaseRef>>>
+inline SubRef Downcast(BaseRef ref) {
+  if (ref.defined()) {
+    if (!ref->template IsInstance<typename SubRef::ContainerType>()) {
+      TVM_FFI_THROW(TypeError) << "Downcast from " << ref->GetTypeKey() << " to "
+                               << SubRef::ContainerType::_type_key << " failed.";
+    }
+    return SubRef(ffi::details::ObjectUnsafe::ObjectPtrFromObjectRef<ffi::Object>(std::move(ref)));
+  } else {
+    if constexpr (ffi::is_optional_type_v<SubRef> || SubRef::_type_is_nullable) {
+      return SubRef(ffi::ObjectPtr<ffi::Object>(nullptr));
+    }
+    TVM_FFI_THROW(TypeError) << "Downcast from undefined(nullptr) to `"
+                             << SubRef::ContainerType::_type_key
+                             << "` is not allowed. Use Downcast<Optional<T>> instead.";
+    TVM_FFI_UNREACHABLE();
+  }
+}
+
+/*!
+ * \brief Downcast any to a specific type
+ *
+ * \param ref The input reference
+ * \return The corresponding SubRef.
+ * \tparam T The target specific reference type.
+ */
+template <typename T>
+inline T Downcast(const ffi::Any& ref) {
+  if constexpr (std::is_same_v<T, Any>) {
+    return ref;
+  } else {
+    return ref.cast<T>();
+  }
+}
+
+/*!
+ * \brief Downcast any to a specific type
+ *
+ * \param ref The input reference
+ * \return The corresponding SubRef.
+ * \tparam T The target specific reference type.
+ */
+template <typename T>
+inline T Downcast(ffi::Any&& ref) {
+  if constexpr (std::is_same_v<T, Any>) {
+    return std::move(ref);
+  } else {
+    return std::move(ref).cast<T>();
+  }
+}
+
+/*!
+ * \brief Downcast std::optional<Any> to std::optional<T>
+ *
+ * \param ref The input reference
+ * \return The corresponding SubRef.
+ * \tparam OptionalType The target optional type
+ */
+template <typename OptionalType, typename = std::enable_if_t<ffi::is_optional_type_v<OptionalType>>>
+inline OptionalType Downcast(const std::optional<ffi::Any>& ref) {
+  if (ref.has_value()) {
+    if constexpr (std::is_same_v<OptionalType, ffi::Any>) {
+      return *ref;
+    } else {
+      return (*ref).cast<OptionalType>();
+    }
+  } else {
+    return OptionalType(std::nullopt);
+  }
+}
+}  // namespace tvm
+#endif  // TVM_NODE_CAST_H_

include/tvm/node/node.h

@@ -35,6 +35,7 @@
 #define TVM_NODE_NODE_H_
 
 #include <tvm/ffi/memory.h>
+#include <tvm/node/cast.h>
 #include <tvm/node/repr_printer.h>
 #include <tvm/node/structural_equal.h>
 #include <tvm/node/structural_hash.h>
@@ -57,8 +58,6 @@ using ffi::ObjectPtrHash;
 using ffi::ObjectRef;
 using ffi::PackedArgs;
 using ffi::TypeIndex;
-using runtime::Downcast;
-using runtime::GetRef;
 
 }  // namespace tvm
 #endif  // TVM_NODE_NODE_H_

include/tvm/runtime/disco/session.h

@@ -124,6 +124,8 @@ inline std::string DiscoAction2String(DiscoAction action) {
   LOG(FATAL) << "ValueError: Unknown DiscoAction: " << static_cast<int>(action);
 }
 
+class SessionObj;
+
 /*!
  * \brief An object that exists on all workers.
  *
@@ -156,6 +158,9 @@ class DRefObj : public Object {
   int64_t reg_id;
   /*! \brief Back-pointer to the host controler session */
   ObjectRef session{nullptr};
+
+ private:
+  inline SessionObj* GetSession();
 };
 
 /*!
@@ -321,18 +326,22 @@ class WorkerZeroData {
 
 // Implementation details
 
+inline SessionObj* DRefObj::GetSession() {
+  return const_cast<SessionObj*>(static_cast<const SessionObj*>(session.get()));
+}
+
 DRefObj::~DRefObj() {
   if (this->session.defined()) {
-    Downcast<Session>(this->session)->DeallocReg(reg_id);
+    GetSession()->DeallocReg(reg_id);
   }
 }
 
 ffi::Any DRefObj::DebugGetFromRemote(int worker_id) {
-  return Downcast<Session>(this->session)->DebugGetFromRemote(this->reg_id, worker_id);
+  return GetSession()->DebugGetFromRemote(this->reg_id, worker_id);
 }
 
 void DRefObj::DebugCopyFrom(int worker_id, ffi::AnyView value) {
-  return Downcast<Session>(this->session)->DebugSetRegister(this->reg_id, value, worker_id);
+  return GetSession()->DebugSetRegister(this->reg_id, value, worker_id);
 }
 
 template <typename... Args>

include/tvm/runtime/object.h

@@ -39,7 +39,6 @@ using tvm::ffi::ObjectPtrEqual;
 using tvm::ffi::ObjectPtrHash;
 using tvm::ffi::ObjectRef;
 
-using tvm::ffi::Downcast;
 using tvm::ffi::GetObjectPtr;
 using tvm::ffi::GetRef;
 

include/tvm/runtime/vm/vm.h

@@ -189,10 +189,12 @@ class VirtualMachine : public runtime::ModuleNode {
     using ContainerType = typename T::ContainerType;
     uint32_t key = ContainerType::RuntimeTypeIndex();
     if (auto it = extensions.find(key); it != extensions.end()) {
-      return Downcast<T>((*it).second);
+      ffi::Any value = (*it).second;
+      return value.cast<T>();
     }
     auto [it, _] = extensions.emplace(key, T::Create());
-    return Downcast<T>((*it).second);
+    ffi::Any value = (*it).second;
+    return value.cast<T>();
   }
 
   /*!
@@ -224,7 +226,7 @@ class VirtualMachine : public runtime::ModuleNode {
   std::vector<Device> devices;
   /*! \brief The VM extensions. Mapping from the type index of the extension to the extension
    * instance. */
-  std::unordered_map<uint32_t, VMExtension> extensions;
+  std::unordered_map<uint32_t, Any> extensions;
 };
 
 }  // namespace vm

src/node/container_printing.cc

@@ -22,6 +22,7 @@
  * \file node/container_printint.cc
  */
 #include <tvm/ffi/function.h>
+#include <tvm/node/cast.h>
 #include <tvm/node/functor.h>
 #include <tvm/node/repr_printer.h>
 
@@ -62,6 +63,6 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
     .set_dispatch<ffi::ShapeObj>([](const ObjectRef& node, ReprPrinter* p) {
-      p->stream << ffi::Downcast<ffi::Shape>(node);
+      p->stream << Downcast<ffi::Shape>(node);
     });
 }  // namespace tvm

src/node/repr_printer.cc

@@ -23,6 +23,7 @@
  */
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/node/cast.h>
 #include <tvm/node/repr_printer.h>
 #include <tvm/runtime/device_api.h>