Skip to content

[Bug][Relax] LambdaLift crashes when a captured local function is called through an alias #20497

Description

@Yuhx141

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

  • needs-triage
  • type: bug

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions