Add C++ API for computing type key from type index

apache · Aug 11, 2021 · 1dc6bc0 · 1dc6bc0
1 parent 09b989d
commit 1dc6bc0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/include/tvm/runtime/c_runtime_api.h b/include/tvm/runtime/c_runtime_api.h
@@ -520,6 +520,14 @@ TVM_DLL int TVMObjectGetTypeIndex(TVMObjectHandle obj, unsigned* out_tindex);
  */
 TVM_DLL int TVMObjectTypeKey2Index(const char* type_key, unsigned* out_tindex);
 
+/*!
+ * \brief Convert type key to type index.
+ * \param type_key The key of the type.
+ * \param out_tindex the corresponding type index.
+ * \return 0 when success, nonzero when failure happens
+ */
+TVM_DLL int TVMObjectTypeIndex2Key(unsigned tindex, char** out_type_key);
+
 /*!
  * \brief Increase the reference count of an object.
  *

diff --git a/src/runtime/object.cc b/src/runtime/object.cc
@@ -262,3 +262,11 @@ int TVMObjectTypeKey2Index(const char* type_key, unsigned* out_tindex) {
   out_tindex[0] = tvm::runtime::ObjectInternal::ObjectTypeKey2Index(type_key);
   API_END();
 }
+
+int TVMObjectTypeIndex2Key(unsigned tindex, char** out_type_key) {
+  API_BEGIN();
+  auto key =  tvm::runtime::Object::TypeIndex2Key(tindex);
+  *out_type_key = static_cast<char*>(malloc(key.size() + 1));
+  strncpy(*out_type_key, key.c_str(), key.size());
+  API_END();
+}