[High level OPT][RFC] NNVMv2 IR - Relay

CMakeLists.txt

@@ -104,6 +104,12 @@ file(GLOB COMPILER_SRCS
     src/schedule/*.cc
     )
 
+file(GLOB_RECURSE RELAY_SRCS
+    src/relay/*.cc
+    )
+list(APPEND COMPILER_SRCS ${RELAY_SRCS})
+
+
 if(NOT MSVC)
   file(GLOB COMPILER_VERILOG_SRCS src/codegen/verilog/*.cc)
   list(APPEND COMPILER_SRCS ${COMPILER_VERILOG_SRCS})

docs/conf.py

@@ -33,7 +33,7 @@
 # General information about the project.
 project = u'tvm'
 author = u'%s developers' % project
-copyright = u'2017, %s' % author
+copyright = u'2018, %s' % author
 github_doc_root = 'https://github.com/tqchen/tvm/tree/master/docs/'
 
 # add markdown parser

include/tvm/relay/base.h

@@ -0,0 +1,203 @@
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file tvm/relay/base.h
+ * \brief Base classes for the Relay IR.
+ */
+#ifndef TVM_RELAY_BASE_H_
+#define TVM_RELAY_BASE_H_
+
+#include <tvm/api_registry.h>
+#include <tvm/ir.h>
+#include <tvm/node.h>
+#include <string>
+#include <vector>
+
+namespace tvm {
+/*!
+ * \brief Relay: a high level functional IR for TVM.
+ *
+ * This namespace contains the abstract syntax tree, and other
+ * essential data structures for the Relay IR.
+ *
+ * You can find more about Relay by reading the language reference.
+ */
+namespace relay {
+/*!
+ * \brief we always used NodeRef for referencing nodes.
+ *
+ *  By default, NodeRef is a std::shared_ptr of node
+ */
+using NodeRef = tvm::NodeRef;
+
+/*!
+ * \brief Content data type.
+ */
+using DataType = ::tvm::Type;
+
+/*!
+ * \brief Symbolic expression for tensor shape.
+ */
+using ShapeExpr = ::tvm::Expr;
+
+/*!
+ * \brief Hash function for nodes.
+ * e.g. std::unordered_map<Expr, Value, NodeHash, NodeEqual>
+ */
+using NodeHash = ::tvm::NodeHash;
+/*!
+ * \brief Equality check function for nodes.
+ */
+using NodeEqual = ::tvm::NodeEqual;
+
+/*!
+ * \brief Macro to make it easy to define node ref type given node
+ * \param TypeName The name of the reference type.
+ * \param NodeName The internal container name.
+ * \param NodeRefBase The base type.
+ */
+#define RELAY_DEFINE_NODE_REF(TypeName, NodeName, NodeRefBase)            \
+  class TypeName : public NodeRefBase {                                   \
+   public:                                                                \
+    TypeName() {}                                                         \
+    explicit TypeName(std::shared_ptr<::tvm::Node> n) : NodeRefBase(n) {} \
+    const NodeName* operator->() const {                                  \
+      return static_cast<const NodeName*>(node_.get());                   \
+    }                                                                     \
+    operator bool() { return this->defined(); }                           \
+    using ContainerType = NodeName;                                       \
+  };
+
+/*!
+ * \brief The source name in the Span
+ * \sa SourceNameNode, Span
+ */
+class SourceName;
+/*!
+ * \brief The name of a source fragment.
+ */
+class SourceNameNode : public Node {
+ public:
+  /*! \brief The source name. */
+  std::string name;
+  // override attr visitor
+  void VisitAttrs(AttrVisitor* v) final { v->Visit("name", &name); }
+
+  TVM_DLL static SourceName make(std::string name);
+
+  static constexpr const char* _type_key = "relay.SourceName";
+  TVM_DECLARE_NODE_TYPE_INFO(SourceNameNode, Node);
+};
+
+/*!
+ * \brief The source name of a file span.
+ * \sa SourceNameNode, Span
+ */
+class SourceName : public NodeRef {
+ public:
+  /*! \brief default constructor  */
+  SourceName() {}
+
+  /*! \brief constructor from node pointer */
+  explicit SourceName(std::shared_ptr<Node> n) : NodeRef(n) {}
+  /*!
+   * \brief access the internal node container
+   * \return the pointer to the internal node container
+   */
+  inline const SourceNameNode* operator->() const;
+
+  /*!
+   * \brief Get an SourceName for a given operator name.
+   *  Will raise an error if the source name has not been registered.
+   * \param name Name of the operator.
+   * \return Reference to a SourceName valid throughout program lifetime.
+   */
+  TVM_DLL static const SourceName& Get(const std::string& name);
+
+  /*! \brief specify container node */
+  using ContainerType = SourceNameNode;
+};
+
+/*!
+ * \brief Span information for debugging purposes
+ */
+class Span;
+/*!
+ * \brief Stores locations in frontend source that generated a node.
+ */
+class SpanNode : public Node {
+ public:
+  /*! \brief The source name */
+  SourceName source;
+  /*! \brief Line number */
+  int lineno;
+  /*! \brief column offset */
+  int col_offset;
+  // override attr visitor
+  void VisitAttrs(AttrVisitor* v) final {
+    v->Visit("source", &source);
+    v->Visit("lineno", &lineno);
+    v->Visit("col_offset", &col_offset);
+  }
+
+  TVM_DLL static Span make(SourceName source, int lineno, int col_offset);
+
+  static constexpr const char* _type_key = "relay.Span";
+  TVM_DECLARE_NODE_TYPE_INFO(SpanNode, Node);
+};
+
+RELAY_DEFINE_NODE_REF(Span, SpanNode, NodeRef);
+
+/*!
+ * \brief This is the base node container of all relay structures.
+ */
+class RelayNode : public Node {
+ public:
+  /*! \brief The location of the program in a SourceFragment can be null,
+   * check with span.defined() */
+  mutable Span span;
+
+  static constexpr const char* _type_key = "relay.Node";
+  TVM_DECLARE_BASE_NODE_INFO(RelayNode, Node);
+};
+
+/*!
+ * \brief Get a reference type from a Node ptr type
+ *
+ *  It is always important to get a reference type
+ *  if we want to return a value as reference or keep
+ *  the node alive beyond the scope of the function.
+ *
+ * \param ptr The node pointer
+ * \tparam RefType The reference type
+ * \tparam NodeType The node type
+ * \return The corresponding RefType
+ */
+template <typename RefType, typename NodeType>
+RefType GetRef(const NodeType* ptr) {
+  static_assert(std::is_same<typename RefType::ContainerType, NodeType>::value,
+                "Can only cast to the ref of same container type");
+  return RefType(const_cast<NodeType*>(ptr)->shared_from_this());
+}
+
+// TODO(@tqchen, @jroesch): can we move these semantics to HalideIR
+template <typename T>
+inline const T* As(const NodeRef& node) {
+  const Node* ptr = static_cast<const Node*>(node.get());
+  if (ptr && (ptr->is_type<T>() || ptr->derived_from<T>())) {
+    return static_cast<const T*>(ptr);
+  }
+  return nullptr;
+}
+
+template <typename SubRef, typename BaseRef>
+SubRef Downcast(BaseRef ref) {
+  CHECK(ref->template is_type<typename SubRef::ContainerType>())
+      << "Downcast from " << ref->type_key() << " to "
+      << SubRef::ContainerType::_type_key << " failed.";
+  return SubRef(ref.node_);
+}
+
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_BASE_H_

include/tvm/relay/environment.h

@@ -0,0 +1,121 @@
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file tvm/relay/environment.h
+ * \brief The global environment: contains information needed to
+ * compile & optimize Relay programs.
+ */
+#ifndef TVM_RELAY_ENVIRONMENT_H_
+#define TVM_RELAY_ENVIRONMENT_H_
+
+#include <tvm/relay/error.h>
+#include <tvm/relay/expr.h>
+#include <tvm/relay/op.h>
+#include <tvm/relay/type.h>
+#include <string>
+#include <vector>
+
+namespace tvm {
+namespace relay {
+
+struct Environment;
+
+/*! \brief The global environment of Relay programs.
+ *
+ *  The global environment contains the global
+ *  information needed to compile a Relay program.
+ *
+ *  It contains all global functions, and configuration
+ *  options.
+ *
+ *  Many operations require access to the global
+ *  Environment. We pass the Environment by value
+ *  in a functional style as an explicit argument,
+ *  but we mutate the Environment while optimizing
+ *  Relay programs.
+ *
+ *  The functional style allows users to construct custom
+ *  environments easily, for example each thread can store
+ *  an Environment while auto-tuning.
+ * */
+
+class EnvironmentNode : public RelayNode {
+ public:
+  /*! \brief A map from ids to all global functions. */
+  tvm::Map<GlobalVar, Function> functions;
+
+  EnvironmentNode() {}
+
+  void VisitAttrs(tvm::AttrVisitor* v) final {
+    v->Visit("functions", &functions);
+    v->Visit("global_map_", &global_map_);
+  }
+
+  TVM_DLL static Environment make(tvm::Map<GlobalVar, Function> global_funcs);
+
+  /*! \brief Add a function to the global environment.
+   * \param var The name of the global function.
+   * \param func The function.
+   * \param update Controls whether you can replace a definition in the
+   * environment.
+   */
+  void Add(const GlobalVar& var, const Function& func, bool update = false);
+
+  /*! \brief Update a function in the global environment.
+   * \param var The name of the global function to update.
+   * \param func The new function.
+   */
+  void Update(const GlobalVar& var, const Function& func);
+
+  /*! \brief Remove a function from the global environment.
+   * \param var The name of the global function to update.
+   */
+  void Remove(const GlobalVar& var);
+
+  /*! \brief Lookup a global function by its variable.
+   * \param str The unique string specifying the global variable.
+   * \returns The global variable.
+   */
+  GlobalVar GetGlobalVar(const std::string& str);
+
+  /*! \brief Lookup a global function by its variable.
+   * \param var The global var to lookup.
+   * \returns The function named by the variable argument.
+   */
+  Function Lookup(const GlobalVar& var);
+
+  /*! \brief Lookup a global function by its string name
+   * \param name The name of the function.
+   * \returns The function named by the argument.
+   */
+  Function Lookup(const std::string& name);
+
+  /*! \brief Combine with another Environment.
+   * \param other The other environment.
+   */
+  void Merge(const Environment& other);
+
+  static constexpr const char* _type_key = "relay.Environment";
+  TVM_DECLARE_NODE_TYPE_INFO(EnvironmentNode, Node);
+
+ private:
+  /*! \brief A map from string names to global variables that 
+   * ensures global uniqueness. 
+   */
+  tvm::Map<std::string, GlobalVar> global_map_;
+};
+
+struct Environment : public NodeRef {
+  Environment() {}
+  explicit Environment(std::shared_ptr<tvm::Node> p) : NodeRef(p) {}
+
+  inline EnvironmentNode* operator->() const {
+    return static_cast<EnvironmentNode*>(node_.get());
+  }
+
+  using ContainerType = EnvironmentNode;
+};
+
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_ENVIRONMENT_H_

include/tvm/relay/error.h

@@ -0,0 +1,36 @@
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file error.h
+ * \brief The set of errors raised by Relay.
+ */
+#ifndef TVM_RELAY_ERROR_H_
+#define TVM_RELAY_ERROR_H_
+
+#include <string>
+#include "./base.h"
+
+namespace tvm {
+namespace relay {
+
+struct Error : dmlc::Error {
+  explicit Error(const std::string &msg) : dmlc::Error(msg) {}
+};
+
+struct InternalError : Error {
+  explicit InternalError(const std::string &msg) : Error(msg) {}
+};
+
+// TODO(@jroesch): we should change spanned errors to report
+// errors against the Environment, inverting control to error definition.
+struct FatalTypeError : dmlc::Error {
+  explicit FatalTypeError(const std::string &s) : dmlc::Error(s) {}
+};
+
+struct TypecheckerError : public dmlc::Error {
+  explicit TypecheckerError(const std::string &msg) : Error(msg) {}
+};
+
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_ERROR_H_