This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e7109fce6b [Relax] Exclude R.null_value()-bound vars from
KillAfterLastUse (#20267)
e7109fce6b is described below
commit e7109fce6b1d708c7e854ea88cb6cefe701c671d
Author: Neo Chien <[email protected]>
AuthorDate: Fri Sep 11 03:00:23 2026 +0800
[Relax] Exclude R.null_value()-bound vars from KillAfterLastUse (#20267)
Hi Committers,
This PR addresses issue https://github.com/apache/tvm/issues/20200. Any
suggestions would be appreciated if you are available.
### Root Cause
`KillAfterLastUse` decides whether a variable is a legal
`R.vm.kill_object` target via `stored_in_vm_register`, which only
excludes constants and FuncType/ShapeType/PrimType vars. A variable
bound to `R.null_value()` has type `AnyType` - not in that exclusion
list - so the pass wrongly treats it as killable. But both `CodeGenVM`
and `CodeGenVMTIR` special-case `null_value` to bypass real VM
register/anylist-slot allocation entirely. In `CodeGenVMTIR`, this
mismatch hits an `ICHECK` requiring the kill target to resolve to an
`anylist_getitem()` call, causing a hard crash (`CodeGenVM`'s bytecode
backend happens to swallow it silently via a sentinel "void register,"
which is why only `exec_mode="compiled"` crashes).
### Solution
Track variables bound to `relax.null_value` in `CollectLastUsage`
(mirroring the existing constant_tensors_ pattern) and add them to the
`stored_in_vm_register` exclusion, so the pass never emits a kill for
them - closing the gap at its sources rather than patching either
codegen.
---------
Co-authored-by: cchung100m <[email protected]>
---
src/relax/transform/kill_after_last_use.cc | 30 +++++--
tests/python/relax/test_kill_after_last_use.py | 110 +++++++++++++++++++++++++
2 files changed, 133 insertions(+), 7 deletions(-)
diff --git a/src/relax/transform/kill_after_last_use.cc
b/src/relax/transform/kill_after_last_use.cc
index ce1d4b1892..78f93f8bfa 100644
--- a/src/relax/transform/kill_after_last_use.cc
+++ b/src/relax/transform/kill_after_last_use.cc
@@ -112,14 +112,22 @@ class CollectLastUsage : public ExprVisitor {
bool already_killed = visitor.killed_objects_.count(var);
// Currently, the VM requires that objects to be killed
- // objects only exist in VM registers. This requires
- // KillAfterLastUse to have more knowledge about the VM
- // implementation than should exist at this stage of lowering.
- // In the future, this may be handled more easily at the
- // CodeGenVM level.
+ // only exist in VM registers. This requires KillAfterLastUse
+ // to have more knowledge about the VM implementation than
+ // should exist at this stage of lowering. In the future,
+ // this may be handled more easily at the CodeGenVM level.
+ //
+ // Variables bound to `relax.null_value` are excluded for the
+ // same reason as constants: both CodeGenVM and CodeGenVMTIR
+ // special-case `null_value` to bypass register/anylist-slot
+ // allocation, so such a variable is never a valid target for
+ // R.vm.kill_object. It is currently the only operator either
+ // codegen special-cases this way; a new special case added to
+ // either codegen should be reflected here as well.
bool stored_in_vm_register =
- !(visitor.constant_tensors_.count(var) ||
var->ty.as<FuncTypeNode>() ||
- var->ty.as<ShapeTypeNode>() || var->ty.as<PrimTypeNode>());
+ !(visitor.constant_tensors_.count(var) ||
visitor.null_value_objects_.count(var) ||
+ var->ty.as<FuncTypeNode>() || var->ty.as<ShapeTypeNode>() ||
+ var->ty.as<PrimTypeNode>());
if (!is_output && !already_killed) {
if (visitor.storage_objects_.count(var)) {
@@ -156,6 +164,7 @@ class CollectLastUsage : public ExprVisitor {
void VisitBinding_(const VarBindingNode* binding, const CallNode* val)
override {
static const Op& vm_alloc_storage = Op::Get("relax.vm.alloc_storage");
static const Op& mem_alloc_storage = Op::Get("relax.memory.alloc_storage");
+ static const Op& null_value_op = Op::Get("relax.null_value");
static const Op& mem_kill_tensor = Op::Get("relax.memory.kill_tensor");
static const Op& mem_kill_storage = Op::Get("relax.memory.kill_storage");
@@ -163,6 +172,8 @@ class CollectLastUsage : public ExprVisitor {
if (val->op.same_as(vm_alloc_storage) ||
val->op.same_as(mem_alloc_storage)) {
storage_objects_.insert(binding->var.get());
+ } else if (val->op.same_as(null_value_op)) {
+ null_value_objects_.insert(binding->var.get());
} else if (val->op.same_as(mem_kill_tensor) ||
val->op.same_as(mem_kill_storage) ||
val->op.same_as(vm_kill_object)) {
TVM_FFI_ICHECK_EQ(val->args.size(), 1)
@@ -204,6 +215,11 @@ class CollectLastUsage : public ExprVisitor {
// R.builtin.kill_tensor called on them.
std::unordered_set<const VarNode*> constant_tensors_;
+ // Variables bound to `relax.null_value`, which do not occupy a VM
+ // register in either CodeGenVM or CodeGenVMTIR, and therefore must
+ // never be passed to R.vm.kill_object.
+ std::unordered_set<const VarNode*> null_value_objects_;
+
// Set of objects that already have a call node to kill them. Should not
have a duplicate
std::unordered_set<const VarNode*> killed_objects_;
diff --git a/tests/python/relax/test_kill_after_last_use.py
b/tests/python/relax/test_kill_after_last_use.py
index c69263977f..6247782c0d 100644
--- a/tests/python/relax/test_kill_after_last_use.py
+++ b/tests/python/relax/test_kill_after_last_use.py
@@ -15,12 +15,15 @@
# specific language governing permissions and limitations
# under the License.
+import pytest
+
import tvm
import tvm.relax
import tvm.testing
from tvm.relax.transform import KillAfterLastUse
from tvm.script import ir as I
from tvm.script import relax as R
+from tvm.testing import env
def test_basic():
@@ -102,5 +105,112 @@ def
test_track_usage_across_trivial_rebindings_in_match_cast():
tvm.ir.assert_structural_equal(Expected, After)
+def test_no_kill_for_null_value():
+ """R.null_value() must never be targeted by R.vm.kill_object
+
+ A variable bound to `R.null_value()` is never assigned a real VM
+ register/anylist slot by either CodeGenVM or CodeGenVMTIR (both
+ special-case `null_value` to a sentinel value instead).
+ KillAfterLastUse must therefore never insert `R.vm.kill_object`
+ for such a variable, even though its type is `R.Any` (the same
+ type used by legitimate killable objectws such as VM storage).
+ """
+
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor([16, 32], "float32")):
+ storage = R.memory.alloc_storage(R.shape([2048]), 0, "global",
"uint8")
+ y = R.memory.alloc_tensor(storage, 0, R.shape([16, 32]), "float32")
+ shape_heap: R.Any = R.null_value()
+ _dummy = R.call_packed("use_shape_heap", [shape_heap],
ty_args=(R.Tuple,))
+ z = R.add(x, y)
+ return z
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor([16, 32], "float32")):
+ storage = R.memory.alloc_storage(R.shape([2048]), 0, "global",
"uint8")
+ y = R.memory.alloc_tensor(storage, 0, R.shape([16, 32]), "float32")
+ _ = R.memory.kill_storage(storage)
+ shape_heap: R.Any = R.null_value()
+ _dummy = R.call_packed("use_shape_heap", [shape_heap],
ty_args=(R.Tuple,))
+ z = R.add(x, y)
+ _ = R.memory.kill_tensor(y)
+ return z
+
+ After = KillAfterLastUse()(Before)
+ tvm.ir.assert_structural_equal(Expected, After)
+
+
+def _assert_no_kill_of_null_value(func: tvm.relax.Function):
+ """Assert no R.vm.kill_object call in `func` targets a null_value()-bound
var
+
+ An `R.null_value()`-bound variable never occupies a VM register in
+ either CodeGenVM or CodeGenVMTIR, so passing one to
+ R.vm.kill_object is always invalid. Checking this structurally
+ (rather than only checking that relax.build succeeds) ensures the
+ test fails if a future change merely makes codegen tolerant of the
+ invalid kill, instead of preventing KillAfterLastUse from
+ inserting it in the first place.
+ """
+ null_value_op = tvm.ir.Op.get("relax.null_value")
+ kill_object_op = tvm.ir.Op.get("relax.vm.kill_object")
+
+ null_value_vars = []
+ killed_args = []
+
+ body = func.body
+ assert isinstance(body, tvm.relax.SeqExpr)
+ for block in body.blocks:
+ for binding in block.bindings:
+ value = binding.value
+ if not isinstance(value, tvm.relax.Call):
+ continue
+ if value.op.same_as(null_value_op):
+ null_value_vars.append(binding.var)
+ elif value.op.same_as(kill_object_op):
+ killed_args.append(value.args[0])
+
+ for killed in killed_args:
+ for null_value_var in null_value_vars:
+ assert not killed.same_as(null_value_var), (
+ f"R.vm.kill_object was called on a variable bound to
R.null_value(): {killed}"
+ )
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_reapply_after_default_pipeline_builds_successfully():
+ """KillAfterLastUse may be re-applied to an already-lowered module
+
+ Applying KillAfterLastUse a second time, on top of the output of
+ `relax.get_default_pipeline` (which itself ends with a
+ KillAfterLastUse application, after VMShapeLower has introduced a
+ `shape_heap: R.Any = R.null_value()` binding), must not insert an
+ invalid `R.vm.kill_object(shape_heap)`, and the resulting module
+ must still build successfully under every exec_mode.
+ """
+
+ @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
+
+ target = tvm.target.Target("llvm")
+ preoptimized = tvm.relax.get_default_pipeline(target)(Mod)
+ second_kill = KillAfterLastUse()(preoptimized)
+
+ # Prove the invalid kill is gone, not merely that codegen tolerates it.
+ _assert_no_kill_of_null_value(second_kill["main"])
+
+ for exec_mode in ["bytecode", "compiled"]:
+ tvm.relax.build(second_kill, target=target, relax_pipeline="zero",
exec_mode=exec_mode)
+
+
if __name__ == "__main__":
tvm.testing.main()
