Cleaning up Arm(R) Ethos(TM)-U codegen

This is a follow up commit to address the
comments of apache#8849

Change-Id: I02d8de64f3bce0e7b544d652eee8737ec1ecbb80

python/tvm/relay/backend/contrib/ethosu/codegen.py

@@ -24,33 +24,34 @@
 from tvm.relay.backend.contrib.ethosu import util
 
 
-@tvm._ffi.register_func("relay.ext.ethosu.constant_updater")
-def constant_updater(expr, symbol):  # pylint: disable=unused-argument
-    """
-    We dont want the build process to extract constants to be loaded in
-    the runtime as we are embedding them inside the C runtime.Module.
-    """
-    return dict()
-
-
 @tvm._ffi.register_func("relay.ext.ethosu")
-def ethosu_compiler(ref):
-    """Main function to a compile a given relay function of
+def ethosu_compiler(external_function):
+    """The entry-point to a compile a external relay function of
     NPU compatible operators to generated command stream.
-    Such generated command stream would be loaded to the runtime
-    module that interfaces with NPU driver.
+    Such generated command stream would be used to create c-source r
+    runtime module that interfaces with NPU driver.
     """
-    assert isinstance(ref, tvm.ir.function.BaseFunc)
-    func_name = ref.attrs["global_symbol"]
+    assert isinstance(external_function, tvm.ir.function.BaseFunc)
+    func_name = external_function.attrs["global_symbol"]
     # There should only be a single input
-    assert len(ref.params) == 1
-    input_size = util.calculate_size_bytes(ref.params[0])
-    output_size = util.calculate_size_bytes(ref.body)
-    cmms, encoded_constants, scratch_size = _compile(ref)
+    assert len(external_function.params) == 1
+    input_size = util.calculate_size_bytes(external_function.params[0])
+    output_size = util.calculate_size_bytes(external_function.body)
+    cmms, encoded_constants, scratch_size = _compile(external_function)
     ethosu_runtime = tvm._ffi.get_global_func("runtime.module.ethosu.create")
     return ethosu_runtime(func_name, cmms, encoded_constants, scratch_size, input_size, output_size)
 
 
+@tvm._ffi.register_func("relay.ext.ethosu.constant_updater")
+def constant_updater(expr, symbol):  # pylint: disable=unused-argument
+    """
+    The constant updater process happen after lowering in the core compiler.
+    For the NPU, we dont want the build process to extract constants to be loaded in
+    the runtime as we are embedding them inside the C runtime.Module.
+    """
+    return dict()
+
+
 def _compile(ext_func):
     """
     This is the main wrapper that accepts an external

src/relay/backend/contrib/ethosu/source_module.cc

@@ -16,6 +16,11 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
+/*!
+ * \file source_module.cc
+ * \brief Source code module for the host to invoke the NPU
+ */
 #include <dmlc/filesystem.h>
 #include <dmlc/logging.h>
 #include <dmlc/memory_io.h>
