[Relax] Exclude R.null_value()-bound vars from KillAfterLastUse

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,13 +164,16 @@ 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");
     static const Op& vm_kill_object = Op::Get("relax.vm.kill_object");
 
     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_;
 

tests/python/relax/test_kill_after_last_use.py

@@ -102,5 +102,111 @@ def main(w: R.Tensor([16, 32], "float32")):
     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}"
+            )
+
+
+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()