[WIP][TVM] Bring Your Own Codegen to TVM

CMakeLists.txt

@@ -252,6 +252,7 @@ include(cmake/modules/LLVM.cmake)
 include(cmake/modules/Micro.cmake)
 include(cmake/modules/ANTLR.cmake)
 include(cmake/modules/contrib/BLAS.cmake)
+include(cmake/modules/contrib/Extern.cmake)
 include(cmake/modules/contrib/Random.cmake)
 include(cmake/modules/contrib/MicroStandaloneRuntime.cmake)
 include(cmake/modules/contrib/Sort.cmake)

cmake/config.cmake

@@ -163,6 +163,11 @@ set(USE_ROCBLAS OFF)
 # Whether use contrib sort
 set(USE_SORT ON)
 
+# Whether use contrib extern (use ";" to separate multiple externs)
+# Available externs:
+#   dnnl
+set(USE_EXTERN none)
+
 # Build ANTLR parser for Relay text format
 # Possible values:
 # - ON: enable ANTLR by searching default locations (cmake find_program for antlr4 and /usr/local for jar)

cmake/modules/contrib/Extern.cmake

@@ -0,0 +1,34 @@
+# 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.
+
+message(STATUS "Build with relay.backend.contrib")
+
+file(GLOB GCC_RELAY_CONTRIB_SRC src/relay/backend/contrib/gcc/codegen.cc)
+list(APPEND COMPILER_SRCS ${GCC_RELAY_CONTRIB_SRC})
+
+list(FIND USE_EXTERN "dnnl" DNNL_IDX)
+if(DNNL_IDX GREATER -1)
+  file(GLOB DNNL_RELAY_CONTRIB_SRC src/relay/backend/contrib/dnnl/codegen.cc)
+  list(APPEND COMPILER_SRCS ${DNNL_RELAY_CONTRIB_SRC})
+
+  find_library(EXTERN_LIBRARY_DNNL dnnl)
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${EXTERN_LIBRARY_DNNL})
+  file(GLOB DNNL_CONTRIB_SRC src/runtime/contrib/dnnl/*)
+  list(APPEND RUNTIME_SRCS ${DNNL_CONTRIB_SRC})
+  message(STATUS "Use extern library: MKLDNN" ${EXTERN_LIBRARY_DNNL})
+endif()
+

include/tvm/relay/attrs/annotation.h

@@ -57,6 +57,19 @@ struct CastHintAttrs : public tvm::AttrsNode<CastHintAttrs> {
   }
 };
 
+/*!
+ * \brief Options for the subgraph operators.
+ */
+struct SubgraphAttrs : public tvm::AttrsNode<SubgraphAttrs> {
+  /*! \brief The 3rd party compiler for subgraph code generation. */
+  std::string compiler;
+
+  TVM_DECLARE_ATTRS(SubgraphAttrs, "relay.attrs.SubgraphAttrs") {
+    TVM_ATTR_FIELD(compiler)
+      .describe("The 3rd compiler used for subgraph code generation.");
+  }
+};
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_ANNOTATION_H_

include/tvm/relay/expr.h

@@ -268,6 +268,14 @@ class FunctionNode : public ExprNode {
    */
   bool IsPrimitive() const;
 
+  /*!
+   * \brief Check whether the function is an external function.
+   * External functions are subgraphes that supported by external libraries.
+   *
+   * \return Whether the function is external or not.
+   */
+  bool IsExternal() const;
+
   TVM_DLL static Function make(tvm::Array<Var> params,
                                Expr body,
                                Type ret_type,
@@ -588,6 +596,25 @@ std::string AsText(const NodeRef& node,
                    bool show_meta_data = true,
                    runtime::TypedPackedFunc<std::string(Expr)> annotate = nullptr);
 
+/*! \brief namespace of the attributes that are attached to a function. */
+namespace attr {
+/*! \brief Mark the function as a primitive function. */
+constexpr const char* kPrimitive = "Primitive";
+/*!
+ * \brief Mark the function as an external function that needs to be handled by
+ * the external codegen tool/backend.
+ */
+constexpr const char* kExternal = "External";
+/*! \brief Indicate if the function is a closure. */
+constexpr const char* kClosure = "Closure";
+/*! \brief Store a Var to parameter/Constant mapping on a Function. */
+constexpr const char* kParams = "__params__";
+/*! \brief Store the function name. */
+constexpr const char* kFuncName = "FuncName";
+/*! \brief Mark if the function should be avoided being optimized. */
+constexpr const char* kSkipOptimization = "SkipOptimization";
+}  // namespace attr
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_EXPR_H_

include/tvm/relay/op_attr_types.h

@@ -29,6 +29,7 @@
 #include <tvm/build_module.h>
 #include <tvm/relay/type.h>
 #include <tvm/relay/expr.h>
+#include <string>
 
 namespace tvm {
 namespace relay {
@@ -122,7 +123,7 @@ using FTVMSchedule = runtime::TypedPackedFunc<
  *  operator with other expressions. This function will be invoked
  *  in AlterOpLayout pass.
  * \param attrs The attribute of the original node.
- * \param inputs The input symbols of the original node.
+ * \param args The input symbols of the original node.
  * \param tinfos An array of placeholders, use for getting the inferred shape
  *               and dtype of the inputs.
  * \return new_expr The modified expression.
@@ -136,8 +137,8 @@ using FTVMAlterOpLayout = runtime::TypedPackedFunc<
  * \brief Legalizes an expression with another expression. This function will be
  *  invoked in Legalize pass. It is a target-dependent pass.
  * \param attrs The attribute of the original node.
- * \param inputs The input symbols of the original node.
- * \param tinfos An array of placeholders, use for getting the inferred shape
+ * \param args The input symbols of the original node.
+ * \param arg_types An array of placeholders, use for getting the inferred shape
  *               and dtype of the inputs.
  * \return new_expr The modified expression.
  */
@@ -146,6 +147,22 @@ using FTVMLegalize = runtime::TypedPackedFunc<
        const Array<Expr>& args,
        const Array<tvm::relay::Type>& arg_types)>;
 
+/*!
+ * \brief Annotates an expression to indicate which external codegen tool an op
+ * should be scheduled to. It is a hardware dependent pass.
+ *
+ * \param attrs The attribute of the original expr.
+ * \param args The arguments of the original expr.
+ * \param compiler The external compiler that is used for external ops.
+ *
+ * \return true if this op should be registered with external codegen tool,
+ * otherwise, false.
+ */
+using FTVMExternOp = runtime::TypedPackedFunc<
+  bool(const Attrs& attrs,  // NOLINT(*)
+       const Array<Expr>& args,
+       const std::string& compiler)>;
+
 /*!
  * \brief Forward rewriting rule for a specific op.
  *

include/tvm/relay/transform.h

@@ -576,6 +576,14 @@ TVM_DLL Pass EtaExpand(bool expand_constructor, bool expand_global_var);
  */
 TVM_DLL Pass PrintIR(bool show_meta_data = true);
 
+/*!
+ * \brief Partition a Relay program into regions that can be executed on
+ * different backends.
+ *
+ * \return The pass.
+ */
+TVM_DLL Pass PartitionGraph();
+
 }  // namespace transform
 
 /*!

include/tvm/runtime/contrib/dnnl/dnnl_kernel.h

@@ -0,0 +1,57 @@
+/*
+ * 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 include/tvm/runtime/contrib/dnnl/dnnl_kernel.h
+ * \brief Use external dnnl library kernels.
+ */
+
+#ifndef TVM_RUNTIME_CONTRIB_DNNL_DNNL_KERNEL_H_
+#define TVM_RUNTIME_CONTRIB_DNNL_DNNL_KERNEL_H_
+
+#include "dnnl.hpp"
+
+namespace tvm {
+namespace runtime {
+namespace contrib {
+
+using namespace dnnl;
+
+extern "C" void dnnl_conv2d(float* data, float* weights, float* out, int p_N_,
+                            int p_C_, int p_H_, int p_W_, int p_O_, int p_G_,
+                            int p_Ph_, int p_Pw_, int p_Kh_, int p_Kw_,
+                            int p_Sh_, int p_Sw_);
+
+extern "C" void dnnl_dense(float* data, float* weight, float* out, int p_B_,
+                           int p_I_, int p_O_);
+
+extern "C" void dnnl_relu(float* data, float* out, int p_N_, int p_C_, int p_H_,
+                          int p_W_);
+
+extern "C" void dnnl_bn(float* data, float* gamma, float* beta, float* mean,
+                        float* variance, float* out, int p_n_, int p_c_,
+                        int p_h_, int p_w_, int p_e_);
+
+extern "C" void dnnl_add(float* data, float* weight, float* out, int p_n_,
+                         int p_c_, int p_h_, int p_w_);
+
+}  // namespace contrib
+}  // namespace runtime
+}  // namespace tvm
+#endif  // TVM_RUNTIME_CONTRIB_DNNL_DNNL_KERNEL_H_

python/tvm/module.py

@@ -133,7 +133,16 @@ def export_library(self,
         self.save(path_obj)
         files = [path_obj]
         is_system_lib = self.type_key == "llvm" and self.get_function("__tvm_is_system_module")()
+        has_imported_c_file = False
         if self.imported_modules:
+            for i, m in enumerate(self.imported_modules):
+                if m.type_key == "c":
+                    has_imported_c_file = True
+                    c_file_name = "tmp_" + str(i) + ".cc"
+                    path_cc = temp.relpath(c_file_name)
+                    with open(path_cc, "w") as f:
+                        f.write(m.get_source())
+                    files.append(path_cc)
             path_cc = temp.relpath("devc.cc")
             with open(path_cc, "w") as f:
                 f.write(_PackImportsToC(self, is_system_lib))
@@ -143,7 +152,7 @@ def export_library(self,
                 fcompile = _tar.tar
             else:
                 fcompile = _cc.create_shared
-        if self.type_key == "c":
+        if self.type_key == "c" or has_imported_c_file:
             options = []
             if "options" in kwargs:
                 opts = kwargs["options"]

python/tvm/relay/__init__.py

@@ -29,7 +29,7 @@
 from . import adt
 from . import analysis
 from . import transform
-from .build_module import build, create_executor, optimize
+from .build_module import build, create_executor, optimize, build_extern
 from .transform import build_config
 from . import prelude
 from . import parser

python/tvm/relay/build_module.py

@@ -30,6 +30,7 @@
 from .module import Module as _Module
 from .backend import interpreter as _interpreter
 from .backend.vm import VMExecutor
+from . import transform as _transform
 
 def _update_target(target):
     target = target if target else _target.current_target()
@@ -296,6 +297,33 @@ def optimize(mod, target=None, params=None):
     return mod, params
 
 
+def build_extern(mod, target):
+    """Helper function that builds a Relay function to run on external codegen
+    tools.
+
+    Parameters
+    ----------
+    mod : relay.Module
+        The module to build. Using relay.Function is deprecated.
+
+    target : str
+        The name of the external compilation target.
+
+    Returns
+    -------
+    mod : relay.Module
+        The relay module contains partitioned subgraphes for external codegen
+        tools.
+    """
+    if isinstance(mod, _expr.Function):
+        mod = _Module.from_expr(mod)
+
+    seq = _transform.Sequential([_transform.ExternOp(target),
+                                 _transform.PartitionGraph()])
+    mod = seq(mod)
+    return mod
+
+
 class GraphExecutor(_interpreter.Executor):
     """Wrapper around Executor interface.
 

python/tvm/relay/op/__init__.py

@@ -19,7 +19,7 @@
 # operator defs
 from .op import get, register, register_schedule, register_compute, register_gradient, \
     register_pattern, register_alter_op_layout, register_legalize, \
-    schedule_injective, Op, OpPattern, debug
+    register_extern_op, schedule_injective, Op, OpPattern, debug
 
 # Operators
 from .reduce import *

python/tvm/relay/op/annotation/annotation.py

@@ -62,6 +62,7 @@ def stop_fusion(data):
     """
     return _make.stop_fusion(data)
 
+
 def checkpoint(data):
     """Annotate an expression to be a checkpoint for the checkpointing memory optimization.
 
@@ -78,3 +79,42 @@ def checkpoint(data):
     return _make.checkpoint(data)
 
 register_schedule("annotation.checkpoint", schedule_injective)
+
+
+def subgraph_begin(data, compiler):
+    """Annotate an expression to indicate that it is the beginning of
+    a subgraph.
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        The expression to be annotated.
+
+    compiler : Str
+        The compiler used to generate code of a subgraph.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The annotated expression.
+    """
+    return _make.subgraph_begin(data, compiler)
+
+
+def subgraph_end(data, compiler):
+    """Annotate an expression to indicate that it is the end of a subgraph.
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        The expression to be annotated.
+
+    compiler : Str
+        The compiler used to generate code of a subgraph.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The annotated expression.
+    """
+    return _make.subgraph_end(data, compiler)

python/tvm/relay/op/contrib/__init__.py

@@ -18,4 +18,5 @@
 """Neural network related operators."""
 from __future__ import absolute_import as _abs
 from .contrib import *
+from .extern_op import *
 from . import _contrib