grant-arm commented on a change in pull request #9597:
URL: https://github.com/apache/tvm/pull/9597#discussion_r758278217
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -80,5 +63,50 @@ def _compile(ext_func):
# that can perform scheduling based on user inputs such as
# scratch memory size.
tir_mod, params = lower_to_tir(mod["main"], copy_constants())
- cmms, encoded_constants, scratch_size =
tir_to_cs_translator.translate(tir_mod, params)
- return cmms, encoded_constants, scratch_size
+
+ for idx in params.keys():
+ params[idx] = tvm.nd.array(params[idx])
+
+ primfunc = tir_mod["main"]
+ primfunc = primfunc.with_attr("global_symbol",
ext_func.attrs["global_symbol"])
+ primfunc = primfunc.with_attr("ethos-u.constants", params)
+ primfunc = primfunc.with_attr("ethos-u.input_size", input_size)
+ primfunc = primfunc.with_attr("ethos-u.output_size", output_size)
+ return primfunc
+
+
+@tvm._ffi.register_func("relay.ext.ethos-u.primfunc_to_artifact")
+def primfunc_to_artifact(primfunc: tvm.tir.PrimFunc) ->
util.CompilationArtifact:
+ """
+ This is hook for python-based lowering of TIR PrimFunc
+ that has undergone unified optimization to Compilation
+ Artifact destined for the microNPU.
+
+ Parameters
+ ----------
+ primfunc : tir.PrimFunc
+ TIR PrimFuncthat has undergone unified optimization
Review comment:
```suggestion
TIR PrimFunc that has undergone unified optimization
```
##########
File path: src/relay/backend/contrib/ethosu/source_module.cc
##########
@@ -79,7 +82,6 @@ class EthosUModuleNode : public ModuleNode {
*/
void SaveToFile(const std::string& file_name, const std::string& format)
final {
std::string fmt = GetFileFormat(file_name, format);
- LOG(INFO) << "format=" << fmt << ";;\n";
Review comment:
Thanks Manupa, I've been meaning to remove this for a while.
##########
File path: src/relay/backend/contrib/ethosu/source_module.cc
##########
@@ -41,34 +41,37 @@
#include <vector>
#include "../../../../runtime/file_utils.h"
+#include "utils.h"
namespace tvm {
namespace runtime {
+using CompilationArtifact = relay::contrib::ethosu::CompilationArtifact;
+
// 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.
+ * \brief The microNPU runtime module.
*
- * \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
+ * \param compilation_artifacts
+ * This is an array of CompilationArtifacts that is produced via
+ * lowering each PrimFunc to command stream. Here, those artifacts
+ * will be used to create the c-source.
*/
- 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_name = func_name_;
- cmms_hex = std::move(cmms_hex_);
- weights_bias_hex = std::move(weights_bias_hex_);
- scratch_size = scratch_size_->value;
- input_size = input_size_->value;
- output_size = output_size_->value;
- c_source = GenerateSource();
+ explicit EthosUModuleNode(Array<CompilationArtifact> compilation_artifacts)
+ : compilation_artifacts_(compilation_artifacts) {
+ c_source += "#include <stdio.h>\n";
+ c_source += "#include <stdlib.h>\n";
+ c_source += "#include <tvm/runtime/crt/module.h>\n";
+ c_source += "#include <tvm_ethosu_runtime.h>\n";
+ c_source += "\n";
Review comment:
```suggestion
c_source += "#include <tvm_ethosu_runtime.h>\n\n";
```
##########
File path: src/relay/backend/contrib/ethosu/source_module.cc
##########
@@ -41,34 +41,37 @@
#include <vector>
#include "../../../../runtime/file_utils.h"
+#include "utils.h"
namespace tvm {
namespace runtime {
+using CompilationArtifact = relay::contrib::ethosu::CompilationArtifact;
+
// 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.
+ * \brief The microNPU runtime module.
*
- * \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
+ * \param compilation_artifacts
+ * This is an array of CompilationArtifacts that is produced via
+ * lowering each PrimFunc to command stream. Here, those artifacts
+ * will be used to create the c-source.
*/
- 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_name = func_name_;
- cmms_hex = std::move(cmms_hex_);
- weights_bias_hex = std::move(weights_bias_hex_);
- scratch_size = scratch_size_->value;
- input_size = input_size_->value;
- output_size = output_size_->value;
- c_source = GenerateSource();
+ explicit EthosUModuleNode(Array<CompilationArtifact> compilation_artifacts)
+ : compilation_artifacts_(compilation_artifacts) {
+ c_source += "#include <stdio.h>\n";
+ c_source += "#include <stdlib.h>\n";
+ c_source += "#include <tvm/runtime/crt/module.h>\n";
+ c_source += "#include <tvm_ethosu_runtime.h>\n";
+ c_source += "\n";
+ for (const CompilationArtifact& compilation_artifact :
compilation_artifacts) {
+ c_source += GenerateSource(compilation_artifact);
+ c_source += "\n";
+ c_source += "\n";
Review comment:
```suggestion
c_source += "\n\n";
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
