[HYBRID FRONTEND] Modify hybrid script to new interface; hybrid op supported; enable compilation_database in CMakeList.txt

CMakeLists.txt

@@ -57,6 +57,7 @@ include_directories("3rdparty/compiler-rt")
 # initial variables
 set(TVM_LINKER_LIBS "")
 set(TVM_RUNTIME_LINKER_LIBS "")
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # Generic compilation options
 if(MSVC)

docs/langref/hybrid_script.rst

@@ -22,13 +22,15 @@ you need to use ``tvm.hybrid.script`` decorator to indicate this is a hybrid fun
 
     @tvm.hybrid.script
     def outer_product(a, b, c):
+        c = output_tensor((100, 99), 'float32')
         for i in range(a.shape[0]):
             for j in range(b.shape[0]):
                 c[i, j] = a[i] * b[j]
-    a = numpy.random.rand(100)
-    b = numpy.random.rand(99)
-    c = numpy.zeros((100, 99))
-    outer_product(a, b, c)
+          return c
+    a = numpy.random.randn(100)
+    b = numpy.random.randn(99)
+    c = outer_product(a, b)
+
 
 This decorator will import `Keywords`_ required spontaneously when software emulation.
 After software emulation is done, the imported keywords will be cleaned up. Users do not need
@@ -40,25 +42,25 @@ or ``numpy`` numeric type.
 Backend Compilation
 ~~~~~~~~~~~~~~~~~~~
 
+This function is not encouraged to use, users are encouraged to use the second interface.
 The current parse interface looks like:
 
 .. code-block:: python
 
    a = tvm.placeholder((100, ), name='a')
    b = tvm.placeholder((99, ), name='b')
-   c = tvm.placeholder((100, 99), name='c')
-   tvm.hybrid.parse(outer_product, [a, b, c]) # return an ir root of this function
+   parser = tvm.hybrid.parse(outer_product, [a, b]) # return the parser of this function
 
-If we pass these tvm tensors to this function, it returns a op node:
 
-**Under construction, we are still deciding what kind of node should be returned.**
+If we pass these tvm tensors to this function, it returns a op node:
 
 .. code-block:: python
 
    a = tvm.placeholder((100, ), name='a')
    b = tvm.placeholder((99, ), name='b')
-   c = tvm.placeholder((100, 99), name='c')
-   op = outer_product(a, b, c) # return the corresponding op node
+   c = outer_product(a, b, c) # return the output tensor(s) of the operator
+
+**Under construction, we are still deciding what kind of node should be returned.**
 
 Tuning
 ~~~~~~

include/tvm/operation.h

@@ -450,6 +450,69 @@ class ExternOpNode : public OperationNode {
   TVM_DECLARE_NODE_TYPE_INFO(ExternOpNode, OperationNode);
 };
 
