[Unity][Transform] Handle symbolic variables in LambdaLift #16411
Conversation
Lunderberg
force-pushed
the
unity_lambda_lift_symbolic_vars
branch
from
b9209cb
to
ec4a2ed
Compare
src/relax/transform/lambda_lift.cc
Outdated
| Function(params, body, func_node->ret_struct_info, func_node->is_pure, func_node->attrs); | ||
| // recursive call | ||
| if (is_recursive && is_closure) { | ||
| // it is required by block_blocker, will be updated later |
There was a problem hiding this comment.
It looks like this typo was present in the original, but I presume this comment should refer to the BlockBuilder
There was a problem hiding this comment.
Thank you for the catch. The typo was present in the original, and the comment is also out of date after this change. I've removed it, and added an appropriate comment.
src/relax/transform/lambda_lift.cc
Outdated
| auto cache = current_lambda_var_; | ||
| current_lambda_var_ = binding->var; | ||
|
|
||
| // ExprMutator::VisitBinding_(binding, func_node); |
There was a problem hiding this comment.
Probably subjective, but this comment is a little more cryptic compared to just pointing out that we are visiting a function literal
There was a problem hiding this comment.
Thank you for the catch, and the comment was some test code, and is now removed.
| func->attrs); | ||
| builder_->UpdateFunction(pair.first, func); | ||
| for (auto [gvar, base_func] : glob_funcs) { | ||
| if (auto opt = base_func.as<Function>()) { |
There was a problem hiding this comment.
Cool, I didn't know about using as for ObjectRefs directly 👍
There was a problem hiding this comment.
Thank you! This was was added back in #14522, and I really like the way it avoids the repetition in if(obj->IsInstance<TNode>()) checks and Downcast<T>(obj) casts inside the conditional.
src/relax/transform/lambda_lift.cc
Outdated
| struct Context { | ||
| explicit Context(int* ptr) : ptr(ptr) { (*ptr)++; } | ||
| ~Context() { (*ptr)--; } | ||
| int* ptr; | ||
| } context(&depth); |
There was a problem hiding this comment.
This doesn't appear to be used anywhere.
There was a problem hiding this comment.
Whoops, debug code for indenting prints during nested calls. Removed.
src/relax/transform/lambda_lift.cc
Outdated
| if (bool is_closure = IsClosure(var); | ||
| is_closure && builder_->LookupBinding(var).as<CallNode>()) { |
There was a problem hiding this comment.
Is there a reason to write it this way instead of IsClosure(var) && ...?
There was a problem hiding this comment.
Nope. A previous implementation of IsClosure returned an Optional<_>, and this stuck around. Updated.
| // it is required by block_blocker, will be updated later | ||
| nested_closure_map_.emplace( | ||
| current_lambda_var_.value(), | ||
| Call(gvar_lifted_func, captured_vars.Map([](Var var) -> Expr { return var; }))); |
There was a problem hiding this comment.
I assume the map is so that an Array<Var> is treated as Array<Expr>, correct?
There was a problem hiding this comment.
Yup, that's correct. I'm wondering if there should be an implicit conversion from Array<DerivedClass> to Array<BaseClass>, to avoid needing this type of conversion.
There was a problem hiding this comment.
Might have to be careful to ensure type safety can't get broken that way, that has tended to lead to soundness issues in languages. (I believe that type of thing is a source of unsoundness in TypeScript, IIRC.)
| // Must visit the function itself, and not just the function | ||
| // body, to ensure that EraseToWellDefined recognized symbolic | ||
| // variables that are exposed by the function signature. | ||
| auto func = Downcast<Function>(VisitExpr(opt.value())); | ||
| builder_->UpdateFunction(gvar, func); |
There was a problem hiding this comment.
Is this the only change needed to handle symbolic vars, fundamentally? I ask because I didn't see any logic specific to symbolic vars above.
There was a problem hiding this comment.
Primarily, yes. Most of the other changes are to allow this specific change to work. Previously, the mutator could assume that every FunctionNode encountered would be a lambda to be lifted out, but now it could be a top-level function as well.
Prior to this commit, symbolic variables used by a lambda function would be duplicated between the caller and the lifted-out function. In addition, shape inference within the lifted-out function was performed without access to the symbolic variables, resulting in unnecessary fallback from `R.Tensor([m, n])` to `R.Tensor(ndim=2)`. This commit updates the `LambdaLift` transform to handle symbolic variables. All symbolic variables have unique definitions across the resulting `IRModule`, and shape inference in the lifted-out function is aware of symbolic variables that have been exposed to it.
Lunderberg
force-pushed
the
unity_lambda_lift_symbolic_vars
branch
from
ec4a2ed
to
43b7f8c
Compare
Prior to this commit, symbolic variables used by a lambda function would be duplicated between the caller and the lifted-out function. In addition, shape inference within the lifted-out function was performed without access to the symbolic variables, resulting in unnecessary fallback from
R.Tensor([m, n])toR.Tensor(ndim=2).This commit updates the
LambdaLifttransform to handle symbolic variables. All symbolic variables have unique definitions across the resultingIRModule, and shape inference in the lifted-out function is aware of symbolic variables that have been exposed to it.