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__])