Lunderberg commented on code in PR #16088:
URL: https://github.com/apache/tvm/pull/16088#discussion_r1385918622
##########
python/tvm/relax/transform/lazy_transform_params.py:
##########
@@ -188,24 +199,24 @@ def visit_var_(self, var: relax.Var) -> None:
return super().visit_var_(var)
def visit_var_binding_(self, binding: relax.VarBinding) -> None:
- if binding.var == self.out_tuple_var:
+ if self.fset_item is not None and binding.var == self.out_tuple_var:
Review Comment:
It looks like the handling of `set_item` and of `get_item` is largely
independent. Rather than having them as two features of a single mutator,
which requires the `self.fset_item is not None` check to be repeated
throughout, can we split it into two distinct mutators? One mutator replaces
access of parameters with `get_item`. The other mutator replaces binding of
output with `set_item`. If both features are desired, then both mutators are
applied.
##########
python/tvm/relax/transform/lazy_transform_params.py:
##########
@@ -225,10 +236,38 @@ class LazyTransformParams:
(Load the input to memory on demand, and immediately free it after the
last use.)
Note: ToNonDataflow() and RemovePurityTracking() should be invoked before
this pass.
+
+ Parameters
+ ----------
+ fget_item: str
+ The name of the get_item function.
+ fset_item: str
+ The name of the set_item function.
+ get_item_param: list of relax.Var
+ The parameters of the get_item function except index.
+ The given parameters will be placed before index.
+ For example, if get_item_param is [param1, param2], then the pass will
generate
+ call_packed(fget_item, [param1, param2, index])
+ set_item_param: list of relax.Var
+ The parameters of the set_item function except index and value.
+ The given parameters will be placed before index and value.
+ For example, if set_item_param is [param1, param2], then the pass will
generate
+ call_packed(fset_item, [param1, param2, index, value])
"""
+ def __init__(
+ self, fget_item="get_item", fset_item="set_item", get_item_param=[],
set_item_param=[]
Review Comment:
Using a mutable object as a default value can have unexpected consequences.
(e.g. Accidental mutation of an object, which is reused in the next call.)
Instead of having a list as the default value, it is better to have
`get_item_param=None`, then to construct the list within the body of the
function.
##########
python/tvm/relax/transform/lazy_transform_params.py:
##########
@@ -118,9 +119,15 @@ class LazyTransformParamsMutator(PyExprMutator):
The module to be transformed
"""
- def __init__(self, mod: IRModule = None) -> None:
+ def __init__(
+ self, fget_item, fset_item, get_item_param, set_item_param, mod:
IRModule = None
+ ) -> None:
super().__init__(mod)
self.mod = mod
+ self.fget_item = fget_item
+ self.get_item_param = get_item_param
Review Comment:
What is the purpose of the `get_item_param` and `set_item_param`? I don't
see any unit test with their intended use. Since the same parameters are
forwarded to all `get_item` and `set_item` calls, it seems the same effect
could be achieved by having a stateful callback, rather than passing an
argument which is then passed right back to you.
##########
python/tvm/relax/transform/lazy_transform_params.py:
##########
@@ -149,9 +156,13 @@ def transform(self, func: relax.Function) ->
relax.Function:
# Step 3. rewrite get item and set item
new_body = self.visit_expr(func.body)
- # Step 4. Find all shape parameters that should be retained as
- # parameters.
+ # Step 4. Add parameters of get_item and set_item (except index) to
the function.
params = []
+ for param in itertools.chain(self.get_item_param, self.set_item_param):
Review Comment:
Python's `*` expansion avoids the need for a loop.
```python
params = [*self.get_item_param, *self.set_item_param]
```
It's a bit more "magic", but avoids mutation of a variable after assignment.
##########
python/tvm/relax/transform/lazy_transform_params.py:
##########
@@ -118,9 +119,15 @@ class LazyTransformParamsMutator(PyExprMutator):
The module to be transformed
"""
- def __init__(self, mod: IRModule = None) -> None:
+ def __init__(
+ self, fget_item, fset_item, get_item_param, set_item_param, mod:
IRModule = None
Review Comment:
From the usage, the `get_item_param` and `set_item_param` are lists of
elements, so they should be plural `get_item_params` and `set_item_params`.
(If they are required overall. See other comment asking if they are.)
This would also be an advantage of breaking the functionality apart into two
separate mutators: A `LazyInput` mutator could have an `extra_params` argument,
and a `LazyOutput` mutator could have an `extra_params` argument, and both
would be clear from the context. With both mutations implemented in the same
class, if we wanted the name to indicate that these are additional arguments,
we'd need to make it be `extra_get_item_params` and `extra_set_item_params`,
which is a bit long for ease of use.
--
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]
