Propagate Stream parallel type in allocation

[Priya2698]

Issue #5309
Unlike device parallelization, a stream parallel tensorview (in loop) may or may not have a stream-parallel allocation domain.

We propagate based on the following:

1. If it is a device parallel type -> always propagate
2. If it is a fusion input or output -> id is not stream parallelized
3. If the stream ID in a tensorview is not mapped to stream ID in all of its consumers -> id is not stream parallelized

For cases like:

Fuser/tests/cpp/test_overlap.cpp

Line 155 in f8e84e5

// allocation. This dimension should eventually be parallelized on `Stream`

, we want to start with replicating Stream-parallel ID, that is the allocation is not parallelized. However, this ID will appear in the logical domain due to rfactor and with the current contract, be allocated fully regardless of parallelization. So I am not making this a condition in the pass, yet.

This can be changed in future when we need.

Depends on #5363

[Priya2698]

!test

[github-actions]

Review updated until commit 46253fb

Description

• Propagate Stream parallelization to allocation domain conditionally

• Prevent allocation sharding for non-device, non-Stream tensors

• Add tests for Stream-parallel allocation behavior

• Print debug transforms in finalize pass for diagnostics

---

Changes walkthrough 📝

	Relevant files
Enhancement	finalize_multidevice_domains.cpp Implement conditional Stream allocation sharding                  csrc/preseg_passes/finalize_multidevice_domains.cpp Introduced shardAllocation to handle device and Stream parallelization Added shouldParallelizeAllocationOnStream to check Stream consumer consistency Added isLoopStreamParallelized to detect Stream in loop domain Skip sharding if no device mesh and not Stream-loop parallelized Print debug transform logging when enabled +57/-30
Bug fix	test_multidevice_lower_communication.cpp Fix allgather test device mesh setup                                          tests/cpp/test_multidevice_lower_communication.cpp Move setDeviceMesh call before split for correctness Remove manual split and allocation on output tensor Use setDeviceMesh on output to enable proper propagation +2/-3
Tests	test_stream.cpp Add Stream allocation propagation tests                                    tests/cpp/test_stream.cpp Add #include Add ShardedAllocation test for Stream allocation in loops Add ReplicatedAllocation test when Stream not in consumers Verify allocation domain matches logical or loop domain +58/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Loop Stream Check The function isLoopStreamParallelized checks if any loop domain ID is stream-parallel, but it may not account for nested or conditional loop structures where stream parallelization is context-dependent. This could lead to incorrect propagation decisions in complex loop scenarios. bool isLoopStreamParallelized(const TensorView* tv) { return std::any_of( tv->getLoopDomain().begin(), tv->getLoopDomain().end(), [](IterDomain* id) { return id->isStream(); }); } Allocation Sharding Logic The shardAllocation function skips splitting for stream-parallel outer dimensions when shouldParallelizeAllocationOnStream returns false, but it does not handle cases where partial stream parallelization exists across consumer tensorviews, potentially leading to inconsistent memory layouts. if (split->outer()->isStream() && !shouldParallelizeAllocationOnStream(tv)) { continue; } Test Coverage The new tests ShardedAllocation and ReplicatedAllocation verify basic behavior but do not test edge cases such as tensorviews with mixed parallel types or multiple stream-parallel dimensions, which could expose flaws in the propagation logic. TEST_F(StreamTest, ShardedAllocation) { auto fusion = std::make_unique<Fusion>(); FusionGuard fg(fusion.get()); const int64_t s = 2; TensorView* tv0 = makeContigTensor(3); TensorView* tv1 = add(tv0, IrBuilder::create<Val>(1.0)); TensorView* tv2 = sum(tv1, {2}); TensorView* tv3 = div(tv1, IrBuilder::create<Val>(2.0)); fusion->addInput(tv0); fusion->addOutput(tv2); fusion->addOutput(tv3); tv0->outer_split(0, s); tv0->axis(0)->parallelize(ParallelType::Stream); preseg_passes::OptimizationPass<preseg_passes::PreSegmenter>::runPass( fusion.get()); for (auto* tv : {tv0, tv1, tv2, tv3}) { EXPECT_TRUE(tv->axis(0)->isStream()) << tv; if (tv->isFusionOutput() || tv->isFusionInput()) { EXPECT_EQ(tv->getAllocationDomain(), tv->getLogicalDomain()); } else { EXPECT_EQ(tv->getAllocationDomain(), tv->getLoopDomain()); } } } TEST_F(StreamTest, ReplicatedAllocation) { auto fusion = std::make_unique<Fusion>(); FusionGuard fg(fusion.get()); const int64_t s = 2; TensorView* tv0 = makeContigTensor(3); TensorView* tv1 = add(tv0, IrBuilder::create<Val>(1.0)); TensorView* tv2 = sum(tv1, {2}); TensorView* tv3 = div(tv1, IrBuilder::create<Val>(2.0)); fusion->addInput(tv0); fusion->addOutput(tv2); fusion->addOutput(tv3); tv0->outer_split(0, s); tv0->axis(0)->parallelize(ParallelType::Stream); tv2->outer_split(1, s); tv2->axis(1)->parallelize(ParallelType::Stream); preseg_passes::OptimizationPass<preseg_passes::PreSegmenter>::runPass( fusion.get()); for (auto* tv : {tv0, tv1, tv2, tv3}) { EXPECT_TRUE(tv->axis(0)->isStream()) << tv; EXPECT_EQ(tv->getAllocationDomain(), tv->getLogicalDomain()); } }

[Priya2698]

!test

[Priya2698]

!test

[Priya2698]

!test

[wujingyue]

LGTM otherwise

[Priya2698]

!test

[wujingyue]

Looks great!

[Priya2698]

!test

[Priya2698]

!test

[Priya2698]

!test

[wujingyue]

Nit: I believe you can move this filter to loop_stream_device_view as well. This way, we put all the filters in one location.

[Priya2698]

This PR was merged, but I'll do it in a follow-up!