relay op strategy

fix lint

bitpack strategy

bitserial_dense (#6)

* update strategy

* address comments

fix a few topi test

Dense strategy (#5)

* dense

* add biforst; remove comments

* address comment

Refactor x86 conv2d_NCHWc (#4)

* Refactor x86 conv2d

* Add x86 depthwise_conv2d_NCHWc

* Add back topi x86 conv2d_nchw

* Merge x86 conv2d_nchw and conv2d_NCHWc

* Minor fix for x86 conv2d

fix more strategy

Add x86 conv2d_NCHWc_int8 strategy (#8)

* Add x86 conv2d_NCHWc_int8 strategy

* Remove contrib_conv2d_nchwc_int8

* Fix generic conv2d_NCHWc for int8

* Fix topi arm_cpu conv2d_NCHWc_int8

update x86 conv2d

Author: icemelon

include/tvm/relay/op_attr_types.h

@@ -101,8 +101,7 @@ using TShapeDataDependant = bool;
 using FTVMCompute = runtime::TypedPackedFunc<
   Array<top::Tensor>(const Attrs& attrs,
                      const Array<top::Tensor>& inputs,
-                     const Type& out_type,
-                     const Target& target)>;
+                     const Type& out_type)>;
 
 /*!
  * \brief Build the computation schedule for
@@ -118,6 +117,16 @@ using FTVMSchedule = runtime::TypedPackedFunc<
                 const Array<top::Tensor>& outs,
                 const Target& target)>;
 
+/*!
+ * \brief Generate the strategy of operators. This function is a generic
+ * function and can be re-defined for different targets.
+ *
+ * The function signature of generic function is:
+ *   OpStrategy(const Attrs& attrs, const Array<Tensor>& inputs,
+ *              const Type& out_type, const Target& target)
+ */
+using FTVMStrategy = GenericFunc;
+
 /*!
  * \brief Alternate the layout of operators or replace the
  *  operator with other expressions. This function will be invoked
@@ -131,7 +140,8 @@ using FTVMSchedule = runtime::TypedPackedFunc<
 using FTVMAlterOpLayout = runtime::TypedPackedFunc<
   Expr(const Attrs& attrs,
        const Array<Expr>& args,
-       const Array<top::Tensor>& tinfos)>;
+       const Array<top::Tensor>& tinfos,
+       const Type& out_type)>;
 
 /*!
  * \brief Convert the layout of operators or replace the
@@ -186,9 +196,7 @@ using FForwardRewrite = runtime::TypedPackedFunc<
  * \brief Gradient for a specific op.
  *
  * \param orig_call the original Expr.
- *
  * \param output_grad the gradient of the Expr.
- *
  * \return the gradient for each parameters.
  */
 using FPrimalGradient = runtime::TypedPackedFunc<tvm::Array<Expr>(const Expr& orig_call,
@@ -202,14 +210,182 @@ enum AnyCodegenStrategy {
   kVariableDimensions
 };
 
-/* \brief A runtime representation of shape. */
+/*! \brief A runtime representation of shape. */
 using Shape = Array<IndexExpr>;
 
 using FShapeFunc = runtime::TypedPackedFunc<
   Array<top::Tensor>(const Attrs& attrs,
                      const Array<top::Tensor>& inputs,
                      const Array<IndexExpr>& out_ndims)>;
 
+/*!
+ * \brief Operator implementation in TVM.
+ */
+class OpImplementNode : public Object {
+ public:
+  /*! \brief Compute function */
+  FTVMCompute fcompute;
+  /*! \brief Schedule function */
+  FTVMSchedule fschedule;
+  /*! \brief Priority level */
+  Integer plevel;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("plevel", &plevel);
+  }
+
+  static constexpr const char* _type_key = "relay.OpImplement";
+  TVM_DECLARE_FINAL_OBJECT_INFO(OpImplementNode, Object);
+};
+
+/*!
+ * \brief Operator implementation class.
+ */
+class OpImplement : public ObjectRef {
+ public:
+  /*! \brief default constructor  */
+  OpImplement() {}
+  /*! \brief constructor from node pointer */
+  explicit OpImplement(ObjectPtr<Object> n) : ObjectRef(n) {}
+  /*!
+   * \brief access the internal node container
+   * \return the pointer to the internal node container
+   */
+  inline const OpImplementNode* operator->() const;
+  /*!
+   * \brief Invoke the operator compute function.
+   * \param attrs The attribute of the primitive
+   * \param inputs The input tensors.
+   * \param out_type The output type information.
+   * \return The output compute description of the operator.
+   */
+  Array<top::Tensor> Compute(const Attrs& attrs,
+                             const Array<top::Tensor>& inputs,
+                             const Type& out_type);
+  /*!
+   * \brief Build the computation schedule.
+   * \param attrs The attribute of the node.
+   * \param outs The output tensors.
+   * \param target The build target.
+   * \return The computation schedule.
+   */
+  top::Schedule Schedule(const Attrs& attrs,
+                         const Array<top::Tensor>& outs,
+                         const Target& target);
+};
+
+/*!
+ * \brief Specialized implementations for operators under certain conditions.
+ */
+class OpSpecializationNode : public Object {
+ public:
+  /*! \brief List of implementations. */
+  Array<OpImplement> implements;
+  /*! \brief Condition to enable the specialization.
+   *    Could be undefined to represent generic case. */
+  top::SpecializedCondition condition;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("condition", &condition);
+    v->Visit("implements", &implements);
+  }
+
+  static constexpr const char* _type_key = "relay.OpSpecialization";
+  TVM_DECLARE_FINAL_OBJECT_INFO(OpSpecializationNode, ExprNode);
+};
+
+/*!
+ * \brief Operator specialization class.
+ */
+class OpSpecialization : public ObjectRef {
+ public:
+  OpSpecialization() {}
+  explicit OpSpecialization(ObjectPtr<Object> n) : ObjectRef(n) {}
+  /*!
+   * \brief access the internal node container
+   * \return the pointer to the internal node container
+   */
+  inline const OpSpecializationNode* operator->() const;
+  /*!
+   * \brief access the internal node container
+   * \return the pointer to the internal node container
+   */
+  inline OpSpecializationNode* operator->();
+  /*!
+   * \brief Add an implementation.
+   * \param compute Compute function
+   * \param schedule Schedule function
+   * \param plevel Priority level of this implemntation.
+   */
+  void AddImplement(FTVMCompute fcompute, FTVMSchedule fschedule,
+                    int plevel);
+};
+
+/*!
+ * \brief Operator strategy to choose implementation.
+ */
+class OpStrategyNode : public Object {
+ public:
+  /*! \brief List of operator specializations. */
+  Array<OpSpecialization> specializations;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("specializations", &specializations);
+  }
+
+  static constexpr const char* _type_key = "relay.OpStrategy";
+  TVM_DECLARE_FINAL_OBJECT_INFO(OpStrategyNode, ExprNode);
+};
+
+/*!
+ * \brief Operator strategy class.
+ */
+class OpStrategy : public ObjectRef {
+ public:
+  /*! \brief default constructor  */
+  OpStrategy() {}
+  /*! \brief constructor from node pointer */
+  explicit OpStrategy(ObjectPtr<Object> n) : ObjectRef(n) {}
+  /*!
+   * \brief access the internal node container
+   * \return the pointer to the internal node container
+   */
+  inline const OpStrategyNode* operator->() const;
+  /*!
+   * \brief access the internal node container
+   * \return the pointer to the internal node container
+   */
+  inline OpStrategyNode* operator->();
+  /*!
+   * \brief Add an implementation.
+   * \param compute Compute function
+   * \param schedule Schedule function
+   * \param plevel Priority level of this implementation.
+   */
+  void AddImplement(FTVMCompute fcompute, FTVMSchedule fschedule, int plevel);
+};
+
+// implementations
+inline const OpImplementNode* OpImplement::operator->() const {
+  return static_cast<const OpImplementNode*>(get());
+}
+
+inline const OpSpecializationNode* OpSpecialization::operator->() const {
+  return static_cast<const OpSpecializationNode*>(get());
+}
+
+inline OpSpecializationNode* OpSpecialization::operator->() {
+  return static_cast<OpSpecializationNode*>(get_mutable());
+}
+
+inline const OpStrategyNode* OpStrategy::operator->() const {
+  return static_cast<const OpStrategyNode*>(get());
+}
+
+inline OpStrategyNode* OpStrategy::operator->() {
+  return static_cast<OpStrategyNode*>(get_mutable());
+}
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_OP_ATTR_TYPES_H_

include/tvm/top/schedule.h

@@ -28,7 +28,7 @@
 #include <tvm/expr.h>
 #include <tvm/top/tensor.h>
 #include <tvm/top/tensor_intrin.h>
-
+#include <tvm/support/with.h>
 
 #include <string>
 #include <unordered_map>
@@ -744,6 +744,55 @@ class SingletonNode : public IterVarRelationNode {
   TVM_DECLARE_FINAL_OBJECT_INFO(SingletonNode, IterVarRelationNode);
 };
 
+class SpecializedConditionNode;
+
+/*!
+ * \brief Specialized condition to enable op specialization
+ */
+class SpecializedCondition : public ObjectRef {
+ public:
+  SpecializedCondition() {}
+  explicit SpecializedCondition(ObjectPtr<Object> n) : ObjectRef(n) {}
+  /*!
+   * \brief Get the current specialized condition.
+   * \return The current specialized condition.
+   */
+  TVM_DLL static SpecializedCondition Current();
+
+  const SpecializedConditionNode* operator->() const;
+
+  using ContainerType = SpecializedConditionNode;
+  class Internal;
+ private:
+  // enable with syntax.
+  friend class Internal;
+  friend class With<SpecializedCondition>;
+  /*! \brief Push a new specialized condition onto the thread local stack. */
+  TVM_DLL void EnterWithScope();
+  /*! \brief Pop a specialized condition off the thread local context stack. */
+  TVM_DLL void ExitWithScope();
+};
+
+/*! \brief Container for specialization conditions. */
+class SpecializedConditionNode : public Object {
+ public:
+  /*!
+   * \brief List of conditions in conjunctive joint form (CNF).
+   *   Each condition should be a simple expression, e.g., n > 16, m % 8 == 0, etc.,
+   *   where n, m are tvm::Var that represents a dimension in the tensor shape.
+   */
+  Array<PrimExpr> clauses;
+
+  void VisitAttrs(AttrVisitor* v) {
+    v->Visit("clauses", &clauses);
+  }
+
+  static SpecializedCondition make(Array<PrimExpr> conditions);
+
+  static constexpr const char* _type_key = "SpecializedCondition";
+  TVM_DECLARE_FINAL_OBJECT_INFO(SpecializedConditionNode, Object);
+};
+
 
 // implementations
 inline const StageNode* Stage::operator->() const {
@@ -767,6 +816,11 @@ inline const IterVarRelationNode* IterVarRelation::operator->() const {
 inline const IterVarAttrNode* IterVarAttr::operator->() const {
   return static_cast<const IterVarAttrNode*>(get());
 }
+
+inline const SpecializedConditionNode* SpecializedCondition::operator->() const {
+  return static_cast<const SpecializedConditionNode*>(get());
+}
+
 }  // namespace top
 }  // namespace tvm
 #endif  // TVM_TOP_SCHEDULE_H_

python/tvm/__init__.py

@@ -57,7 +57,7 @@
 from .tensor_intrin import decl_tensor_intrin
 from .object import register_object
 from .ndarray import register_extension
-from .schedule import create_schedule
+from .schedule import create_schedule, current_specialization
 from .build_module import build, lower, build_config
 from .tag import tag_scope
 

python/tvm/autotvm/__init__.py

@@ -43,6 +43,7 @@
 from .tuner import callback
 from .task import template, get_config, create, ConfigSpace, ConfigEntity, \
     register_topi_compute, register_topi_schedule, \
+    register_topi_compute2, register_topi_schedule2, register_customized_task, \
     DispatchContext, FallbackContext, ApplyHistoryBest as apply_history_best, \
     ApplyGraphBest as apply_graph_best
 from .env import GLOBAL_SCOPE

python/tvm/autotvm/graph_tuner/base_graph_tuner.py

@@ -25,7 +25,7 @@
 import tvm
 from tvm import autotvm, relay
 from tvm.autotvm.task import get_config
-from tvm.autotvm.task.topi_integration import deserialize_args, serialize_args
+from tvm.autotvm.task.topi_integration import serialize_args
 from tvm.autotvm.record import encode, load_from_file
 from tvm.autotvm.measure import MeasureResult, MeasureInput
 
@@ -42,11 +42,17 @@
     "topi_nn_depthwise_conv2d_nchw": topi.nn.depthwise_conv2d_infer_layout,
 }
 
