Propagate stream in loop irrespective of device mesh

[Priya2698]

When filtering the reference inputs, inputs without device mesh were removed. This caused fusions with only stream-parallel tensorviews to skip propagation.

[Priya2698]

!test

[github-actions]

Review updated until commit f356951

Description

• Propagate stream parallelization in loop domains

• Skip sharding propagation for scatter op outputs

• Sort TensorViews by device/stream dimensions

• Add tests for stream propagation

---

Changes walkthrough 📝

	Relevant files
Enhancement	propagate_shardings.cpp Enhance sharding propagation with stream support                  csrc/preseg_passes/propagate_shardings.cpp Introduced sortTvsByParallelDims to sort TensorViews by device/stream dimensions Updated input/output ordering to consider both device and stream dimensions Added check to skip sharding propagation for scatter op outputs Removed redundant null and device mesh checks +45/-63
Tests	test_stream.cpp Add tests for stream propagation                                                  tests/cpp/test_stream.cpp Added ForwardPropagation test for stream parallelization Added BackwardPropagation test for stream sharding Updated includes to use all_ops.h and propagate_shardings.h +45/-1

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Issue The function sortTvsByParallelDims includes both device and stream dimensions in its count, but the previous version sortTvsByDeviceDims only considered device dimensions. This change may alter the priority ordering of TensorViews during sharding propagation, potentially affecting correctness when stream and device dimensions interact. auto num_parallel_dims = [](TensorView* tv) { return std::count_if( tv->getLoopDomain().begin(), tv->getLoopDomain().end(), [](IterDomain* id) { return !id->isReduction() && (id->isStream() || id->isDeviceDim()); }); }; std::vector<TensorView*> tvs_vec(tvs.begin(), tvs.end()); std::ranges::stable_sort(tvs_vec, [&num_parallel_dims](auto a, auto b) { return std::make_pair(num_parallel_dims(a), a->getDeviceMesh().rank()) > std::make_pair(num_parallel_dims(b), b->getDeviceMesh().rank()); }); return tvs_vec; Logic Gap The removal of NVF_ERROR checks for device mesh presence in getOrderedReferenceInputs may allow propagation from inputs without a device mesh, leading to undefined behavior during sharding propagation if such inputs are encountered. for (auto* ref_input : reference_inputs) { NVF_ERROR(ref_input != nullptr); // Consider out [M, N] = linear (inp [M, K], weight (N, // K)) with inp sharded on M ([DIDx(d), M/d, K]) and weight sharded on N // ([DIDy(d), N/d, K]). We propagate from weights first, so the output Missing Validation The backward propagation loop skips user-sharded tensors, but does not validate whether the reference output has a device mesh before use, which could lead to dereferencing a null or invalid mesh in transformLoopDomain. for (Expr* expr : exprs | std::views::reverse) { const auto& outputs = ir_utils::filterByType<TensorView>(expr->outputs()); if (outputs.empty()) { continue; } // All outputs of an expression (Welford, SDPA) should be uniformly sharded. // We pick the most parallel output as the reference. // This is to avoid picking seed/offset tvs in SDPA. std::vector<TensorView*> sorted_outputs = sortTvsByParallelDims(outputs); TensorView* ref_output = sorted_outputs.front();

[Priya2698]

!test

[Priya2698]

!test

[Priya2698]

!test

[Priya2698]

PR #5365

[Priya2698]

!test