jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586979630
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel,
build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided
in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
By the way, my colleague is going to add Ansor support for
sparse_conv2d. We'll add extra check to this entry first, and see if there's
any better way to merge them.
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel,
build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided
in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
Yeah, I've also had some discussions in our weekly sync while didn't
figure out any better solutions.
There're several reasons:
1. Different ops have different requirements over specific inputs;
2. While the problems is in a subgraph generated in Relay Integration, the
placeholder are all the same, we can not differentiate them from tag, name or
any other way, even the order of inputs are not guaranteed.
Current approach is to merge all specific inputs checking to this function,
at least they have a same entry here. For the other ops, you have to add their
own check functions below.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
