[Unity][Transform] Update LambdaLift to use name of lifted lambda #16306
Conversation
Prior to this commit, the `LambdaLift` pass named each function as `"lifted_func_" + i`, in incremental order of occurrence. This provided unique names for each function, but could be difficult to read, or to refer to the lifted functions. This commit updates the naming scheme to use the location at which the lifted lambda occurs to generate a unique name for the new `GlobalVar`.
Pull in unity head to get the bugfix from apache#16322
src/relax/transform/lambda_lift.cc
Outdated
| std::unordered_set<String> unavailable_names = previous_global_vars_; | ||
|
|
||
| // A helper function to check de-duplicate names | ||
| auto use_if_unique = [&](const auto& generate_proposed_name) { |
There was a problem hiding this comment.
This might be a bit of a nitpick but I don't think use_if_unique is necessarily a fully descriptive name. What is important to convey is that it takes a scheme for generating names and applies it to all remaining lambdas. It had me scratching my head until I looked through the implementation. Not strictly necessary to change but it might be easier to understand at a glance. Possible alternative: attempt_generate_names?
There was a problem hiding this comment.
Good call on the naming. I've updated the name of use_if_unique to attempt_name_generation, the argument from generate_proposed_name to proposed_name_generation_func, and added a comment describing the intended arguments and return value for proposed_name_generation_func.
slyubomirsky
left a comment
slyubomirsky
left a comment
There was a problem hiding this comment.
Very comprehensive change and it should lead to more human-readable code. Thanks!
Thank you! Because I tend to run into errors late in a lowering pipeline, having the names preserved throughout the lowering pipeline, whenever possible, helps out debugging quite a bit. |
| if (Optional<String> opt_proposed_name = proposed_name_generation_func(func, location)) { | ||
| auto proposed_name = opt_proposed_name.value(); | ||
|
|
||
| if (unavailable_names.count(proposed_name)) { |
There was a problem hiding this comment.
the body of the if statement is empty, probably we can merge it with else if/ else
There was a problem hiding this comment.
Yeah, I went back and forth on it. It could be merged with the else if/else, but that would require repeating (and re-evaluating) the condition in the other two branches. It could be pulled out as if(!unavailable_names.count(proposed_name)) to wrap around the other two cases, but this utility is already getting a bit deeply nested for readability.
|
|
||
| @R.function(private=True) | ||
| def lifted_func_1( | ||
| def glob_func_2_inner( |
There was a problem hiding this comment.
how about adding a test case to cover the dedup of possible duplicate lifted names like:
@ir_module
class Before:
@R.function
def foo_foo(x1: R.Tensor):
@R.function
def foo(x2: R.Tensor):
...
@R.function
def foo(x1: R.Tensor):
@R.function
def foo_foo(x2: R.Tensor):
...
There was a problem hiding this comment.
Good call, and added.
The `relax.transform.LiftTransformParams` pass splits apart a relax function, extracting the steps that could be performed at compile-time. Prior to this commit, the transformed parameters were named `param0`, `param1`, and so on. This commit updates the `LiftTransformParams` pass to preserve any human-readable parameter names. The parameter names for the updated function are taken from the original parameter names, if no transformation is performed, or from the internal variable binding, if a transformation is applied. This implementation uses `LambdaLift` internally, relying on the changes made in apache#16306.
* [Relax][Transform] Preserve param names in LiftTransformParams The `relax.transform.LiftTransformParams` pass splits apart a relax function, extracting the steps that could be performed at compile-time. Prior to this commit, the transformed parameters were named `param0`, `param1`, and so on. This commit updates the `LiftTransformParams` pass to preserve any human-readable parameter names. The parameter names for the updated function are taken from the original parameter names, if no transformation is performed, or from the internal variable binding, if a transformation is applied. This implementation uses `LambdaLift` internally, relying on the changes made in #16306. * Update based on review comments
3 participants
Prior to this commit, the
LambdaLiftpass named each function as"lifted_func_" + i, in incremental order of occurrence. This provided unique names for each function, but could be difficult to read, or to refer to the lifted functions. This commit updates the naming scheme to use the location at which the lifted lambda occurs to generate a unique name for the newGlobalVar.