[API/Refactor] Unified PackedFunc for API and Generated Functions

.travis.yml

@@ -4,7 +4,7 @@ language: cpp
 
 os:
  - linux
- - osx
+ # - osx
 
 env:
  # code analysis

include/tvm/api_registry.h

@@ -0,0 +1,85 @@
+/*!
+ * Copyright (c) 2016 by Contributors
+ * \file api_registry.h
+ * \brief This file defines the TVM API registry.
+ *
+ * The API registry stores type-erased functions.
+ * Each registered function is automatically exposed
+ * to front-end language(e.g. python).
+ * Front-end can also pass callbacks as PackedFunc, or register
+ * then into the same global registry in C++.
+ * The goal is to mix the front-end language and the TVM back-end.
+ *
+ * \code
+ * // register the function as MyAPIFuncName
+ * TVM_REGISTER_API(MyAPIFuncName)
+ * .set_body([](TVMArgs args, TVMRetValue* rv) {
+ * // my code.
+ * });
+ * \endcode
+ */
+#ifndef TVM_API_REGISTRY_H_
+#define TVM_API_REGISTRY_H_
+
+#include <dmlc/base.h>
+#include <string>
+#include "./base.h"
+#include "./runtime/packed_func.h"
+#include "./packed_func_ext.h"
+
+namespace tvm {
+
+/*! \brief Utility to register API. */
+class APIRegistry {
+ public:
+ /*!
+ * \brief set the body of the function to be f
+ * \param f The body of the function.
+ */
+ APIRegistry& set_body(PackedFunc f); // NOLINT(*)
+ /*!
+ * \brief set the body of the function to be f
+ * \param f The body of the function.
+ */
+ APIRegistry& set_body(PackedFunc::FType f) { // NOLINT(*)
+ return set_body(PackedFunc(f));
+ }
+ /*!
+ * \brief Register a function with given name
+ * \param name The name of the function.
+ */
+ static APIRegistry& __REGISTER__(const std::string& name); // NOLINT(*)
+
+ private:
+ /*! \brief name of the function */
+ std::string name_;
+};
+
+/*!
+ * \brief Get API function by name.
+ *
+ * \param name The name of the function.
+ * \return the corresponding API function.
+ * \note It is really PackedFunc::GetGlobal under the hood.
+ */
+inline PackedFunc GetAPIFunc(const std::string& name) {
+ return PackedFunc::GetGlobal(name);
+}
+
+#define _TVM_REGISTER_VAR_DEF_ \
+ static DMLC_ATTRIBUTE_UNUSED ::tvm::APIRegistry& __make_TVMRegistry_
+
+/*!
+ * \brief Register API function globally.
+ * \code
+ * TVM_REGISTER_API(MyPrint)
+ * .set_body([](TVMArgs args, TVMRetValue* rv) {
+ * // my code.
+ * });
+ * \endcode
+ */
+#define TVM_REGISTER_API(OpName) \
+ DMLC_STR_CONCAT(_TVM_REGISTER_VAR_DEF_, __COUNTER__) = \
+ ::tvm::APIRegistry::__REGISTER__(#OpName)
+} // namespace tvm
+#endif // TVM_API_REGISTRY_H_

include/tvm/c_api.h

@@ -2,83 +2,23 @@
  * Copyright (c) 2016 by Contributors
  * \file c_api.h
  * \brief C API of TVM DSL
+ *
+ * \note The API is designed in a minimum way.
+ * Most of the API functions are registered and can be pulled out.
+ *
+ * The common flow is:
+ * - Use TVMFuncListGlobalNames to get global function name
+ * - Use TVMFuncCall to call these functions.
  */
 #ifndef TVM_C_API_H_
 #define TVM_C_API_H_
 
 #include "./runtime/c_runtime_api.h"
 
 TVM_EXTERN_C {
-/*! \brief handle to functions */
-typedef void* APIFuncHandle;
 /*! \brief handle to node */
 typedef void* NodeHandle;
 
-/*!
- * \brief List all the node function name
- * \param out_size The number of functions
- * \param out_array The array of function names.
- * \return 0 when success, -1 when failure happens
- */
-TVM_DLL int TVMListAPIFuncNames(int *out_size,
- const char*** out_array);
-/*!
- * \brief get function handle by name
- * \param name The name of function
- * \param handle The returning function handle
- * \return 0 when success, -1 when failure happens
- */
-TVM_DLL int TVMGetAPIFuncHandle(const char* name,
- APIFuncHandle *handle);
-
-/*!
- * \brief Get the detailed information about function.
- * \param handle The operator handle.
- * \param real_name The returned name of the function.
- * This name is not the alias name of the atomic symbol.
- * \param description The returned description of the symbol.
- * \param num_doc_args Number of arguments that contain documents.
- * \param arg_names Name of the arguments of doc args
- * \param arg_type_infos Type informations about the arguments.
- * \param arg_descriptions Description information about the arguments.
- * \param return_type Return type of the function, if any.
- * \return 0 when success, -1 when failure happens
- */
-TVM_DLL int TVMGetAPIFuncInfo(APIFuncHandle handle,
- const char **real_name,
- const char **description,
- int *num_doc_args,
- const char ***arg_names,
- const char ***arg_type_infos,
- const char ***arg_descriptions,
- const char **return_type);
-
-/*!
- * \brief Push an argument to the function calling stack.
- * If push fails, the stack will be reset to empty
- *
- * \param arg The argument
- * \param type_code The type_code of argument as in TVMTypeCode
- * \return 0 when success, -1 when failure happens
- * \note API calls always exchanges with type bits=64, lanes=1
- */
-TVM_DLL int TVMAPIPushStack(TVMValue arg,
- int type_code);
-
-/*!
- * \brief call a function by using arguments in the stack.
- * The stack will be cleanup to empty after this call, whether the call is successful.
- *
- * \param handle The function handle
- * \param ret_val The return value.
- * \param ret_type_code the type code of return value.
- * \return 0 when success, -1 when failure happens
- * \note API calls always exchanges with type bits=64, lanes=1
- */
-TVM_DLL int TVMAPIFuncCall(APIFuncHandle handle,
- TVMValue* ret_val,
- int* ret_type_code);
-
 /*!
  * \brief free the node handle
  * \param handle The node handle to be freed.

include/tvm/expr.h

@@ -12,6 +12,7 @@
 #include <string>
 #include <algorithm>
 #include "./base.h"
+#include "./runtime/packed_func.h"
 
 namespace tvm {
 

include/tvm/packed_func_ext.h

@@ -0,0 +1,196 @@
+/*!
+ * Copyright (c) 2016 by Contributors
+ * \file packed_func_ext.h
+ * \brief Extension package to PackedFunc
+ * This enales pass NodeRef types into/from PackedFunc.
+ */
+#ifndef TVM_PACKED_FUNC_EXT_H_
+#define TVM_PACKED_FUNC_EXT_H_
+
+#include <sstream>
+#include <string>
+#include <memory>
+#include <type_traits>
+
+#include "./base.h"
+#include "./expr.h"
+
+namespace tvm {
+using runtime::TVMArgs;
+using runtime::TVMRetValue;
+using runtime::PackedFunc;
+
+namespace runtime {
+/*!
+ * \brief Runtime type checker for node type.
+ * \tparam T the type to be checked.
+ */
+template<typename T>
+struct NodeTypeChecker {
+ static inline bool Check(Node* sptr) {
+ // This is the only place in the project where RTTI is used
+ // It can be turned off, but will make non strict checking.
+ // TODO(tqchen) possibly find alternative to turn of RTTI
+ using ContainerType = typename T::ContainerType;
+ return (dynamic_cast<ContainerType*>(sptr) != nullptr);
+ }
+ static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
+ using ContainerType = typename T::ContainerType;
+ os << ContainerType::_type_key;
+ }
+};
+
+template<typename T>
+struct NodeTypeChecker<Array<T> > {
+ static inline bool Check(Node* sptr) {
+ if (sptr == nullptr) return false;
+ if (!sptr->is_type<ArrayNode>()) return false;
+ ArrayNode* n = static_cast<ArrayNode*>(sptr);
+ for (const auto& p : n->data) {
+ if (!NodeTypeChecker<T>::Check(p.get())) return false;
+ }
+ return true;
+ }
+ static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
+ os << "array<";
+ NodeTypeChecker<T>::PrintName(os);
+ os << ">";
+ }
+};
+
+template<typename K, typename V>
+struct NodeTypeChecker<Map<K, V> > {
+ static inline bool Check(Node* sptr) {
+ if (sptr == nullptr) return false;
+ if (!sptr->is_type<MapNode>()) return false;
+ MapNode* n = static_cast<MapNode*>(sptr);
+ for (const auto& kv : n->data) {
+ if (!NodeTypeChecker<K>::Check(kv.first.get())) return false;
+ if (!NodeTypeChecker<V>::Check(kv.second.get())) return false;
+ }
+ return true;
+ }
+ static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
+ os << "map<";
+ NodeTypeChecker<K>::PrintName(os);
+ os << ',';
+ NodeTypeChecker<V>::PrintName(os);
+ os << '>';
+ }
+};
+
+template<typename T>
+inline std::string NodeTypeName() {
+ std::ostringstream os;
+ NodeTypeChecker<T>::PrintName(os);
+ return os.str();
+}
+
+// extensions for tvm arg value
+
+template<typename TNodeRef, typename>
+inline TVMArgValue::operator TNodeRef() const {
+ static_assert(
+ std::is_base_of<NodeRef, TNodeRef>::value,
+ "Conversion only works for NodeRef");
+ if (type_code_ == kNull) return TNodeRef();
+ TVM_CHECK_TYPE_CODE(type_code_, kNodeHandle);
+ std::shared_ptr<Node>& sptr = *ptr<std::shared_ptr<Node> >();
+ CHECK(NodeTypeChecker<TNodeRef>::Check(sptr.get()))
+ << "Expected type " << NodeTypeName<TNodeRef>()
+ << " but get " << sptr->type_key();
+ return TNodeRef(sptr);
+}
+
+inline TVMArgValue::operator Halide::Expr() const {
+ if (type_code_ == kNull) return Expr();
+ if (type_code_ == kInt) {
+ return Expr(static_cast<int>(value_.v_int64));
+ }
+ if (type_code_ == kFloat) {
+ return Expr(static_cast<float>(value_.v_float64));
+ }
+ TVM_CHECK_TYPE_CODE(type_code_, kNodeHandle);
+ std::shared_ptr<Node>& sptr = *ptr<std::shared_ptr<Node> >();
+ if (sptr->is_type<IterVarNode>()) {
+ return IterVar(sptr)->var;
+ }
+ CHECK(NodeTypeChecker<Expr>::Check(sptr.get()))
+ << "Expected type " << NodeTypeName<Expr>()
+ << " but get " << sptr->type_key();
+ return Expr(sptr);
+}
+
+inline std::shared_ptr<Node>& TVMArgValue::node_sptr() {
+ TVM_CHECK_TYPE_CODE(type_code_, kNodeHandle);
+ return *ptr<std::shared_ptr<Node> >();
+}
+
+
+template<typename TNodeRef, typename>
+inline bool TVMArgValue::IsNodeType() const {
+ TVM_CHECK_TYPE_CODE(type_code_, kNodeHandle);
+ std::shared_ptr<Node>& sptr =
+ *ptr<std::shared_ptr<Node> >();
+ return NodeTypeChecker<TNodeRef>::Check(sptr.get());
+}
+
+// extensions for TVMRetValue
+inline TVMRetValue& TVMRetValue::operator=(
+ const std::shared_ptr<Node>& other) {
+ SwitchToClass<std::shared_ptr<Node> >(kNodeHandle, other);
+ return *this;
+}
+
+inline TVMRetValue& TVMRetValue::operator=(const NodeRef& other) {
+ SwitchToClass<std::shared_ptr<Node> >(kNodeHandle, other.node_);
+ return *this;
+}
+
+template<typename TNodeRef, typename>
+inline TVMRetValue::operator TNodeRef() const {
+ static_assert(
+ std::is_base_of<NodeRef, TNodeRef>::value,
+ "Conversion only works for NodeRef");
+ if (type_code_ == kNull) return TNodeRef();
+ TVM_CHECK_TYPE_CODE(type_code_, kNodeHandle);
+ return TNodeRef(*ptr<std::shared_ptr<Node> >());
+}
+
+inline void TVMArgsSetter::operator()(size_t i, NodeRef& other) const { // NOLINT(*)
+ values_[i].v_handle = &(other.node_);
+ type_codes_[i] = kNodeHandle;
+}
+
+// Type related stuffs
+inline Type TVMType2Type(TVMType t) {
+ return Type(static_cast<halide_type_code_t>(t.code), t.bits, t.lanes);
+}
+
+inline TVMType Type2TVMType(Type t) {
+ TVMType ret;
+ ret.code = static_cast<uint8_t>(t.code());
+ ret.bits = static_cast<uint8_t>(t.bits());
+ ret.lanes = static_cast<uint16_t>(t.lanes());
+ return ret;
+}
+
+inline TVMRetValue& TVMRetValue::operator=(const Halide::Type& t) {
+ return this->operator=(Type2TVMType(t));
+}
+
+inline TVMRetValue::operator Halide::Type() const {
+ return TVMType2Type(operator TVMType());
+}
+
+inline TVMArgValue::operator Halide::Type() const {
+ return TVMType2Type(operator TVMType());
+}
+
+inline void TVMArgsSetter::operator()(
+ size_t i, const Halide::Type& t) const {
+ this->operator()(i, Type2TVMType(t));
+}
+} // namespace runtime
+} // namespace tvm
+#endif // TVM_PACKED_FUNC_EXT_H_