Skip to content

[Relax][VM] Re-running KillAfterLastUse on a default-lowered module crashes compiled VM codegen[Bug] #20200

Description

@lrcyyds1

Expected behavior

The compiled VM build should handle a Relax module produced by the default CPU pipeline after KillAfterLastUse is applied again, or report a clear diagnostic if this IR is outside the accepted VM-lowered
input contract.

In the reproducer below, the ordinary one-shot compiled build succeeds, and the same twice-lowered module is accepted by the bytecode VM build.

Actual behavior

The compiled VM build fails in CodeGenVMTIR::EmitKillObject:

tvm.error.InternalError: Check failed: (tir_call->op == tirx::builtin::anylist_getitem()) is false:

The failing stack reaches:

python/tvm/relax/vm_build.py, line 78, in _vmcodegen
  return _ffi_api.VMTIRCodeGen(builder, mod)

src/relax/backend/vm/codegen_vm_tir.cc, line 245, in VisitExpr_(CallNode)
  dst_reg = EmitKillObject(call);

src/relax/backend/vm/codegen_vm_tir.cc, line 436, in EmitKillObject
  TVM_FFI_ICHECK(tir_call->op == tirx::builtin::anylist_getitem());

Bisecting the second CPU default pipeline shows that the first pass causing the compiled build failure is relax.transform.KillAfterLastUse(). It inserts an additional R.vm.kill_object(shape_heap) into a
module already produced by the default pipeline.

After that pass:

- relax.build(..., relax_pipeline="zero", exec_mode="bytecode") succeeds.
- relax.build(..., relax_pipeline="zero", exec_mode="compiled") fails with the anylist_getitem assertion above.

I also checked seven small high-level Relax modules (add, multiply, reshape, permute, matmul, where, and concat variants). Reapplying VM finalization passes other than KillAfterLastUse did not find another
signature; reapplying KillAfterLastUse reproduced this same compiled-VM assertion in all seven modules.

### Environment

- OS: Linux
- Target: llvm
- TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05
- Commit subject: [Relax][Frontend][TFLite] Support static hashtable find (#19879)

### Steps to reproduce

import tvm
from tvm import relax
from tvm.script import from_source

src = r'''
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Mod:
    @R.function
    def main(x: R.Tensor((1, 4), "float32")) -> R.Tensor((1, 4), "float32"):
        R.func_attr({"global_symbol": "main", "num_input": 1})
        y: R.Tensor((1, 4), "float32") = R.add(x, R.const(1.0, "float32"))
        z: R.Tensor((1, 4), "float32") = R.add(y, R.const(2.0, "float32"))
        return z
'''

mod = from_source(src)
target = tvm.target.Target("llvm")
default_pipeline = relax.get_default_pipeline(target)

# Control: the ordinary one-shot compiled build succeeds.
relax.build(mod, target=target, relax_pipeline=default_pipeline, exec_mode="compiled")
print("one-shot default compiled build: ok")

preoptimized = default_pipeline(mod)
relax.build(preoptimized, target=target, relax_pipeline="zero", exec_mode="compiled")
print("preoptimized + zero compiled build: ok")

# Running KillAfterLastUse again inserts R.vm.kill_object(shape_heap).
second_kill = relax.transform.KillAfterLastUse()(preoptimized)
print("kill_object_count_after_second_kill:", second_kill.script(show_meta=False).count("kill_object"))

# Bytecode VM accepts the result, but compiled VM crashes in CodeGenVMTIR::EmitKillObject.
relax.build(second_kill, target=target, relax_pipeline="zero", exec_mode="bytecode")
print("second KillAfterLastUse + bytecode build: ok")
relax.build(second_kill, target=target, relax_pipeline="zero", exec_mode="compiled")

Observed output before the final failure:

one-shot default compiled build: ok
preoptimized + zero compiled build: ok
kill_object_count_after_second_kill: 4
second KillAfterLastUse + bytecode build: ok

Then the final compiled build fails with:

tvm.error.InternalError: Check failed: (tir_call->op == tirx::builtin::anylist_getitem()) is false:

During pass bisection, the failing module contained:

shape_heap: R.Object = R.null_value()
R.call_packed("vm.builtin.match_shape", x, shape_heap, ...)
R.vm.kill_object(shape_heap)

### Triage

- needs-triage
- relax
- vm

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

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions