Hack for inserting start/end point of quantizing. (apache#33)

* Hack for inserting start/end point of quantizing.

* Fix.

Author: ZihengJiang

include/tvm/relay/attrs/transform.h

@@ -150,6 +150,18 @@ struct SliceLikeAttrs : public tvm::AttrsNode<SliceLikeAttrs> {
   }
 };
 
+
+/*! \brief Attributes for Annotate operator */
+struct AnnotateAttrs : public tvm::AttrsNode<AnnotateAttrs> {
+  std::string info;
+
+  TVM_DECLARE_ATTRS(AnnotateAttrs, "relay.attrs.AnnotateAttrs") {
+    TVM_ATTR_FIELD(info)
+      .describe("The annotation info.");
+  }
+};
+
+
 /*! \brief Attributes for Clip operator */
 struct ClipAttrs : public tvm::AttrsNode<ClipAttrs> {
   double a_min;

python/tvm/relay/quantize/quantize.py

@@ -8,7 +8,7 @@
 from .. import op as _op
 from .. import build_module as _build
 from ..base import register_relay_node
-from ..._ffi.function import register_func
+from ..._ffi.function import register_func, get_global_func
 
 
 class QFieldKind(object):
@@ -72,14 +72,27 @@ def simulated_quantize(data, dom_scale, bit, clip_min, clip_max,
                                         sign, rounding, kind)
 
 
+
+SQ_CACHE_MAP = {}
+
+
 @register_func("relay.quantize.attach_simulated_quantize")
 def attach_simulated_quantize(data, kind):
+    global SQ_CACHE_MAP
+    key = data
+    if data in SQ_CACHE_MAP:
+        return SQ_CACHE_MAP[data]
+    if len(SQ_CACHE_MAP) == 0:
+        f = get_global_func("relay._quantize.make_annotate_op")
+        data = f(data, "quantize_start")
     dom_scale = _expr.var("dom_scale")
     bit = _expr.var("bit")
     clip_min = _expr.var("clip_min")
     clip_max = _expr.var("clip_max")
-    return simulated_quantize(data, dom_scale, bit, clip_min, clip_max,
-                              True, "round", kind)
+    ret = simulated_quantize(data, dom_scale, bit, clip_min, clip_max,
+                             True, "round", kind)
+    SQ_CACHE_MAP[key] = ret
+    return ret
 
 
 def register_qfield_rewrite(op_name, frewrite=None, level=10):

src/relay/op/tensor/unary.cc

@@ -81,6 +81,24 @@ RELAY_REGISTER_UNARY_OP("copy")
 .set_support_level(3)
 .set_attr<FTVMCompute>("FTVMCompute", RELAY_UNARY_COMPUTE(topi::identity));
 
+
+// relay.annotate
+TVM_REGISTER_NODE_TYPE(AnnotateAttrs);
+
+RELAY_REGISTER_OP("annotate")
+.describe(R"code(Copy a tensor with annotation information
+)code" TVM_ADD_FILELINE)
+.set_num_inputs(1)
+.add_argument("data", "Tensor", "The input tensor.")
+.add_type_rel("Identity", IdentityRel)
+.set_attr<TOpPattern>("TOpPattern", kElemWise)
+.set_attr<TOpIsStateful>("TOpIsStateful", false)
+.set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout)
+.set_attrs_type_key("relay.attrs.AnnotateAttrs")
+.set_support_level(10)
+.set_attr<FTVMCompute>("FTVMCompute", RELAY_UNARY_COMPUTE(topi::identity));
+
+
 // relay.clip
 TVM_REGISTER_NODE_TYPE(ClipAttrs);
 

src/relay/pass/quantize.cc

@@ -66,13 +66,26 @@ TVM_REGISTER_API("relay._quantize.simulated_quantize")
    });
 
 
+Expr MakeAnnotateOp(Expr data, std::string info) {
+  auto attrs = make_node<AnnotateAttrs>();
+  attrs->info = info;
+  static const Op& op = Op::Get("annotate");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay._quantize.make_annotate_op")
+.set_body_typed<Expr(Expr, std::string)>(MakeAnnotateOp);
+
+
 // =============
 // annotate pass
 
 Expr QFieldExprNode::Realize() const {
   // store low bit output back
   const PackedFunc* f = runtime::Registry::Get("relay.quantize.attach_simulated_quantize");
-  return (*f)(expr, static_cast<int>(kQInput));
+  Expr ret = (*f)(expr, static_cast<int>(kQInput));
+  ret = MakeAnnotateOp(ret, "quantize_end");
+  return ret;
 }
 
 QFieldExpr QFieldExprNode::make(Expr expr, QFieldKind kind) {

src/relay/pass/quantize.h

@@ -36,6 +36,7 @@ struct SimulatedQuantizeAttrs : public tvm::AttrsNode<SimulatedQuantizeAttrs> {
   }
 };
 
+
 class QFieldExpr;
 
 class QFieldExprNode : public TempExprNode {