Fix liveness of TMem address smem TV

[zasdfgbnm]

To use TMem, we must have a smem tensor for storing the address of TMem, and uses smem_tv[0] = tcgen05::alloc() to store the allocated TMem address. The liveness of this smem tv should be from the tcgen05::alloc to the tcgen05::dealloc. However, we are currently unable to get that due to tcgen05::dealloc(smem_tv[0]) is not a "TV Op". Due to this bug, the smem address tensor is allocated overlapping with other smem tensors. This PR fixes this bug by handling both TV Op or non-TV op correctly.

[zasdfgbnm]

!test

[github-actions]

Description

• Handle IfThenElse nodes in BufferReuseDebugPrinter

• Improve liveness analysis for IfThenElse nodes

• Update memory type for tv1 in TMemTest

---

Changes walkthrough 📝

	Relevant files
Enhancement	alias_memory.cpp Improve liveness analysis and handle IfThenElse nodes        csrc/device_lower/pass/alias_memory.cpp Implement handling of IfThenElse nodes in BufferReuseDebugPrinter Add scope management for IfThenElse nodes in AllocationInfoMap Update liveness analysis to handle non-TensorView inputs and outputs +18/-10
Bug fix	test_tmem.cpp Update memory type for tv1 in TMemTest                                      tests/cpp/test_tmem.cpp Change memory type of tv1 from Global to Shared in TMemTest +2/-2

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Issue The handle method for kir::IfThenElse now throws an exception if the debug printer is not set. This could lead to runtime errors if the debug printer is not initialized. indent(); os_ << "IF " << node->predicate()->toString() << ":\n"; } Code Duplication The loop for handling inputs and outputs in collectLivenessInfoOfExpr has duplicate code for checking if a Val is a TensorView. This could be refactored to avoid duplication. for (Val* input : expr->inputs()) { TensorView* input_tv = ir_utils::getTv(input); if (!input_tv) { continue; } auto alloc_info = getAllocInfoFromTV(input_tv); if (alloc_info) { if (!isSerialBroadcastResolution(input_tv, for_loops_)) { alloc_info->inner_live_interval->markRead(expr_pos); Code Duplication The loop for handling inputs and outputs in collectLivenessInfoOfExpr has duplicate code for checking if a Val is a TensorView. This could be refactored to avoid duplication. // Allocate is inlined in the innermost loop, // so outer live interval is the same as inner. alloc_info->outer_live_interval->markRead(expr_pos); } } } for (Val* output : expr->outputs()) { TensorView* output_tv = ir_utils::getTv(output); if (!output_tv) { continue; }

[zasdfgbnm]

Using smem or gmem makes no difference in testing "InexactParallelType", but using smem could trigger the bug of liveness.

[zasdfgbnm]

!test

[jacobhinkle]

LGTM. I guess another way to go would be to create a DeAllocTMem node which has the TV as an input and translate it in the inline PTX pass. Then that expr would be a TV op instead of the current Asm node that is inserted.

[jacobhinkle]

It seems like the only difference here is in properly handling kir::TensorIndex inputs. Is that right? I thought that should not be needed since this pass runs before indexing.

[zasdfgbnm]

Yes exactly.

I thought that should not be needed since this pass runs before indexing.

In general yes, but there are always exceptions, such as if you are using smem for non-tensors like mbarrier or TMem address.

[jacobhinkle]

Ah right, because it's a kir::Asm, it has TensorIndex inputs upon creation.