+def get_infer_layout(task_name):
+    if task_name.startswith("conv2d"):
+        return topi.nn.conv2d_infer_layout
+    elif task_name.startswith("depthwise_conv2d"):
+        return topi.nn.depthwise_conv2d_infer_layout
+    else:
+        raise ValueError("Cannot find infer layout for task %s" % task_name)
 
-@autotvm.template
+@autotvm.register_customized_task("layout_transform")
 def layout_transform(*args):
     """Autotvm layout transform template."""
-    args = deserialize_args(args)
     cfg = get_config()
     cfg.add_flop(-1)
     data = args[0]
@@ -212,7 +218,7 @@ def _fetch_cfg(self):
                 node_entry["record_candidates"] = cache_dict[workload]
                 continue
             record_candidates = []
-            infer_layout_func = OP2LAYOUT[node_entry["topi_op"][0]]
+            infer_layout_func = get_infer_layout(node_entry["topi_op"][0])
             layout_tracking_dict = {}
             for record in cfg_dict[workload]:
                 in_measure, out_measure = record
@@ -264,7 +270,7 @@ def _iterate_layout_transform(self, callback):
 
                 if node_entry["op"] in self._target_ops:
                     o_idx = key
-                    o_infer_layout_func = OP2LAYOUT[node_entry["topi_op"][0]]
+                    o_infer_layout_func = get_infer_layout(node_entry["topi_op"][0])
                     o_wkl = node_entry["workloads"][0]
                     i_topi_op = in_node_entry["topi_op"][0]
                     i_wkl = in_node_entry["workloads"][0]
