[GitHub] [tvm] areusch commented on a change in pull request #9331: [4/10] Code generation for Conv2D via CMSIS-NN

[GitBox <gi...@apache.org>]

areusch commented on a change in pull request #9331:
URL: https://github.com/apache/tvm/pull/9331#discussion_r743906641



##########
File path: tests/python/contrib/test_cmsisnn/test_extract_constants.py
##########
@@ -0,0 +1,175 @@
+# 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.
+
+"""CMSIS-NN integration tests: extract_constants pass"""
+import itertools
+import math
+import numpy as np
+import pytest
+import tvm
+from tvm import relay
+
+from utils import (
+    make_module,
+    count_num_calls,
+    get_range_for_dtype_str,
+    get_same_padding,
+    get_conv2d_qnn_params,
+    make_qnn_relu,
+)
+
+tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
+
+
+class CheckFunctionsForConstants(tvm.relay.ExprVisitor):
+    def __init__(self):
+        super().__init__()
+        self.num_constants_ = 0
+
+    def visit_call(self, call):
+        super().visit_call(call)
+        for arg in call.args:
+            if isinstance(arg, relay.Constant) and arg.data.numpy().ndim > 0:
+                self.num_constants_ += 1
+
+    def visit_function(self, func):
+        super().visit_function(func)
+        assert self.num_constants_ == 0, "Functions should not have constant arguments in Calls"
+
+
+def set_external_func_attr(func, compiler, ext_symbol):
+    func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+    func = func.with_attr("Compiler", compiler)
+    func = func.with_attr("global_symbol", ext_symbol)
+    return func
+
+
+def test_external_function():
+    y0_data = np.random.uniform(0, 1, (8, 8)).astype("float32")
+    x0 = relay.var("x0", shape=(8, 8))
+    y0_const = relay.const(y0_data, "float32")
+    z0 = x0 + y0_const
+    ef = relay.Function([x0], z0, relay.TensorType((8, 8), "float32"))
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "external_compiler", ev.name_hint)
+
+    x = relay.var("x", shape=(8, 8))
+    c = relay.Call(ev, [x])
+    mf = relay.Function([x], c, relay.TensorType((8, 8), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = ExtractConstantsFromPartitionedFunction()(mod)
+    CheckFunctionsForConstants().visit_function(mod[ev])
+    relay.transform.InferType()(mod)
+
+
+def test_nested_function():
+    y1_data = np.random.uniform(0, 1, (8, 8)).astype("float32")
+    x1 = relay.var("x1", shape=(8, 8))
+    y1_const = relay.const(y1_data, "float32")
+    z1 = x1 + y1_const
+    w1 = z1 * relay.const(5.0, "float32")
+    lf = relay.Function([x1], w1, relay.TensorType((8, 8), "float32"))
+
+    x0 = relay.var("x0", shape=(8, 8))
+    c0 = relay.Call(lf, [x0])
+    ef = relay.Function([x0], c0, relay.TensorType((8, 8), "float32"))
+
+    x = relay.var("x", shape=(8, 8))
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "external_compiler", ev.name_hint)
+    c = relay.Call(ev, [x])
+    mf = relay.Function([x], c, relay.TensorType((8, 8), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = ExtractConstantsFromPartitionedFunction()(mod)
+    CheckFunctionsForConstants().visit_function(mod[ev])
+    relay.transform.InferType()(mod)
+
+
+def test_multiple_functions():
+    y20_data = np.random.uniform(0, 1, (8, 8)).astype("float32")
+    x20 = relay.var("x20", shape=(8, 8))
+    y20_const = relay.const(y20_data, "float32")
+    z20 = x20 + y20_const
+    f20 = relay.Function([x20], z20, relay.TensorType((8, 8), "float32"))
+
+    y21_data = np.random.uniform(0, 1, (8, 8)).astype("float32")
+    x21 = relay.var("x21", shape=(8, 8))
+    y21_const = relay.const(y21_data, "float32")
+    z21 = x21 + y21_const
+    f21 = relay.Function([x21], z21, relay.TensorType((8, 8), "float32"))
+
+    x10 = relay.var("x10", shape=(8, 8))
+    c10 = relay.Call(f20, [x10])
+    c11 = relay.Call(f21, [c10])
+    ef = relay.Function([x10], c11, relay.TensorType((8, 8), "float32"))
+
+    x0 = relay.var("x0", shape=(8, 8))
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "external_compiler", ev.name_hint)
+    c = relay.Call(ev, [x0])
+    mf = relay.Function([x0], c, relay.TensorType((8, 8), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = ExtractConstantsFromPartitionedFunction()(mod)
+    CheckFunctionsForConstants().visit_function(mod[ev])
+    relay.transform.InferType()(mod)

Review comment:
       should we assert anything here? same question above

##########
File path: src/relay/backend/contrib/cmsisnn/tir_to_runtime.cc
##########
@@ -54,6 +54,162 @@ class CodeGenCMSISNN : public CodeGenC {
   }
 
  private:
+  /*!  * \brief Emit the CMSIS-NN context buffer */
+  void VisitStmt_(const AllocateNode* op) {
+    context_buffer_name_ = op->buffer_var->name_hint;
+    context_buffer_size_ = op->constant_allocation_size();
+    CodeGenC::VisitStmt_(op);
+  }
+
+  /*!  * \brief Emits CMSIS-NN APIs for every call_extern */
+  void VisitExpr_(const CallNode* op, std::ostream& os) {  // NOLINT(*)
+    if (!op->op.same_as(builtin::call_extern())) {
+      return;
+    }
+    std::string cmsis_func_name = op->args[0].as<StringImmNode>()->value;
+    if (cmsis_func_name == "arm_softmax_s8" || cmsis_func_name == "arm_elementwise_mul_s8" ||
+        cmsis_func_name == "arm_elementwise_add_s8") {
+      CodeGenC::VisitExpr_(op, os);
+    } else if (cmsis_func_name == "arm_convolve_wrapper_s8") {
+      EmitConv2D(op);
+    }
+    return;
+  }
+
+  /*!  * \brief Emits cmsis_nn_context struct */
+  std::string EmitCMSISNNContext(std::ostream& os, std::string buf_name, int buf_size) {
+    std::string struct_name = "context";
+    PrintIndent();
+    os << "cmsis_nn_context " << struct_name << "= {" << buf_name << "," << buf_size << "};\n";
+    return struct_name;
+  }
+
+  /*!  * \brief Emits cmsis_nn_conv_params struct */
+  std::string EmitCMSISNNConvParams(std::ostream& os, int32_t input_offset, int32_t output_offset,
+                                    int32_t stride_w, int32_t stride_h, int32_t padding_w,
+                                    int32_t padding_h, int32_t dilation_w, int32_t dilation_h,
+                                    int32_t clip_min, int32_t clip_max) {
+    std::string struct_name = "conv_params";
+    PrintIndent();
+    os << "cmsis_nn_tile stride = {" << stride_w << "," << stride_h << "};\n";
+    PrintIndent();
+    os << "cmsis_nn_tile padding = {" << padding_w << "," << padding_h << "};\n";
+    PrintIndent();
+    os << "cmsis_nn_tile dilation = {" << dilation_w << "," << dilation_h << "};\n";
+    PrintIndent();
+    os << "cmsis_nn_activation activation = {" << clip_min << "," << clip_max << "};\n";
+    PrintIndent();
+    os << "cmsis_nn_conv_params " << struct_name << " = {" << input_offset << ", " << output_offset
+       << ", stride, padding, dilation, activation};\n";
+    return struct_name;
+  }
+
+  /*!  * \brief Emits cmsis_nn_per_channel_quant_params struct */
+  std::string EmitCMSISNNPerChannelQuantParams(std::ostream& os, std::string multiplier,
+                                               std::string shift) {
+    std::string struct_name = "quant_params";
+    PrintIndent();
+    os << "cmsis_nn_per_channel_quant_params " << struct_name << " = {" << multiplier << ", "
+       << shift << "};\n";
+    return struct_name;
+  }
+
+  /*!  * \brief Emits cmsis_nn_dims struct */
+  std::string EmitCMSISNNDims(std::ostream& os, std::string tensor_type, int32_t n, int32_t h,
+                              int32_t w, int32_t c) {
+    std::string struct_name = tensor_type + "_dims";
+    PrintIndent();
+    os << "cmsis_nn_dims " << struct_name << " = {" << n << "," << h << "," << w << "," << c
+       << "};\n";
+    return struct_name;
+  }
+
+  /*!  * \brief Emits CMSIS-NN APIs for every call_extern */
+  void EmitConv2D(const CallNode* op) {
+    static const int max_num_args = 33;
+    std::string cmsis_func_name = op->args[0].as<StringImmNode>()->value;
+
+    bool bias_enabled = false;
+    if (op->args.size() == max_num_args) {
+      bias_enabled = true;
+    }
+
+    auto get_var_name = [](const CallNode* op, int id) {
+      return op->args[id].as<VarNode>()->name_hint.c_str();
+    };
+    auto get_arg_value = [](const CallNode* op, int id) {
+      return op->args[id].as<IntImmNode>()->value;
+    };
+    int arg_id = 0;
+    std::string input_data = get_var_name(op, ++arg_id);
+    std::string filter_data = get_var_name(op, ++arg_id);
+    std::string multiplier = get_var_name(op, ++arg_id);
+    std::string bias_data("0x0");
+    if (bias_enabled) {
+      bias_data = get_var_name(op, ++arg_id);
+    }
+    std::string shift = get_var_name(op, ++arg_id);
+    std::string output_data = get_var_name(op, ++arg_id);
+
+    int input_offset = get_arg_value(op, ++arg_id);

Review comment:
       is there a way to use the encoding here?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org