[API/JIT] Enable registerable global function, introduce StackVM intepreter

Author: tqchen

include/tvm/codegen.h

@@ -10,7 +10,7 @@
 #include "./base.h"
 #include "./expr.h"
 #include "./module.h"
-#include "./runtime/runtime.h"
+#include "./runtime/packed_func.h"
 
 
 namespace tvm {

include/tvm/ir.h

@@ -81,11 +81,13 @@ constexpr const char* tvm_handle_is_null = "tvm_handle_is_null";
 /*!
  * \brief See pesudo code
  *
- *  bool tvm_print(VType value) {
- *     LOG(INFO) << value;
+ *  int tvm_call_global(name, TVMValue* args) {
+ *     PackedFunc f = PackedFunc::GetGlobal(name);
+ *     f (args, type_code_of(args), len(args));
+ *     return 0;
  *  }
  */
-constexpr const char* tvm_print = "tvm_print";
+constexpr const char* tvm_call_global = "tvm_call_global";
 
 /*! \brief The field id of each field in array */
 enum TVMArrayFieldKind {

include/tvm/module.h

@@ -20,9 +20,6 @@ namespace tvm {
 // Internal node container of lowered function.
 class LoweredFuncNode;
 
-// Internal node container of module.
-class ModuleNode;
-
 /*!
  * \brief LoweredFunc represents function after lowering.
  *  This is the final IR representation before codegen.

include/tvm/runtime/c_runtime_api.h

@@ -161,7 +161,7 @@ TVM_DLL const char *TVMGetLastError(void);
  * \param option_vals Additional option values to pass
  * \param num_options Number of options to be passed into it.
  * \param out_code 1: success, 0: already initialized
- * \return Whether the function is successful.
+ * \return 0 when success, -1 when failure happens
  */
 TVM_DLL int TVMDeviceInit(int dev_mask,
                           const char** option_keys,
@@ -188,7 +188,7 @@ TVM_DLL int TVMContextEnabled(TVMContext ctx,
  * \param dtype The array data type.
  * \param ctx The ctx this array sits on.
  * \param out The output handle.
- * \return Whether the function is successful.
+ * \return 0 when success, -1 when failure happens
  */
 TVM_DLL int TVMArrayAlloc(const tvm_index_t* shape,
                           tvm_index_t ndim,
@@ -198,6 +198,7 @@ TVM_DLL int TVMArrayAlloc(const tvm_index_t* shape,
 /*!
  * \brief Free the TVM Array.
  * \param handle The array handle to be freed.
+ * \return 0 when success, -1 when failure happens
  */
 TVM_DLL int TVMArrayFree(TVMArrayHandle handle);
 
@@ -206,6 +207,7 @@ TVM_DLL int TVMArrayFree(TVMArrayHandle handle);
  * \param from The array to be copied from.
  * \param to The target space.
  * \param stream The stream where the copy happens, can be NULL.
+ * \return 0 when success, -1 when failure happens
  */
 TVM_DLL int TVMArrayCopyFromTo(TVMArrayHandle from,
                                TVMArrayHandle to,
@@ -214,13 +216,14 @@ TVM_DLL int TVMArrayCopyFromTo(TVMArrayHandle from,
  * \brief Wait until all computations on stream completes.
  * \param ctx The ctx to be synchronized.
  * \param stream The stream to be synchronized.
+ * \return 0 when success, -1 when failure happens
  */
 TVM_DLL int TVMSynchronize(TVMContext ctx, TVMStreamHandle stream);
 
 /*!
  * \brief Free the function when it is no longer needed.
  * \param func The function handle
- * \return whether
+ * \return 0 when success, -1 when failure happens
  */
 TVM_DLL int TVMFuncFree(TVMFunctionHandle func);
 
@@ -239,6 +242,57 @@ TVM_DLL int TVMFuncCall(TVMFunctionHandle func,
                         TVMValue* args,
                         int* type_codes,
                         int num_args);
+
+/*!
+ * \brief C type of packed function.
+ *
+ * \param args The arguments
+ * \param type_codes The type codes of the arguments
+ * \param num_args Number of arguments.
+ * \param resource_handle The handle additional resouce handle from fron-end.
+ */
+typedef void (*TVMPackedCFunc)(
+    TVMValue* args, int* type_codes, int num_args, void* resource_handle);
+
+/*!
+ * \brief C callback to free the resource handle in C packed function.
+ * \param resource_handle The handle additional resouce handle from fron-end.
+ */
+typedef void (*TVMPackedCFuncFinalizer)(void* resource_handle);
+
+/*!
+ * \brief Wrap a TVMPackedCFunc to become a FunctionHandle.
+ *
+ * The resource_handle will be managed by TVM API, until the function is no longer used.
+ *
+ * \param func The packed C function.
+ * \param resource_handle The resource handle from front-end, can be NULL.
+ * \param fin The finalizer on resource handle when the FunctionHandle get freed, can be NULL
+ * \param out the result function handle.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
+                                   void* resource_handle,
+                                   TVMPackedCFuncFinalizer fin,
+                                   TVMFunctionHandle *out);
+
+/*!
+ * \brief Register the function to runtime's global table.
+ *
+ * The registered function then can be pulled by the backend by the name.
+ *
+ * \param name The name of the function.
+ * \param f The function to be registered.
+ */
+TVM_DLL int TVMFuncRegisterGlobal(const char* name, TVMFunctionHandle f);
+
+/*!
+ * \brief Get a global function.
+ *
+ * \param name The name of the function.
+ * \param out the result function pointer.
+ */
+TVM_DLL int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out);
 }  // TVM_EXTERN_C
 
 #endif  // TVM_RUNTIME_C_RUNTIME_API_H_

include/tvm/runtime/runtime.h renamed to include/tvm/runtime/packed_func.h

@@ -1,35 +1,43 @@
 /*!
  *  Copyright (c) 2016 by Contributors
- * \file runtime.h
+ * \file packed_func.h
  * \brief Runtime related c++ class.
  */
-#ifndef TVM_RUNTIME_RUNTIME_H_
-#define TVM_RUNTIME_RUNTIME_H_
+#ifndef TVM_RUNTIME_PACKED_FUNC_H_
+#define TVM_RUNTIME_PACKED_FUNC_H_
 
 #include <functional>
 #include <tuple>
+#include <vector>
+#include <string>
 #include "./c_runtime_api.h"
 
 namespace tvm {
 namespace runtime {
 
 /*!
- * \brief Packed function is a runtime function
- *  whose argument type_codes are erased by packed format.
+ * \brief Packed function is a type-erased function.
+ *  The arguments are passed by packed format.
  *
- *  This is an useful unified interface to call generated functions.
+ *  This is an useful unified interface to call generated functions,
+ *  It is the unified function function type of TVM.
+ *  It corresponds to TVMFunctionHandle in C runtime API.
  */
 class PackedFunc {
  public:
   /*! \brief The internal std::function */
   using FType = std::function<void(const TVMValue* args, const int* type_codes, int num_args)>;
+  /*! \brief default constructor */
   PackedFunc() {}
+  /*!
+   * \brief constructing a packed function from a std::function.
+   * \param body the internal container of packed function.
+   */
   explicit PackedFunc(FType body) : body_(body) {}
   /*!
-   * \brief invoke the packed function by directly passing in arguments.
+   * \brief Call packed function by directly passing in unpacked format.
    * \param args Arguments to be passed.
    * \tparam Args arguments to be passed.
-   * \return The first return value.
    */
   template<typename... Args>
   inline void operator()(Args&& ...args) const;
@@ -41,9 +49,25 @@ class PackedFunc {
    */
   inline void CallPacked(const TVMValue* args, const int* type_codes, int num_args) const;
   /*! \return the internal body function */
-  inline FType body() const {
-    return body_;
-  }
+  inline FType body() const;
+  /*!
+   * \brief Register f as into global function table
+   * \param name The name of the function.
+   * \param f The function to be registered.
+   * \return Reference to the registered function.
+   * \note The returned reference is valid until the end of the program
+   */
+  static const PackedFunc& RegisterGlobal(const std::string& name, PackedFunc f);
+  /*!
+   * \brief Get the global function by name.
+   * \param name The name of the function.
+   * \return reference to the registered function.
+   */
+  static const PackedFunc& GetGlobal(const std::string& name);
+  /*!
+   * \brief Get the names of currently registered global function.
+   */
+  static std::vector<std::string> ListGlobalNames();
 
  private:
   /*! \brief internal container of packed function */
@@ -56,6 +80,10 @@ inline void PackedFunc::CallPacked(
   body_(args, type_codes, num_args);
 }
 
+inline PackedFunc::FType PackedFunc::body() const {
+  return body_;
+}
+
 template<bool stop, std::size_t I, typename F, typename ...Args>
 struct for_each_dispatcher_ {
   static inline void run(const std::tuple<Args...>& args, F f) {
@@ -124,4 +152,4 @@ inline void PackedFunc::operator()(Args&& ...args) const {
 }
 }  // namespace runtime
 }  // namespace tvm
-#endif  // TVM_RUNTIME_RUNTIME_H_
+#endif  // TVM_RUNTIME_PACKED_FUNC_H_