[Unity] make LazyTransformParam more general #16088
Conversation
jinhongyii
commented
Nov 7, 2023
jinhongyii
commented
- user can now customize the name of packed func get_item and set_item
- user can now only rewrite get_item but not set_item. This is useful when we want a tuple as output.
| super().__init__(mod) | ||
| self.mod = mod | ||
| self.fget_item = fget_item | ||
| self.get_item_param = get_item_param |
There was a problem hiding this 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.
There was a problem hiding this comment.
My bad. I forget to add a test on this. The case is that I need to call get_item(loader, index). Since shard loader is created at runtime and is a DRef object, it cannot be a stateful callback and I have to modify the interface of get_item. Does this make sense to you?
| @@ -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: | |||
There was a problem hiding this 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.
There was a problem hiding this comment.
sounds good. let me try
| @@ -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 | |||
There was a problem hiding this 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.
There was a problem hiding this comment.
I cannot fully separate these mutators because they need to share input_tuple_param and func.params[0] will change after one mutator.
There was a problem hiding this comment.
so what I can do is create a class that does all the analysis, calls into the two mutators, and create function. In this case, we have to use extra_get_item_params and extra_set_item_params
There was a problem hiding this comment.
Hmm, looks like the analysis would need to be split apart as well. In general, I think having multiple smaller changes is better than a single mutator, but this seems like a refactor than would make sense to ask in a feature addition PR. The current revision looks good, and any additional changes can be in a follow-up PR if desired.
|
@Lunderberg I've tried my best to address your comment, please review again |
|
Looks good, and approved. Current CI failures are due to use of |
