diff --git a/python/tvm/relay/backend/contrib/ethosu/codegen.py b/python/tvm/relay/backend/contrib/ethosu/codegen.py index e821ea8bf0c4b..827146f77c804 100644 --- a/python/tvm/relay/backend/contrib/ethosu/codegen.py +++ b/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 diff --git a/src/relay/backend/contrib/ethosu/source_module.cc b/src/relay/backend/contrib/ethosu/source_module.cc index 61a880e17ffba..e3f48bc276175 100644 --- a/src/relay/backend/contrib/ethosu/source_module.cc +++ b/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 #include #include @@ -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& 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{this->func_name}; + }); } return PackedFunc(); } @@ -109,7 +119,7 @@ class EthosUModuleNode : public ModuleNode { private: String c_source; - Array 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 \n"; ss << "#include \n"; - ss << "#include \n"; ss << "#include \n"; ss << "#include \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(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(); diff --git a/tests/python/contrib/test_ethosu/infra.py b/tests/python/contrib/test_ethosu/infra.py index aeed64ad4aec3..8b0d3063a6964 100644 --- a/tests/python/contrib/test_ethosu/infra.py +++ b/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", diff --git a/tests/python/contrib/test_ethosu/test_codegen.py b/tests/python/contrib/test_ethosu/test_codegen.py index 2e0fb4b78cf1a..1944de5f94c05 100644 --- a/tests/python/contrib/test_ethosu/test_codegen.py +++ b/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 diff --git a/tests/python/contrib/test_ethosu/test_networks.py b/tests/python/contrib/test_ethosu/test_networks.py index a4be923e09dbe..a0acf12439503 100644 --- a/tests/python/contrib/test_ethosu/test_networks.py +++ b/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) diff --git a/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py b/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py index 479a1032453a7..1754dc385e056 100644 --- a/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py +++ b/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__])