You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/12/01 02:08:11 UTC
[GitHub] [tvm] huajsj commented on a change in pull request #9597: [microNPU] Move the compilation to use Target Hooks.
huajsj commented on a change in pull request #9597:
URL: https://github.com/apache/tvm/pull/9597#discussion_r759767264
##########
File path: include/tvm/tir/transform.h
##########
@@ -287,6 +287,11 @@ TVM_DLL Pass LowerThreadAllreduce();
*/
TVM_DLL Pass InferFragment();
+/*!
+ * \brief This annotation for nodes to be disabled for builtin lowering
Review comment:
"* \brief This annotation is for nodes where you want to disable the built-in lowering."
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -52,25 +35,25 @@ def constant_updater(expr, symbol): # pylint: disable=unused-argument
return dict()
-def _compile(ext_func):
+@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir_func")
Review comment:
This name is dazzling, how about use "camel-case" like what you did for relay.ext.ethos-u.CompilationArtifact?
##########
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:
The TIR PrimFunc which has undergone the unified optimization.
##########
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
+
+ Returns
+ -------
+ CompilationArtifact
+ This is a structure that holds the binary artifacts
+ for the microNPU
+ """
+ symbol = str(primfunc.attrs["global_symbol"])
+ params = primfunc.attrs["ethos-u.constants"]
+ input_size = primfunc.attrs["ethos-u.input_size"]
+ output_size = primfunc.attrs["ethos-u.output_size"]
+ tir_mod = tvm.IRModule()
+ tir_mod[symbol] = primfunc
+
+ params_with_int_keys = dict()
+ for idx in params.keys():
+ params_with_int_keys[int(idx)] = params[idx].numpy()
Review comment:
is it possible that the key "idx" is a "str"?
##########
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)
Review comment:
ethos-u.constants is little confused because params is a dict instead of constant, should it be "params"?
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -52,25 +35,25 @@ def constant_updater(expr, symbol): # pylint: disable=unused-argument
return dict()
-def _compile(ext_func):
+@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir_func")
+def relay_to_tir_func(ext_func: relay.Function) -> tvm.tir.PrimFunc:
"""
- This is the main wrapper that accepts an external
- relay function and runs all the passes to lower it down
- to command stream
+ This is hook for python-based lowering of relay function
+ that gets offloaded to the microNPU.
Review comment:
"This is a hook for the python-based lowering of relay function that gets offloaded to the microNPU."
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -52,25 +35,25 @@ def constant_updater(expr, symbol): # pylint: disable=unused-argument
return dict()
-def _compile(ext_func):
+@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir_func")
+def relay_to_tir_func(ext_func: relay.Function) -> tvm.tir.PrimFunc:
"""
- This is the main wrapper that accepts an external
- relay function and runs all the passes to lower it down
- to command stream
+ This is hook for python-based lowering of relay function
+ that gets offloaded to the microNPU.
+
Parameters
----------
- ext_func : tvm.relay.function.Function
- The partitioned relay function
+ ext_func : relay.Function
+ This is the partitioned relay function
+
Returns
-------
- cs : str
- An hex string of the bytes of command stream
- encoded_constants : str
- An hex string of the bytes that includes concat'd
- encoded weights, encoded biases and scales.
- scratch_size : int
- The size of the scratch buffer needed.
+ primfunc : tir.PrimFunc
+ This returns the scheduled PrimFunc
"""
+ assert len(ext_func.params) == 1
Review comment:
why len(ext_func.params) only can be 1, is this a limitation?
##########
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")
Review comment:
"camel-case" for primfunc_to_artifact?
##########
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])
Review comment:
```
for idx, val in params.items():
params[idx] = tvm.nd.array(val)
```
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -52,25 +35,25 @@ def constant_updater(expr, symbol): # pylint: disable=unused-argument
return dict()
-def _compile(ext_func):
+@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir_func")
+def relay_to_tir_func(ext_func: relay.Function) -> tvm.tir.PrimFunc:
"""
- This is the main wrapper that accepts an external
- relay function and runs all the passes to lower it down
- to command stream
+ This is hook for python-based lowering of relay function
+ that gets offloaded to the microNPU.
+
Parameters
----------
- ext_func : tvm.relay.function.Function
- The partitioned relay function
+ ext_func : relay.Function
+ This is the partitioned relay function
+
Returns
-------
- cs : str
- An hex string of the bytes of command stream
- encoded_constants : str
- An hex string of the bytes that includes concat'd
- encoded weights, encoded biases and scales.
- scratch_size : int
- The size of the scratch buffer needed.
+ primfunc : tir.PrimFunc
+ This returns the scheduled PrimFunc
"""
+ assert len(ext_func.params) == 1
Review comment:
better to use RuntimeError with a error report message.
##########
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.
Review comment:
"This is a hook for the python-based lowering of the TIR PrimFunc which has undergone the unified optimization for the microNPU"
not sure if the said sentence match with original meaning, but at least need to add "a", "the" and lowercase the first letter of "Compilation" and "Artifact".
##########
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
+
+ Returns
+ -------
+ CompilationArtifact
+ This is a structure that holds the binary artifacts
+ for the microNPU
+ """
+ symbol = str(primfunc.attrs["global_symbol"])
+ params = primfunc.attrs["ethos-u.constants"]
+ input_size = primfunc.attrs["ethos-u.input_size"]
+ output_size = primfunc.attrs["ethos-u.output_size"]
+ tir_mod = tvm.IRModule()
+ tir_mod[symbol] = primfunc
+
+ params_with_int_keys = dict()
+ for idx in params.keys():
+ params_with_int_keys[int(idx)] = params[idx].numpy()
Review comment:
```
for idx, val in params.items():
params_with_int_keys[int(idx)] = val.numpy()
```
--
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