diff --git a/src/relax/transform/kill_after_last_use.cc b/src/relax/transform/kill_after_last_use.cc index ce1d4b1892dd..78f93f8bfacc 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() || - var->ty.as() || var->ty.as()); + !(visitor.constant_tensors_.count(var) || visitor.null_value_objects_.count(var) || + var->ty.as() || var->ty.as() || + var->ty.as()); 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 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 null_value_objects_; + // Set of objects that already have a call node to kill them. Should not have a duplicate std::unordered_set 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 c69263977f79..a69fdeacb3d5 100644 --- a/tests/python/relax/test_kill_after_last_use.py +++ b/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()