You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2023/10/30 19:17:57 UTC
(tvm) branch unity updated: [Unity] Remove end-of-life handling from StaticPlanBlockMemory (#15841)
This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 6936829aac [Unity] Remove end-of-life handling from StaticPlanBlockMemory (#15841)
6936829aac is described below
commit 6936829aacdb4fb403bc055169a5c751cf010b8e
Author: Eric Lunderberg <Lu...@users.noreply.github.com>
AuthorDate: Mon Oct 30 14:17:52 2023 -0500
[Unity] Remove end-of-life handling from StaticPlanBlockMemory (#15841)
The catch-all in `relax.transform.KillAfterLastUse` can be used for
releasing tensor/storage, not just those that were not previously
handled. This simplifies the `StaticPlanBlockMemory` pass, as it does
not need to track the location at which `R.memory.kill_storage` and
`R.memory.kill_tensor` should be inserted.
---
src/relax/transform/rewrite_cuda_graph.cc | 3 +-
src/relax/transform/static_plan_block_memory.cc | 89 +++++---------------
.../test_transform_static_plan_block_memory.py | 98 +++++-----------------
3 files changed, 43 insertions(+), 147 deletions(-)
diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc
index 402dd55545..719703a3ec 100644
--- a/src/relax/transform/rewrite_cuda_graph.cc
+++ b/src/relax/transform/rewrite_cuda_graph.cc
@@ -222,14 +222,13 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final {
static const auto& mem_alloc_storage_op = Op::Get("relax.memory.alloc_storage");
- static const auto& mem_kill_storage_op = Op::Get("relax.memory.kill_storage");
static const auto& builtin_alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor");
static const auto& call_builtin_with_ctx_op = Op::Get("relax.call_builtin_with_ctx");
if (call->op.same_as(mem_alloc_storage_op) && IsStaticAllocStorage(binding)) {
AddStaticBinding(binding, /*is_alloc_storage=*/true);
return;
- } else if (call->op.same_as(mem_kill_storage_op) || call->op.same_as(builtin_alloc_tensor_op)) {
+ } else if (call->op.same_as(builtin_alloc_tensor_op)) {
return;
}
diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc
index ef2d582548..4a2a1555ff 100644
--- a/src/relax/transform/static_plan_block_memory.cc
+++ b/src/relax/transform/static_plan_block_memory.cc
@@ -38,16 +38,19 @@
* alloc_tensor.
*
* The third stage is IR rewrite. Based on the decision made in the second
- * stage, we insert memory alloc_storage, alloc_tensor, kill_tensor, and
- * kill_storage accordingly. Specifically, we
- * - insert alloc_storage before the site that each storage token is firstly
- * used,
- * - insert memory alloc_tensor for each builtin alloc_tensor,
- * - insert kill_tensor after the site that a tensor created by alloc_tensor
- * is last referenced, and
- * - insert kill_storage at the end of each binding block, for all the storage
- * tokens that are allocated inside the binding block, as the memory planning
- * only works on block level.
+ * stage, we insert memory alloc_storage, alloc_tensor.
+ *
+ * - Insert `memory.alloc_storage` before first usage site of each
+ * storage token.
+ *
+ * - Insert `memory.alloc_tensor` at the site of the
+ * `builtin.alloc_tensor` that it replaces.
+ *
+ * We do not insert `memory.kill_storage` or `memory.kill_tensor`, as
+ * these are handled in the later `KillAfterLastUse` lowering pass.
+ * This ensures that all tensors are killed after their last use,
+ * including dynamically-sized tensors, without requiring that
+ * `StaticPlanBlockMemory` track these dynamic-sized tensors.
*
* The memory planning pass "supports" dynamic shape in the way of TIR variable
* upper bound annotation. To be more specific, we can annotate the attribute
@@ -549,14 +552,10 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor {
* token in the token pool we maintain.
* - For each VM builtin reshape, we reuse the input's tokens.
*
- * After the allocation planning, we
- * - know the token that each builtin alloc_tensor plans to use. Compared
- * with the initialization, here the token is possibly a reuse of some
- * previous token, rather than we having one token for each alloc_tensor.
- * - know the last referenced site for each builtin alloc_tensor. This
- * information is used for inserting kill_tensor in the rewrite stage.
- * - know the tokens allocated in each binding block. This information
- * is used for inserting kill_storage in the rewrite stage.
+ * After the allocation planning, we know the token that each builtin
+ * alloc_tensor plans to use. Compared with the initialization, here
+ * the token is possibly a reuse of some previous token, rather than
+ * we having one token for each alloc_tensor.
*/
class StorageAllocator : public StorageAllocatorBaseVisitor {
public:
@@ -579,8 +578,6 @@ class StorageAllocator : public StorageAllocatorBaseVisitor {
* underlying storage token that it is using.
*/
std::unordered_map<const ExprNode*, StorageToken> alloc_tensor2token;
- /*! \brief The mapping from each Expr to the tensors that need to be killed after it. */
- std::unordered_map<const ExprNode*, std::vector<Var>> expr2killed_tensors;
/*! \brief The mapping from each binding block to the storage tokens that are create inside. */
std::unordered_map<const BindingBlockNode*, std::vector<const StorageTokenNode*>> block2tokens;
@@ -641,10 +638,10 @@ class StorageAllocator : public StorageAllocatorBaseVisitor {
// And release it if so.
for (const Expr& arg : call->args) {
Tokens tokens = GetTokens(arg);
- ForEachLeaf(tokens, [this, call](StorageToken token) {
+ ForEachLeaf(tokens, [this](StorageToken token) {
ICHECK_GT(token->ref_counter, 0);
token->ref_counter -= 1;
- this->CheckForRelease(token, call);
+ this->CheckForRelease(token);
});
}
}
@@ -662,11 +659,8 @@ class StorageAllocator : public StorageAllocatorBaseVisitor {
/*!
* \brief Check if a token has no reference and thus can be released. And release it if so.
* \param token The token to be checked.
- * \param release_site The CallNode where the input token is send for release.
- * If the token is checked to release here, we keep record of the release site so that
- * kill_tensor can be inserted here at the rewrite stage.
*/
- void CheckForRelease(StorageToken token, const CallNode* release_site) {
+ void CheckForRelease(StorageToken token) {
// Sanity check: the token was allocated before and has non-negative reference.
ICHECK_GE(token->storage_id, 0);
ICHECK_GE(token->ref_counter, 0);
@@ -675,10 +669,6 @@ class StorageAllocator : public StorageAllocatorBaseVisitor {
allocator_.Release(token);
auto it = token2cur_tensor_.find(token.get());
ICHECK(it != token2cur_tensor_.end());
- // Record that the tensors that are using this token will be killed
- // immediately after the release site.
- std::vector<Var>& killed_tensors = expr2killed_tensors[release_site];
- killed_tensors.insert(killed_tensors.end(), it->second.begin(), it->second.end());
token2cur_tensor_.erase(it);
}
}
@@ -696,19 +686,15 @@ class StorageAllocator : public StorageAllocatorBaseVisitor {
* \details
* - For each builtin alloc_tensor that was planned, substitute it with a memory
* alloc_tensor. If no memory alloc_storage was created for it before, create one.
- * - Insert memory kill_tensor at the release site of each tensor.
- * - Insert memory kill_storage at the end of each binding block, for the tokens allocated in it.
*/
class StorageAllocationRewriter : public ExprMutator {
public:
explicit StorageAllocationRewriter(
IRModule mod, std::unordered_map<const ExprNode*, StorageToken> alloc_tensor2token,
- std::unordered_map<const ExprNode*, std::vector<Var>> expr2killed_tensors,
std::unordered_map<const BindingBlockNode*, std::vector<const StorageTokenNode*>>
block2tokens)
: ExprMutator(std::move(mod)),
alloc_tensor2token_(std::move(alloc_tensor2token)),
- expr2killed_tensors_(std::move(expr2killed_tensors)),
block2tokens_(std::move(block2tokens)) {}
IRModule Rewrite() {
@@ -728,38 +714,6 @@ class StorageAllocationRewriter : public ExprMutator {
private:
using ExprMutator::VisitExpr_;
- BindingBlock VisitBindingBlock_(const BindingBlockNode* block) final {
- builder_->BeginBindingBlock();
- for (Binding binding : block->bindings) {
- this->VisitBinding(binding);
- }
-
- // Insert `memory.kill_storage` for the storage tokens allocated inside this block.
- for (const StorageTokenNode* token : block2tokens_[block]) {
- auto it_token = token2storage_var_.find(token);
- ICHECK(it_token != token2storage_var_.end());
- static const Op& mem_kill_storage = Op::Get("relax.memory.kill_storage");
- this->builder_->Emit(Call(mem_kill_storage, {it_token->second}), /*name_hint=*/"_");
- }
-
- BindingBlock new_block = builder_->EndBlock();
- return new_block;
- }
-
- void VisitBinding_(const VarBindingNode* binding) final {
- ExprMutator::VisitBinding_(binding);
-
- // Insert `memory.kill_tensor` for the tensors that need to be killed after this binding.
- auto it = expr2killed_tensors_.find(binding->value.get());
- if (it != expr2killed_tensors_.end()) {
- for (const Var& var : it->second) {
- static const Op& mem_kill_tensor = Op::Get("relax.memory.kill_tensor");
- this->builder_->Emit(Call(mem_kill_tensor, {Downcast<Var>(this->VisitExpr(var))}),
- /*name_hint=*/"_");
- }
- }
- }
-
Expr VisitExpr_(const CallNode* call) final {
auto it = alloc_tensor2token_.find(call);
if (it != alloc_tensor2token_.end()) {
@@ -805,8 +759,6 @@ class StorageAllocationRewriter : public ExprMutator {
its corresponding underlying storage token that it is using.
*/
std::unordered_map<const ExprNode*, StorageToken> alloc_tensor2token_;
- /*! \brief The mapping from each Expr to the tensors that need to be killed after it. */
- std::unordered_map<const ExprNode*, std::vector<Var>> expr2killed_tensors_;
/*! \brief The mapping from each binding block to the storage tokens that are create inside. */
std::unordered_map<const BindingBlockNode*, std::vector<const StorageTokenNode*>> block2tokens_;
/*! \brief The mapping from each token to its corresponding storage var in each function. */
@@ -822,7 +774,6 @@ IRModule StaticPlanBlockMemory(IRModule mod) {
// Step 3. Rewrite the function.
StorageAllocationRewriter rewriter(std::move(mod), //
std::move(allocator.alloc_tensor2token),
- std::move(allocator.expr2killed_tensors),
std::move(allocator.block2tokens));
return rewriter.Rewrite();
}
diff --git a/tests/python/relax/test_transform_static_plan_block_memory.py b/tests/python/relax/test_transform_static_plan_block_memory.py
index 451fe4fbe2..0c24f90efc 100644
--- a/tests/python/relax/test_transform_static_plan_block_memory.py
+++ b/tests/python/relax/test_transform_static_plan_block_memory.py
@@ -104,29 +104,22 @@ def test_basic():
cls = Expected
storage: R.Object = R.memory.alloc_storage(R.shape([32]), virtual_device_index=0, storage_scope="global", dtype="float32")
alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), dtype="float32")
- _: R.Tuple() = cls.exp(x, alloc)
+ _ = cls.exp(x, alloc)
lv: R.Tensor((2, 4), dtype="float32") = alloc
lv1: R.Tensor((8,), dtype="float32") = R.reshape(lv, (8,))
storage1: R.Object = R.memory.alloc_storage(R.shape([40]), virtual_device_index=0, storage_scope="global", dtype="float32")
alloc1: R.Tensor((8,), dtype="float32") = R.memory.alloc_tensor(storage1, 0, R.shape([8]), dtype="float32")
- _1: R.Tuple() = cls.relu(lv1, alloc1)
- _2: R.Tuple() = R.memory.kill_tensor(alloc)
- _3: R.Tuple() = R.memory.kill_tensor(lv1)
+ _ = cls.relu(lv1, alloc1)
lv2: R.Tensor((8,), dtype="float32") = alloc1
alloc2: R.Tensor((8,), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([8]), dtype="float32")
- _4: R.Tuple() = cls.add(lv2, R.const(1, "float32"), alloc2)
- _5: R.Tuple() = R.memory.kill_tensor(alloc1)
+ _ = cls.add(lv2, R.const(1, "float32"), alloc2)
lv3: R.Tensor((8,), dtype="float32") = alloc2
alloc3: R.Tensor((10,), dtype="float32") = R.memory.alloc_tensor(storage1, 0, R.shape([10]), dtype="float32")
- _6: R.Tuple() = cls.pad(lv3, alloc3)
- _7: R.Tuple() = R.memory.kill_tensor(alloc2)
+ _ = cls.pad(lv3, alloc3)
lv4: R.Tensor((10,), dtype="float32") = alloc3
alloc4: R.Tensor((10,), dtype="float32") = R.builtin.alloc_tensor(R.shape([10]), dtype="float32", runtime_device_index=0)
- _8: R.Tuple() = cls.log(lv4, alloc4)
- _9: R.Tuple() = R.memory.kill_tensor(alloc3)
+ _ = cls.log(lv4, alloc4)
gv5: R.Tensor((10,), dtype="float32") = alloc4
- _11: R.Tuple() = R.memory.kill_storage(storage)
- _10: R.Tuple() = R.memory.kill_storage(storage1)
return gv5
@I.ir_module
@@ -162,35 +155,32 @@ def test_basic():
storage: R.Object = R.vm.alloc_storage(R.shape([32]), R.prim_value(0), R.dtype("uint8"))
alloc: R.Tensor((2, 4), dtype="float32") = R.vm.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
_: R.Tuple = cls.exp(x, alloc)
- lv: R.Tensor((2, 4), dtype="float32") = alloc
- lv1: R.Tensor((8,), dtype="float32") = R.call_packed("vm.builtin.reshape", lv, R.shape([8]), sinfo_args=(R.Tensor((8,), dtype="float32"),))
+ lv1: R.Tensor((8,), dtype="float32") = R.call_packed("vm.builtin.reshape", alloc, R.shape([8]), sinfo_args=(R.Tensor((8,), dtype="float32"),))
+ _ = R.vm.kill_object(alloc)
storage1: R.Object = R.vm.alloc_storage(R.shape([40]), R.prim_value(0), R.dtype("uint8"))
alloc1: R.Tensor((8,), dtype="float32") = R.vm.alloc_tensor(storage1, R.prim_value(0), R.shape([8]), R.dtype("float32"))
- _1: R.Tuple = cls.relu(lv1, alloc1)
- __1: R.Tuple = R.vm.kill_object(alloc)
- _1_1: R.Tuple = R.vm.kill_object(lv1)
- lv2: R.Tensor((8,), dtype="float32") = alloc1
+ _ = cls.relu(lv1, alloc1)
+ _ = R.vm.kill_object(lv1)
alloc2: R.Tensor((8,), dtype="float32") = R.vm.alloc_tensor(storage, R.prim_value(0), R.shape([8]), R.dtype("float32"))
- _2: R.Tuple = cls.add(lv2, R.const(1, "float32"), alloc2)
- _2_1: R.Tuple = R.vm.kill_object(alloc1)
- lv3: R.Tensor((8,), dtype="float32") = alloc2
+ _ = R.vm.kill_object(storage)
+ _ = cls.add(alloc1, R.const(1, "float32"), alloc2)
+ _ = R.vm.kill_object(alloc1)
alloc3: R.Tensor((10,), dtype="float32") = R.vm.alloc_tensor(storage1, R.prim_value(0), R.shape([10]), R.dtype("float32"))
- _3: R.Tuple = cls.pad(lv3, alloc3)
- _3_1: R.Tuple = R.vm.kill_object(alloc2)
- lv4: R.Tensor((10,), dtype="float32") = alloc3
- storage_1: R.Object = R.vm.alloc_storage(R.shape([40]), R.prim_value(0), R.dtype("uint8"))
- alloc4: R.Tensor((10,), dtype="float32") = R.vm.alloc_tensor(storage_1, R.prim_value(0), R.shape([10]), R.dtype("float32"))
- _4: R.Tuple = cls.log(lv4, alloc4)
- _4_1: R.Tuple = R.vm.kill_object(alloc3)
- gv: R.Tensor((10,), dtype="float32") = alloc4
- _5: R.Tuple = R.vm.kill_object(storage)
- _6: R.Tuple = R.vm.kill_object(storage1)
- return gv
+ _ = R.vm.kill_object(storage1)
+ _ = cls.pad(alloc2, alloc3)
+ _ = R.vm.kill_object(alloc2)
+ storage2: R.Object = R.vm.alloc_storage(R.shape([40]), R.prim_value(0), R.dtype("uint8"))
+ alloc4: R.Tensor((10,), dtype="float32") = R.vm.alloc_tensor(storage2, R.prim_value(0), R.shape([10]), R.dtype("float32"))
+ _ = R.vm.kill_object(storage2)
+ _ = cls.log(alloc3, alloc4)
+ _ = R.vm.kill_object(alloc3)
+ return alloc4
# fmt: on
mod = relax.transform.StaticPlanBlockMemory()(Module)
tvm.ir.assert_structural_equal(mod, Expected)
mod = relax.transform.LowerAllocTensor()(mod)
+ mod = relax.transform.KillAfterLastUse()(mod)
mod = relax.transform.VMBuiltinLower()(mod)
tvm.ir.assert_structural_equal(mod, ExpectedLowered)
@@ -263,7 +253,6 @@ def test_different_dtype():
storage, 0, R.shape([2, 3]), dtype="float32"
)
_: R.Tuple() = cls.add(x, x, alloc)
- _1: R.Tuple() = R.memory.kill_tensor(alloc)
gv1: R.Tensor((2, 3), dtype="float32") = alloc
storage1: R.Object = R.memory.alloc_storage(
R.shape([24]), virtual_device_index=0, storage_scope="global", dtype="int32"
@@ -272,10 +261,7 @@ def test_different_dtype():
storage1, 0, R.shape([2, 3]), dtype="int32"
)
_2: R.Tuple() = cls.add1(y, y, alloc1)
- _3: R.Tuple() = R.memory.kill_tensor(alloc1)
gv12: R.Tensor((2, 3), dtype="int32") = alloc1
- _5: R.Tuple() = R.memory.kill_storage(storage)
- _4: R.Tuple() = R.memory.kill_storage(storage1)
return x
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -325,9 +311,7 @@ def test_dtype_bool():
storage, 0, R.shape([2, 3]), dtype="bool"
)
_2: R.Tuple() = cls.add1(y, y, alloc)
- _3: R.Tuple() = R.memory.kill_tensor(alloc)
gv12: R.Tensor((2, 3), dtype="bool") = alloc
- _4: R.Tuple() = R.memory.kill_storage(storage)
return y
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -386,15 +370,12 @@ def test_same_dtype():
storage, 0, R.shape([2, 3]), dtype="float32"
)
_: R.Tuple() = cls.add(x, x, alloc)
- _1: R.Tuple() = R.memory.kill_tensor(alloc)
gv1: R.Tensor((2, 3), dtype="float32") = alloc
alloc1: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor(
storage, 0, R.shape([2, 3]), dtype="float32"
)
_2: R.Tuple() = cls.add(y, y, alloc1)
- _3: R.Tuple() = R.memory.kill_tensor(alloc1)
gv12: R.Tensor((2, 3), dtype="float32") = alloc1
- _4: R.Tuple() = R.memory.kill_storage(storage)
return x
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -617,27 +598,17 @@ def test_nested_tuple():
storage3, 0, R.shape([2, 3]), dtype="float32"
)
_3: R.Tuple() = cls.exp(y1_, alloc3)
- _4: R.Tuple() = R.memory.kill_tensor(alloc)
- _11: R.Tuple() = R.memory.kill_tensor(alloc3)
z1: R.Tensor((2, 3), dtype="float32") = alloc3
alloc4: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor(
storage, 0, R.shape([2, 3]), dtype="float32"
)
_41: R.Tuple() = cls.exp(y2_, alloc4)
- _21: R.Tuple() = R.memory.kill_tensor(alloc1)
- _31: R.Tuple() = R.memory.kill_tensor(alloc4)
z2: R.Tensor((2, 3), dtype="float32") = alloc4
alloc5: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor(
storage3, 0, R.shape([2, 3]), dtype="float32"
)
_5: R.Tuple() = cls.exp(y3_, alloc5)
- _42: R.Tuple() = R.memory.kill_tensor(alloc2)
- _51: R.Tuple() = R.memory.kill_tensor(alloc5)
z3: R.Tensor((2, 3), dtype="float32") = alloc5
- _9: R.Tuple() = R.memory.kill_storage(storage)
- _7: R.Tuple() = R.memory.kill_storage(storage1)
- _8: R.Tuple() = R.memory.kill_storage(storage2)
- _6: R.Tuple() = R.memory.kill_storage(storage3)
return x
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -696,9 +667,7 @@ def test_call_packed_external_func():
R.shape([2, 3]), R.dtype("float32"), R.prim_value(0)
)
_1: R.Tuple = R.call_packed("extern_func", y, alloc1, sinfo_args=(R.Tuple(),))
- _2: R.Tuple = R.memory.kill_tensor(alloc)
z: R.Tensor((2, 3), dtype="float32") = alloc1
- _3: R.Tuple = R.memory.kill_storage(storage)
return z
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -755,7 +724,6 @@ def test_zero_reference():
alloc: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor(
storage, 0, R.shape([2, 3]), dtype="float32"
)
- _: R.Tuple() = R.memory.kill_storage(storage)
return x
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -878,7 +846,6 @@ def test_multiple_functions():
storage, 0, R.shape([2, 3]), dtype="float32"
)
_: R.Tuple() = cls.add(x, x, alloc)
- _1: R.Tuple() = R.memory.kill_tensor(alloc)
gv1: R.Tensor((2, 3), dtype="float32") = alloc
storage1: R.Object = R.memory.alloc_storage(
R.shape([24]), virtual_device_index=0, storage_scope="global", dtype="int32"
@@ -887,10 +854,7 @@ def test_multiple_functions():
storage1, 0, R.shape([2, 3]), dtype="int32"
)
_2: R.Tuple() = cls.add1(y, y, alloc1)
- _3: R.Tuple() = R.memory.kill_tensor(alloc1)
gv12: R.Tensor((2, 3), dtype="int32") = alloc1
- _5: R.Tuple() = R.memory.kill_storage(storage)
- _4: R.Tuple() = R.memory.kill_storage(storage1)
return x
@R.function
@@ -906,15 +870,12 @@ def test_multiple_functions():
storage, 0, R.shape([2, 3]), dtype="float32"
)
_: R.Tuple() = cls.add(x, x, alloc)
- _1: R.Tuple() = R.memory.kill_tensor(alloc)
gv1: R.Tensor((2, 3), dtype="float32") = alloc
alloc1: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor(
storage, 0, R.shape([2, 3]), dtype="float32"
)
_2: R.Tuple() = cls.add(y, y, alloc1)
- _3: R.Tuple() = R.memory.kill_tensor(alloc1)
gv12: R.Tensor((2, 3), dtype="float32") = alloc1
- _4: R.Tuple() = R.memory.kill_storage(storage)
return x
mod = relax.transform.StaticPlanBlockMemory()(Module)
@@ -1011,23 +972,16 @@ def test_tir_var_upper_bound():
storage1: R.Object = R.memory.alloc_storage(R.shape([40]), R.prim_value(0), R.str("global"), R.dtype("float32"))
alloc1: R.Tensor((2 * n,), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2 * n]), R.dtype("float32"))
_1: R.Tuple = cls.relu(lv1, alloc1)
- __1: R.Tuple = R.memory.kill_tensor(alloc)
- _1_1: R.Tuple = R.memory.kill_tensor(lv1)
lv2: R.Tensor((2 * n,), dtype="float32") = alloc1
alloc2: R.Tensor((2 * n,), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2 * n]), R.dtype("float32"))
_2: R.Tuple = cls.add(lv2, R.const(1, "float32"), alloc2)
- _2_1: R.Tuple = R.memory.kill_tensor(alloc1)
lv3: R.Tensor((2 * n,), dtype="float32") = alloc2
alloc3: R.Tensor((2 * n + 2,), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2 * n + 2]), R.dtype("float32"))
_3: R.Tuple = cls.pad(lv3, alloc3)
- _3_1: R.Tuple = R.memory.kill_tensor(alloc2)
lv4: R.Tensor((2 * n + 2,), dtype="float32") = alloc3
alloc4: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([10]), R.dtype("float32"), R.prim_value(0))
_4: R.Tuple = cls.log(lv4, alloc4)
- _4_1: R.Tuple = R.memory.kill_tensor(alloc3)
gv: R.Tensor((2 * n + 2,), dtype="float32") = alloc4
- _5: R.Tuple = R.memory.kill_storage(storage)
- _6: R.Tuple = R.memory.kill_storage(storage1)
return gv
# fmt: on
@@ -1079,14 +1033,10 @@ def test_tir_var_decreasing_monotone():
storage1: R.Object = R.memory.alloc_storage(R.shape([8000]), R.prim_value(0), R.str("global"), R.dtype("float32"))
alloc1: R.Tensor((n, m, T.max(n - m, 1)), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([n, m, T.max(n - m, 1)]), R.dtype("float32"))
_1: R.Tuple = cls.tir_exp(y, alloc1)
- __1: R.Tuple = R.memory.kill_tensor(alloc)
z: R.Tensor((n, m, T.max(n - m, 1)), dtype="float32") = alloc1
alloc2: R.Tensor((n, m, T.max(n - m, 1)), dtype="float32") = R.builtin.alloc_tensor(R.shape([n, m, T.max(n - m, 1)]), R.dtype("float32"), R.prim_value(0))
_2: R.Tuple = cls.tir_exp(z, alloc2)
- _1_1: R.Tuple = R.memory.kill_tensor(alloc1)
r: R.Tensor((n, m, T.max(n - m, 1)), dtype="float32") = alloc2
- _2_1: R.Tuple = R.memory.kill_storage(storage)
- _3: R.Tuple = R.memory.kill_storage(storage1)
return r
# fmt: on
@@ -1144,14 +1094,10 @@ def test_call_tir_dyn():
storage1: R.Object = R.memory.alloc_storage(R.shape([80]), R.prim_value(0), R.str("global"), R.dtype("float32"))
alloc1: R.Tensor((n,), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([n]), R.dtype("float32"))
_1: R.Tuple = cls.tir_exp(full, alloc1)
- __1: R.Tuple = R.memory.kill_tensor(alloc)
lv2: R.Tensor((n,), dtype="float32") = alloc1
alloc2: R.Tensor((n,), dtype="float32") = R.builtin.alloc_tensor(R.shape([n]), R.dtype("float32"), R.prim_value(0))
_2: R.Tuple = cls.tir_exp(lv2, alloc2)
- _1_1: R.Tuple = R.memory.kill_tensor(alloc1)
lv3: R.Tensor((n,), dtype="float32") = alloc2
- _2_1: R.Tuple = R.memory.kill_storage(storage)
- _3: R.Tuple = R.memory.kill_storage(storage1)
return lv3
# fmt: on