@@ -273,14 +279,14 @@ def _iterate_layout_transform(self, callback):
                         pivot += 1
                         i_topi_op = in_node_entry["topi_op"][pivot]
                         i_wkl = in_node_entry["workloads"][pivot]
-                    i_infer_layout_func = OP2LAYOUT[i_topi_op]
+                    i_infer_layout_func = get_infer_layout(i_topi_op)
                 else:
                     o_idx = target_input_idx
                     if i <= target_input_pos:
                         continue
-                    o_infer_layout_func = OP2LAYOUT[node_entry["topi_op"][0]]
+                    o_infer_layout_func = get_infer_layout(node_entry["topi_op"][0])
                     o_wkl = node_entry["workloads"][target_input_pos]
-                    i_infer_layout_func = OP2LAYOUT[node_entry["topi_op"][i]]
+                    i_infer_layout_func = get_infer_layout(node_entry["topi_op"][i])
                     i_wkl = node_entry["workloads"][i]
 
                 if (i_idx, o_idx) in pair_tracker:

python/tvm/autotvm/graph_tuner/utils/traverse_graph.py

@@ -65,6 +65,9 @@ def expr2graph(expr, target_ops, node_dict, node_list):
                                % op_name)
         topi_funcs += OP2COMPUTE[op_name]
     env.reset(topi_funcs)
+    # TODO(@kevinthesun, @icemelon9): Currently graph tuning pass relies on the fact
+    #   that # autotvm tasks == # ops. But this won't be true after having relay op
+    #   strategy. We need to find a solution to fix this.
     with env:
         _expr2graph_impl(expr, target_ops, node_dict, node_list)
         task_pos = 0