[Bug] AllocateWorkspace returns an ill-formed module with a dangling GlobalVar #20479

Thanks for participating in the TVM community! We use https://discuss.tvm.apache.org/ for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first 😸

Issues that are inactive for a period of time may get closed. We adopt this policy so that we won't lose track of actionable issues that may fall at the bottom of the pile. Feel free to reopen a new one if you feel there is an additional problem that needs attention when an old one gets closed.

Expected behavior

relax.transform.AllocateWorkspace should return an IRModule in which every referenced GlobalVar is defined.

In particular:

• relax.analysis.well_formed should pass.
• relax.build should succeed.
• relax.VirtualMachine construction should succeed.

Actual behavior

The returned module calls the old GlobalVar of each function renamed by the pass, even though that GlobalVar has been removed.

As a result, well_formed fails, and VirtualMachine construction aborts:

=== after AllocateWorkspace ===
global vars: ['inner1', 'inner_tir', 'main', 'outer1']

ValueError: GlobalVar I.GlobalVar("inner") is not defined.

InternalError: Check failed: (func.has_value()) is false: Error: Cannot find ffi::Function inner
in either Relax VM kernel library, or in TVM runtime ffi::Function registry, or in global
Relax functions of the VM executable

AllocateWorkspace itself returns successfully, and relax.build also succeeds.

The failure surfaces only when constructing the VirtualMachine.

No BYOC backend is installed or registered — the backend name is never resolved.

The module is already malformed before any code generation runs.

Environment

• TVM version: 0.26.dev0
• Commit: 8312a17f8734ddfd56e5f3977cd5df25b83ec49f
• OS: Linux x86_64, Ubuntu
• Kernel: 5.15
• Python: 3.11
• Target: llvm
• Device: CPU only
• External codegen backend: Not required
• Reproduces with: relax.transform.AllocateWorkspace alone

Steps to reproduce

import tvm
from tvm import relax
from tvm.script import ir as I, tirx as T, relax as R


@I.ir_module
class Module:
    @T.prim_func
    def inner_tir(
        x: T.Buffer((4,), "float32"),
        y: T.Buffer((4,), "float32"),
    ):
        for i in T.serial(4):
            y[i] = x[i] + T.float32(1)

    @R.function
    def main(
        x: R.Tensor((4,), dtype="float32"),
    ) -> R.Tensor((4,), dtype="float32"):
        cls = Module
        with R.dataflow():
            gv = cls.outer(x)
            R.output(gv)
        return gv

    @R.function
    def outer(
        x: R.Tensor((4,), dtype="float32"),
    ) -> R.Tensor((4,), dtype="float32"):
        R.func_attr({"Codegen": "my_backend"})
        cls = Module
        with R.dataflow():
            gv = cls.inner(x)
            R.output(gv)
        return gv

    @R.function
    def inner(
        x: R.Tensor((4,), dtype="float32"),
    ) -> R.Tensor((4,), dtype="float32"):
        R.func_attr(
            {
                "Composite": "my_backend.op",
                "Primitive": 1,
                "WorkspaceSize": 1024,
            }
        )
        cls = Module
        with R.dataflow():
            gv = R.call_tir(
                cls.inner_tir,
                (x,),
                out_ty=R.Tensor((4,), dtype="float32"),
            )
            R.output(gv)
        return gv


print("=== before ===")
relax.analysis.well_formed(Module)  # passes

after = relax.transform.AllocateWorkspace()(Module)

print("=== after AllocateWorkspace ===")
print(sorted(str(g.name_hint) for g in after.get_global_vars()))

relax.analysis.well_formed(after)  # fails

ex = relax.build(after, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())

The expected output is the block shown under Actual behavior.

The module is already malformed before any codegen runs.

Triage

• needs-triage
• bug