Description
LambdaLift handles a direct call to a captured local function, but crashes when the same function
is first assigned to a variable alias and then called. The source module is well formed. The pass
recognizes the aliased value as a closure, but closure-call lowering also requires the called
variable's immediate binding to be a CallNode; for an alias, that binding is another Var.
Environment
- Ubuntu 24.04.4 LTS, x86-64
- TVM
0.26.dev0
- Runtime reproduction:
e269315c90e3a061c9e1c77b370ce883b1b223f4
- LLVM 18.1.3, CPU bytecode VM
Current upstream main was 2393db61bbd0100f048b1a6d3c2519fb70914564 when checked on
2026-09-30. src/relax/transform/lambda_lift.cc has changed elsewhere, but the same immediate-
binding CallNode guard remains. I did not have a binary build for that exact main revision, so
the isolated pass run below is from the stated fixed revision.
Minimal reproducer
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def main(
x: R.Tensor((2,), "float32"),
y: R.Tensor((2,), "float32"),
):
@R.function
def inner(z: R.Tensor((2,), "float32")):
return R.add(z, x)
alias = inner
return alias(y)
assert relax.analysis.check_well_formed(Module, check_ty=True)
relax.transform.LambdaLift()(Module) # raises InternalError
Observed error:
InternalError: Check failed: (opt) is false: Call->op must contains a function type
Expected behavior
LambdaLift should preserve the captured x and lower the aliased call to the same closure
invocation used for a direct call.
Actual behavior
The source passes check_well_formed(..., check_ty=True), but the isolated pass raises before
producing a target module. Tensor extents 1 through 6 all reproduce the same error.
Removing only alias = inner and returning inner(y) gives six passing controls. The controls are
lifted successfully, the targets are well formed, and their LLVM VM outputs match flat reference
functions.
Modules containing local functions cannot be executed directly by this VM build because
VMShapeLower requires LambdaLift first. I therefore use only the lifted control targets and flat
reference functions as runtime oracles.
Suspected cause
LambdaLifter::VisitExpr_(const CallNode*) contains:
auto bound_value = LookupBinding(var);
if (IsClosure(var) && bound_value.as<CallNode>()) {
// emit invoke_pure_closure / invoke_closure
}
IsClosure(var) follows variable bindings and recognizes the aliased closure, but the immediate
binding of alias is a Var, not a CallNode. The call is therefore not rewritten as a closure
invocation, and rebuilding it later fails because the call operator no longer has a function type.
Duplicate check
I searched existing TVM issues and pull requests for LambdaLift, closure aliases, indirect calls,
and the exact error. I found other LambdaLift and VM-shape reports, but none with the same cause:
closure identity stopping at an aliased call site.
I can send a fix PR and help with follow-up testing if this diagnosis looks right.
Triage
Description
LambdaLifthandles a direct call to a captured local function, but crashes when the same functionis first assigned to a variable alias and then called. The source module is well formed. The pass
recognizes the aliased value as a closure, but closure-call lowering also requires the called
variable's immediate binding to be a
CallNode; for an alias, that binding is anotherVar.Environment
0.26.dev0e269315c90e3a061c9e1c77b370ce883b1b223f4Current upstream
mainwas2393db61bbd0100f048b1a6d3c2519fb70914564when checked on2026-09-30.
src/relax/transform/lambda_lift.cchas changed elsewhere, but the same immediate-binding
CallNodeguard remains. I did not have a binary build for that exactmainrevision, sothe isolated pass run below is from the stated fixed revision.
Minimal reproducer
Observed error:
Expected behavior
LambdaLiftshould preserve the capturedxand lower the aliased call to the same closureinvocation used for a direct call.
Actual behavior
The source passes
check_well_formed(..., check_ty=True), but the isolated pass raises beforeproducing a target module. Tensor extents 1 through 6 all reproduce the same error.
Removing only
alias = innerand returninginner(y)gives six passing controls. The controls arelifted successfully, the targets are well formed, and their LLVM VM outputs match flat reference
functions.
Modules containing local functions cannot be executed directly by this VM build because
VMShapeLowerrequiresLambdaLiftfirst. I therefore use only the lifted control targets and flatreference functions as runtime oracles.
Suspected cause
LambdaLifter::VisitExpr_(const CallNode*)contains:IsClosure(var)follows variable bindings and recognizes the aliased closure, but the immediatebinding of
aliasis aVar, not aCallNode. The call is therefore not rewritten as a closureinvocation, and rebuilding it later fails because the call operator no longer has a function type.
Duplicate check
I searched existing TVM issues and pull requests for
LambdaLift, closure aliases, indirect calls,and the exact error. I found other LambdaLift and VM-shape reports, but none with the same cause:
closure identity stopping at an aliased call site.
I can send a fix PR and help with follow-up testing if this diagnosis looks right.
Triage