code review: tests for new passes, clean up of relay_to_tir for cmsis-nn

Change-Id: I78dafb2be49afcd5b816c1d2828246c213edca85

Author: Ashutosh Parkhi

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

@@ -23,6 +23,8 @@
 from ...dataflow_pattern import is_constant, is_op, wildcard
 from .register import register_pattern_table
 
+tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
+
 
 def enabled():
     return bool(tvm.get_global_func("relay.ext.cmsisnn", True))
@@ -47,8 +49,6 @@ def partition_for_cmsisnn(mod, params=None, **opts):
     if params:
         mod["main"] = bind_params_by_name(mod["main"], params)
 
-    tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
-
     seq = tvm.transform.Sequential(
         [
             transform.InferType(),
@@ -91,7 +91,7 @@ def qnn_conv2d_pattern():
         """Create pattern for qnn.conv2D with optional fused relu."""
         qnn_conv2d = is_op("qnn.conv2d")(
             wildcard(), is_constant(), is_constant(), is_constant(), is_constant(), is_constant()
-        ).has_attr({"kernel_layout": "HWIO"})
+        )
         bias_add = is_op("nn.bias_add")(qnn_conv2d, is_constant())
         req = is_op("qnn.requantize")(
             qnn_conv2d | bias_add, is_constant(), is_constant(), is_constant(), is_constant()
@@ -123,8 +123,7 @@ def check_qnn_conv2d(pattern):
         kernel_zp = [kernel_zp] if kernel_zp.ndim == 0 else kernel_zp
 
         return (
-            conv2d.attrs.kernel_layout == "HWIO"
-            and conv2d.attrs.out_dtype == "int32"
+            conv2d.attrs.out_dtype == "int32"
             and conv2d.attrs.padding[2] == 0
             and conv2d.attrs.padding[3] == 0
             and conv2d_input.checked_type.dtype == "int8"

src/relay/backend/contrib/cmsisnn/extract_constants.cc

@@ -17,6 +17,11 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+/*!
+ * \file extract_constant.cc
+ * \brief Pushes out constants within partitioned functions all the way upto main()
+ */
+
 #include <tvm/relay/attrs/nn.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/transform.h>
@@ -30,44 +35,47 @@ namespace relay {
 namespace contrib {
 namespace cmsisnn {
 
+/*!
+ * \brief This Mutator finds all functions with constants. Constants are replaced with function
+ * parameter variables. Constants are pushed all the way upto main().
+ */
 class ExtractConstantsMutator : public MixedModeMutator {
  public:
-  explicit ExtractConstantsMutator(IRModule& mod) : mod_(mod) {}
+  explicit ExtractConstantsMutator(const IRModule& mod) : mod_(mod) {}
 
  private:
   String gen_var_name() { return "tvm_var_extract_const_" + std::to_string(var_count_++); }
 
-  Expr VisitExpr_(const FunctionNode* func) final {
-    Function final_func = GetRef<Function>(func);
-    ++func_nesting_level_;
+  Expr VisitExpr_(const FunctionNode* function) final {
+    Function func = GetRef<Function>(function);
+    function_to_constants_.Set(func, Array<Constant>{});
+    functions_.push_back(func);
     auto new_body = VisitExpr(func->body);
-    --func_nesting_level_;
-    if (!new_body.same_as(func->body)) {
-      final_func = Function(FreeVars(new_body), new_body, func->ret_type,
-                            FreeTypeVars(new_body, mod_), func->attrs);
-      function_to_constants_.Set(GetRef<Function>(func), constants_within_function_);
-      constants_within_function_.clear();
+    functions_.pop_back();
+    if (function_to_constants_[func].size()) {
+      func = Function(FreeVars(new_body), new_body, func->ret_type, FreeTypeVars(new_body, mod_),
+                      func->attrs);
     }
-    return final_func;
+    return func;
   }
 
   Expr Rewrite_(const CallNode* call, const Expr& post) final {
     Expr final_call = post;
     auto* post_call = post.as<CallNode>();
-    if (post_call == nullptr) {
-      return final_call;
-    }
 
     // Replace Constant arguments with Vars for ML Operators
     // Perform this for non-main Call Nodes only
-    if (func_nesting_level_ && call->op.as<OpNode>()) {
+    if (!functions_.empty() && call->op.as<OpNode>()) {
       Array<Expr> new_args;
       for (auto& arg : post_call->args) {
         auto* const_arg = arg.as<ConstantNode>();
         if (const_arg && !const_arg->is_scalar()) {
           Var var_arg = Var(gen_var_name(), const_arg->tensor_type());
           new_args.push_back(var_arg);
-          constants_within_function_.push_back(GetRef<Constant>(const_arg));
+          const Function& last_func = functions_.back();
+          Array<Constant> fconstants(function_to_constants_[last_func]);
+          fconstants.push_back(GetRef<Constant>(const_arg));
+          function_to_constants_.Set(last_func, fconstants);
         } else {
           new_args.push_back(arg);
         }
@@ -94,17 +102,21 @@ class ExtractConstantsMutator : public MixedModeMutator {
 
     // Since the constants are kicked out of the local partitioned functions
     // a new call to local function is needed
+    // Also, pass on the constants to the callee of this function to support nested functions
     if (auto* func_node = call->op.as<FunctionNode>()) {
       Function func = GetRef<Function>(func_node);
       auto new_func = VisitExpr(func);
       if (!new_func.same_as(func)) {
         Array<Expr> new_args = post_call->args;
         ICHECK(function_to_constants_.find(func) != function_to_constants_.end());
+        const Function& last_func = functions_.back();
+        Array<Constant> fconstants(function_to_constants_[last_func]);
         for (auto constant : function_to_constants_.at(func)) {
-          constants_within_function_.push_back(constant);
+          fconstants.push_back(constant);
           Var var_arg = Var(gen_var_name(), constant->tensor_type());
           new_args.push_back(var_arg);
         }
+        function_to_constants_.Set(last_func, fconstants);
         final_call = Call(new_func, new_args);
       }
     }
@@ -117,16 +129,14 @@ class ExtractConstantsMutator : public MixedModeMutator {
   IRModule mod_;
   /* \brief Maintains mapping of original function to the replaced constants */
   Map<Function, Array<Constant>> function_to_constants_;
-  /* \brief Constants being kicked out of a function during the function visit */
-  Array<Constant> constants_within_function_;
+  /* \brief Stack of functions to determine scope while filling up function_to_constants_ */
+  Array<Function> functions_;
   /* \brief Keeps track of variables being created */
   int var_count_ = 0;
-  /* \brief Keeps track of function scope */
-  int func_nesting_level_ = 0;
 };
 
 /*!  * \brief Kicks out all constants out of the partitioned function into main()  */
-IRModule ExtractConstants(IRModule mod) {
+IRModule ExtractConstants(const IRModule& mod) {
   String func_name;
   Function func;
 
@@ -150,7 +160,7 @@ transform::Pass ExtractConstantsFromPartitionedFunction() {
 }
 
 TVM_REGISTER_GLOBAL("relay.ext.cmsisnn.transform.ExtractConstantsFromPartitionedFunction")
-    .set_body_typed([]() { return ExtractConstantsFromPartitionedFunction(); });
+    .set_body_typed(ExtractConstantsFromPartitionedFunction);
 
 }  // namespace cmsisnn
 }  // namespace contrib

src/relay/backend/contrib/cmsisnn/generate_constants.cc

@@ -17,24 +17,35 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+/*!
+ * \file generate_constant.cc
+ * \brief Generates quantization parameters needed by CMSIS-NN
+ */
+
 #include <tvm/relay/attrs/nn.h>
 #include <tvm/relay/attrs/transform.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/transform.h>
 #include <tvm/runtime/ndarray.h>
 
+#include "../../../op/make_op.h"
 #include "../../../qnn/utils.h"
 #include "../../../transforms/pattern_utils.h"
 
 namespace tvm {
 namespace relay {
-Expr MakeTranspose(Expr data, Array<Integer> axes);
 namespace contrib {
 namespace cmsisnn {
 
+/*!
+ * \brief This Mutator will find all partitioned functions meant for CMSIS-NN Conv2D.
+ * It will substitute original Conv2D's weight zero point and original Requantize's input zero point
+ * with CMSIS-NN's quantization parameters.
+ * https://github.com/tensorflow/tflite-micro/blob/0f40100fc60276e9f345c23282de3baf19a78059/tensorflow/lite/kernels/internal/quantization_util.cc#L53
+ */
 class GenerateConstantsMutator : public MixedModeMutator {
  public:
-  explicit GenerateConstantsMutator(IRModule& mod) : mod_(mod) {}
+  explicit GenerateConstantsMutator(const IRModule& mod) : mod_(mod) {}
 
  private:
   /*!  * \brief Converts Kernel layout from HWIO to OHWI to align to CMSIS-NN requirements */
@@ -52,8 +63,15 @@ class GenerateConstantsMutator : public MixedModeMutator {
     attrs->out_dtype = std::move(conv2d_attrs->out_dtype);
     *new_attrs = tvm::Attrs{attrs};
 
+    std::string kernel_layout = conv2d_attrs->kernel_layout.c_str();
+    int pos_o = kernel_layout.find("O");
+    int pos_h = kernel_layout.find("H");
+    int pos_w = kernel_layout.find("W");
+    int pos_i = kernel_layout.find("I");
+
     IRModule kernel_module;
-    auto func_body = MakeTranspose(kernel_expr, {Integer(3), Integer(0), Integer(1), Integer(2)});
+    auto func_body = MakeTranspose(
+        kernel_expr, {Integer(pos_o), Integer(pos_h), Integer(pos_w), Integer(pos_i)});
     auto kernel_func =
         Function(FreeVars(func_body), func_body, Type(), FreeTypeVars(func_body, kernel_module));
     GlobalVar kernel_var("main");
@@ -158,9 +176,6 @@ class GenerateConstantsMutator : public MixedModeMutator {
   Expr Rewrite_(const CallNode* call, const Expr& post) final {
     Expr final_call = post;
     auto* post_call = post.as<CallNode>();
-    if (post_call == nullptr) {
-      return final_call;
-    }
 
     auto* global_var = call->op.as<GlobalVarNode>();
     if (global_var) {
@@ -196,7 +211,7 @@ class GenerateConstantsMutator : public MixedModeMutator {
   IRModule mod_;
 };
 
-IRModule GenerateConstants(IRModule mod) {
+IRModule GenerateConstants(const IRModule& mod) {
   String func_name;
   Function func;
 
@@ -220,9 +235,8 @@ transform::Pass GenerateCMSISNNConstants() {
   return tvm::transform::CreateModulePass(pass_func, 0, "GenerateCMSISNNConstants", {});
 }
 
-TVM_REGISTER_GLOBAL("relay.ext.cmsisnn.transform.GenerateCMSISNNConstants").set_body_typed([]() {
-  return GenerateCMSISNNConstants();
-});
+TVM_REGISTER_GLOBAL("relay.ext.cmsisnn.transform.GenerateCMSISNNConstants")
+    .set_body_typed(GenerateCMSISNNConstants);
 
 }  // namespace cmsisnn
 }  // namespace contrib