@@ -40,20 +45,24 @@
 namespace tvm {
 namespace runtime {
 
+// The runtime.Module that contains the host-side c code
+// required for invoking the NPU with the command stream
 class EthosUModuleNode : public ModuleNode {
  public:
   /*!
    * \brief The ethos runtime module.
    *
-   * \param cmms A array of external symbol 1, serialized command stream 1
-   * external symbol 2, serialized command stream 2, ....
-   * TODO : if and when FFI support Maps with non-objects OR compound arrays
-   * switch to that.
+   * \param func_name_ name of the should be codegen'd function
+   * \param cmms_hex_ command stream for the NPU in hex
+   * \param weights_bias_hex_ the encoded biases and weights for the NPU in hex
+   * \param scratch_size_ the size of the scratch memory required for command stream
+   * \param input_size_ the size (in bytes) for the input tensor
+   * \param output_size_ the size (in bytes) for the output tensor
    */
   explicit EthosUModuleNode(const String& func_name_, const String& cmms_hex_,
                             const String& weights_bias_hex_, const Integer& scratch_size_,
                             const Integer& input_size_, const Integer& output_size_) {
-    func_names_.push_back(func_name_);
+    func_name = func_name_;
     cmms_hex = std::move(cmms_hex_);
     weights_bias_hex = std::move(weights_bias_hex_);
     scratch_size = scratch_size_->value;
@@ -92,8 +101,9 @@ class EthosUModuleNode : public ModuleNode {
    */
   PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final {
     if (name == "get_func_names") {
-      return PackedFunc(
-          [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->func_names_; });
+      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+        *rv = Array<String>{this->func_name};
+      });
     }
     return PackedFunc();
   }
@@ -109,7 +119,7 @@ class EthosUModuleNode : public ModuleNode {
 
  private:
   String c_source;
-  Array<String> func_names_;
+  String func_name;
   String cmms_hex;
   String weights_bias_hex;
   size_t scratch_size;
@@ -202,13 +212,12 @@ class EthosUModuleNode : public ModuleNode {
    * \return string of code that offloads a subgraph to the NPU
    */
   std::string GenerateSource() {
-    std::string func_no_dashes = func_names_[0];
+    std::string func_no_dashes = func_name;
     std::replace(func_no_dashes.begin(), func_no_dashes.end(), '-', '_');
     std::stringstream ss;
 
     ss << "#include <stdio.h>\n";
     ss << "#include <stdlib.h>\n";
-    ss << "#include <dlpack/dlpack.h>\n";
     ss << "#include <tvm/runtime/crt/module.h>\n";
     ss << "#include <ethosu_driver.h>\n";
     ss << "\n";
@@ -282,7 +291,7 @@ class EthosUModuleNode : public ModuleNode {
     PrintExternCPostfix(ss);
     ss << "\n";
     PrintExternCPrefix(ss);
-    PrintRuntimeFunctionHeader(ss, func_names_[0]);
+    PrintRuntimeFunctionHeader(ss, func_name);
     EnterScope();
     PrintIndents(ss);
     ss << "return " << func_no_dashes << "_wrapper_(input, output);\n";
@@ -308,9 +317,12 @@ inline EthosUModuleNode* EthosUModule::operator->() {
   return static_cast<EthosUModuleNode*>(get_mutable());
 }
 
-TVM_REGISTER_GLOBAL("runtime.module.ethosu.create").set_body([](TVMArgs args, TVMRetValue* rv) {
-  *rv = EthosUModuleNode::Create(args[0], args[1], args[2], args[3], args[4], args[5]);
-});
+TVM_REGISTER_GLOBAL("runtime.module.ethosu.create")
+    .set_body_typed([](String func_name, String cmms_hex, String weights_bias_hex,
+                       Integer scratch_size, Integer input_size, Integer output_size) {
+      return EthosUModuleNode::Create(func_name, cmms_hex, weights_bias_hex, scratch_size,
+                                      input_size, output_size);
+    });
 
 TVM_REGISTER_GLOBAL("runtime.module.ethosu.getcs").set_body_typed([](EthosUModule mod) {
   return mod->GetCS();

tests/python/contrib/test_ethosu/infra.py

@@ -228,14 +228,14 @@ def _create_test_runner(accel):
     )
 
 
-def build_source(module, inputs, outputs, accel="ethos-u55-256"):
+def build_source(module, inputs, outputs, accel="ethos-u55-256", output_tolerance=0):
     test_runner = _create_test_runner(accel)
     return compile_models(
         models=AOTTestModel(
             module=module,
             inputs=inputs,
             outputs=outputs,
-            output_tolerance=10,
+            output_tolerance=output_tolerance,
             extra_memory_in_bytes=16 * 1024 * 1024,
         ),
         interface_api="c",

tests/python/contrib/test_ethosu/test_codegen.py

@@ -152,10 +152,7 @@ def create_graph_activation(input_tensor_name, input_tensor_shape, input_tensor_
         output_data = generate_ref_data(relay_module, input_data)
 
         compiled_models = infra.build_source(
-            mod,
-            input_data,
-            output_data,
-            accel_type,
+            mod, input_data, output_data, accel_type, output_tolerance=1
         )
 
         # Assumes only two runtime.Modules are created -- i.e. single offload module

tests/python/contrib/test_ethosu/test_networks.py

@@ -58,7 +58,9 @@ def test_forward_mobilenet_v1(accel_type="ethos-u55-256"):
     output_data = generate_ref_data(relay_mod, input_data)
 
     mod = partition_for_ethosu(relay_mod, params)
-    compiled_models = infra.build_source(mod, input_data, output_data, accel_type)
+    compiled_models = infra.build_source(
+        mod, input_data, output_data, accel_type, output_tolerance=10
+    )
     infra.verify_source(compiled_models, accel_type)
 
 

tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py

@@ -764,7 +764,4 @@ def check_buffer(address, region, length, buffer_var):
 
 
 if __name__ == "__main__":
-    test_buffer_info_extraction()
-    test_translate_ethosu_conv2d()
-    test_translate_ethosu_copy()
-    test_assign_addresses()
+    pytest.main([__file__])