+/*!
+ * \brief A computation operator that generated by hybrid script.
+ */
+class HybridOpNode : public OperationNode {
+ public:
+  /*! \brief The input tensors */
+  Array<Tensor> inputs;
+  /*! \brief Symbolic placeholder representation of outputs */
+  Array<Tensor> outputs;
+  /*! \brief the statement that generates the computation. This is
+   * slightly different from the body in ExternOpNode. All the output
+   * tensors keep its own name specified by users in the script.
+   * However, when compilation, these tensors will be placed by those
+   * actual output tensors. */
+  Stmt body;
+
+  /*! \brief constructor */
+  HybridOpNode() {}
+  // override functions
+  int num_outputs() const final;
+  Array<IterVar> root_iter_vars() const final;
+  Type output_dtype(size_t i) const final;
+  Array<Expr> output_shape(size_t i) const final;
+  Array<Tensor> InputTensors() const final;
+  Operation ReplaceInputs(
+      const Operation& self,
+      const std::unordered_map<Tensor, Tensor>& rmap) const final;
+  void PropBoundToInputs(
+      const Operation& self,
+      const std::unordered_map<const Variable*, IntSet>& dom_map,
+      std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
+  void GatherBound(
+      const Operation& self,
+      const std::unordered_map<Tensor, TensorDom>& tensor_dom,
+      std::unordered_map<IterVar, Range>* out_dom_map) const final;
+  Stmt BuildRealize(
+      const Stage& stage,
+      const std::unordered_map<IterVar, Range>& realize_map,
+      const Stmt& body) const final;
+  Stmt BuildProvide(
+      const Stage& stage,
+      const std::unordered_map<IterVar, Range>& dom_map,
+      bool debug_keep_trivial_loop) const final;
+
+  void VisitAttrs(AttrVisitor* v) final {
+    v->Visit("name", &name);
+    v->Visit("tag", &tag);
+    v->Visit("attrs", &attrs);
+    v->Visit("inputs", &inputs);
+    v->Visit("outputs", &outputs);
+    v->Visit("body", &body);
+  }
+  EXPORT static Operation make(std::string name,
+                               std::string tag,
+                               Map<std::string, NodeRef> attrs,
+                               Array<Tensor> inputs,
+                               Array<Tensor> outputs,
+                               Stmt body);
+
+  static constexpr const char* _type_key = "HybridOp";
+  TVM_DECLARE_NODE_TYPE_INFO(HybridOpNode, OperationNode);
+};
+
 /*! \brief The compute function to specify the input source of a Tensor */
 using FCompute = std::function<Expr (const Array<Var>& i)>;
 

python/tvm/build_module.py

@@ -340,11 +340,6 @@ def lower(sch,
         bounds = schedule.InferBound(sch)
         stmt = schedule.ScheduleOps(sch, bounds)
         stmt = ir_pass.InjectPrefetch(stmt)
-    else:
-        #So far there is no op for hybrid script, so a plain ir body is given
-        if not isinstance(sch, _stmt.Stmt):
-            raise ValueError("sch should be either a Schedule or a Stmt")
-        stmt = sch
 
     for f in lower_phase0:
         stmt = f(stmt)

python/tvm/hybrid/__init__.py

@@ -7,4 +7,5 @@
 2. Developers can build HalideIR by writing Python code.
 """
 
-from .api import script, parse
+from .api import script
+from .parser import parse_python

python/tvm/hybrid/api.py

@@ -1,9 +1,12 @@
 """APIs of lowering the Python subset to HalideIR"""
 from __future__ import absolute_import as _abs
 
-import types
 from .._ffi.base import decorate
+from .. import _api_internal as _tvm_internal
+from ..tensor import Tensor
+
 from .parser import parse_python
+from .util import _pruned_source
 
 
 def script(pyfunc):
@@ -17,40 +20,26 @@ def script(pyfunc):
     hybrid_func : function
         A decorated hybrid script function.
     """
-    def wrapped_func(func, *args, **kwargs):
+    def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring
         from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types
         if _is_tvm_arg_types(args):
-            return parse(func, args)
+            src = _pruned_source(func)
+            parser = parse_python(src, args)
+
+            input_tensors = []
+            for i in args:
+                if isinstance(i, Tensor):
+                    input_tensors.append(i)
+
+            op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors,
+                                         parser.outputs, parser.parsed_body)
+            res = [op.output(i) for i in range(len(parser.outputs))]
+
+            return res[0] if len(res) == 1 else res
 
         intersect = _enter_hybrid_runtime(func)
         value = func(*args, **kwargs)
         _restore_runtime(func, intersect)
         return value
-    return decorate(pyfunc, wrapped_func)
-
-
-def parse(func, args):
-    """Parse a subset of Python to HalideIR
 
-    Parameters
-    ----------
-    func : str or types.FunctionType
-        If it is a string, parse the source code
-        If it is a function, parse the function
-
-    args : list of Buffer or Tensor or Var
-        The argument lists to the function.
-        Leave it None if no buffer is related to the function to be parsed
-
-    Returns
-    -------
-    root : Stmt
-        The result Halide IR and the parser class instance.
-    """
-    from .util import _pruned_source
-    if isinstance(func, str):
-        src = func
-    else:
-        assert isinstance(func, types.FunctionType)
-        src = _pruned_source(func)
-    return parse_python(src, args)
+    return decorate(pyfunc, wrapped_func)

python/tvm/hybrid/intrin.py

@@ -48,6 +48,7 @@ def allocate(shape, dtype='float32', scope='global'): #pylint: disable=unused-ar
     """
     return numpy.zeros(shape).astype(dtype)
 
+output_tensor = allocate #pylint: disable=invalid-name
 
 def popcount(x):
     """
@@ -87,18 +88,19 @@ def sigmoid(x):
 
 
 HYBRID_GLOBALS = {
-    'unroll'    : unroll,
-    'vectorize' : vectorize,
-    'parallel'  : parallel,
-    'allocate'  : allocate,
-    'bind'      : bind,
-    'sqrt'      : numpy.sqrt,
-    'log'       : numpy.log,
-    'tanh'      : numpy.tanh,
-    'power'     : numpy.power,
-    'exp'       : numpy.exp,
-    'sigmoid'   : sigmoid,
-    'popcount'  : popcount
+    'unroll'       : unroll,
+    'vectorize'    : vectorize,
+    'parallel'     : parallel,
+    'allocate'     : allocate,
+    'output_tensor': output_tensor,
+    'bind'         : bind,
+    'sqrt'         : numpy.sqrt,
+    'log'          : numpy.log,
+    'tanh'         : numpy.tanh,
+    'power'        : numpy.power,
+    'exp'          : numpy.exp,
+    'sigmoid'      : sigmoid,
+    'popcount'     : popcount
 }