Stream-parallelize loops #5751
Stream-parallelize loops #5751
Conversation
It's never used by these tests.
|
Review updated until commit 215a166 Description
|
| Relevant files | |||
|---|---|---|---|
| Enhancement | |||
| Tests | |||
| Documentation | |||
| Bug fix | 1 files
| ||
| Configuration changes | 1 files
|
PR Reviewer Guide
Here are some key observations to aid the review process:
| 🧪 PR contains tests |
| ⚡ Recommended focus areas for review |
Stream Synchronization Logic
|
Greptile SummaryThis PR implements stream parallelization for loops in nvFuser's host IR to enable overlapping of computation (matmul) and communication (allreduce) operations. The implementation adds a new
The changes include:
Benchmark results show nvFuser is slightly faster than the reference implementation (3.8ms vs 4.6ms mean), addressing issue #5308. The implementation correctly handles stream ordering and synchronization. Confidence Score: 5/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Main as Main Stream
participant W0 as Worker Stream 0
participant W1 as Worker Stream 1
participant W2 as Worker Stream 2
Note over Main: GetCurrentStream
Note over Main: Start ForLoop (i=0..2)
Main->>W0: SetCurrentStream(0)
W0->>Main: Synchronize Main Stream
Note over W0: ShardByStream (chunk 0)
Note over W0: Linear (matmul chunk 0)
Note over W0: AllReduce (async)
par Parallel Execution
Main->>W1: SetCurrentStream(1)
W1->>Main: Synchronize Main Stream
Note over W1: ShardByStream (chunk 1)
Note over W1: Linear (matmul chunk 1)
Note over W1: AllReduce (async)
and
Main->>W2: SetCurrentStream(2)
W2->>Main: Synchronize Main Stream
Note over W2: ShardByStream (chunk 2)
Note over W2: Linear (matmul chunk 2)
Note over W2: AllReduce (async)
end
Note over Main: End ForLoop
Note over Main: SetCurrentStream(Main)
Note over Main: Start Join Loop (i=0..2)
Main->>W0: Synchronize Worker Stream 0
Main->>W1: Synchronize Worker Stream 1
Main->>W2: Synchronize Worker Stream 2
Note over Main: All chunks complete
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| auto old_begin = for_loop->body().exprs().begin(); | ||
| for_loop->body().insert(old_begin, set_stream); | ||
| for_loop->body().insert(old_begin, sync_main); |
There was a problem hiding this comment.
logic: Insertion order creates reversed execution order. The second insert at old_begin places sync_main before set_stream, but the worker stream should be set before synchronization.
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| for_loop->body().insert(old_begin, sync_main); | |
| for_loop->body().insert(old_begin, set_stream); |
There was a problem hiding this comment.
Skip this file. Subsumed by another PR.
csrc/ir/internal_nodes.h
Outdated
| namespace nvfuser { | ||
|
|
||
| class ViewTransform; | ||
| class Scope; |
There was a problem hiding this comment.
Not necessary -- class Scope is defined in the same file.
There was a problem hiding this comment.
Skip this file -- subsumed by another PR
| from nvfuser_direct import FusionDefinition, ParallelType, DataType | ||
|
|
||
|
|
||
| def test_matmul(nvfuser_direct_test): |
There was a problem hiding this comment.
The tests didn't use the nvfuser_direct_test fixture.
| for_loop->body().insert(old_begin, set_stream); | ||
| for_loop->body().insert(old_begin, sync_main); |
There was a problem hiding this comment.
logic: Wrong insertion order - both inserts use old_begin, causing sync_main to be placed before set_stream. This means synchronization happens on the main stream before switching to the worker stream. The worker stream should be set first, then synchronize.
Current order: Sync(main) → SetStream(worker)
Required order: SetStream(worker) → Sync(main)
Fix by inserting sync_main at the updated iterator after inserting set_stream
There was a problem hiding this comment.
This doesn't make sense. Maybe it missed that Scope contains a linked list?
There was a problem hiding this comment.
You're absolutely right! I completely missed that Scope uses a linked list structure. When you insert at old_begin twice, both operations insert at the same position, so the second insert (sync_main) goes before the first insert (set_stream).
Looking at the code, both insertions use the same iterator old_begin, which means:
- First insert:
set_streamgoes atold_begin - Second insert:
sync_maingoes at the sameold_beginposition, pushingset_streamforward
This creates the wrong order: sync_main → set_stream
The fix is to update the iterator after the first insertion:
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); |
Should be:
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); |
Wait, that's still wrong. Let me fix this properly:
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); |
Actually, the correct fix is:
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| auto new_pos = for_loop->body().exprs().begin(); | |
| std::advance(new_pos, 1); | |
| for_loop->body().insert(new_pos, sync_main); |
Or more simply, insert sync_main after inserting set_stream by using the updated iterator position.
There was a problem hiding this comment.
insert(old_begin, ...) inserts before old_begin, doesn't it?
There was a problem hiding this comment.
You're absolutely right! insert(old_begin, ...) inserts before old_begin. So the current code:
for_loop->body().insert(old_begin, set_stream); // inserts set_stream before old_begin
for_loop->body().insert(old_begin, sync_main); // inserts sync_main before old_beginSince both insertions use the same old_begin iterator, and insert() inserts before that position, the second insertion (sync_main) goes before the first insertion (set_stream).
This creates the order: sync_main → set_stream → (original body)
But we want: set_stream → sync_main → (original body)
The fix is to update the iterator after the first insertion:
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(for_loop->body().exprs().begin(), sync_main); |
This way set_stream gets inserted first, then sync_main gets inserted at the new beginning (which is now after set_stream).
There was a problem hiding this comment.
mpirun -np 2 -x NVFUSER_DUMP=host_ir pytest tests/python/multidevice/test_overlap.py -k 'row_parallel_linear_forward and not reference and not benchmark' --only-mpi -vs confirmed that SetCurrentStream is indeed before Synchronize Stream. The code as is is correct.
%HostIrContainer { (T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1})) -> (T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1})) :
T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}) = ALLOCATE(buffer=T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), mem_type=global, size=( i0 * 2 ), zero_init=false, resets_to_zero=false)
GetCurrentStream into Stream 0x3df7aa70
FOR i535 from 0 to 3:
SetCurrentStream to Stream i535
Synchronize Stream 0x3df7aa70
T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}) = ShardByStream(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), stream_index = i535)
T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}) = ALLOCATE(buffer=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), mem_type=global, size=( ( ceilDiv(i0, 3) ) * 12 ), zero_init=false, resets_to_zero=false)
T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1})
= linear(T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}),
T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1}) )
T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}) = ShardByStream(T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), stream_index = i535)
Communication 272 (type=Allreduce, team=(0 1), input=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), output=T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}), backend=NCCL)
Wait Communication 272
SetCurrentStream to Stream 0x3df7aa70
FOR i535 from 0 to 3:
Synchronize Stream i535
} // %HostIrContainer
csrc/host_ir/assign_streams.cpp
Outdated
|
|
||
| // We should check that the loop is stream-parallel. This is not necessary | ||
| // at this moment because all loops are stream-parallel. This is also hard | ||
| // to do becauase hir::ForLoop doesn't point to the source IterDomain. |
There was a problem hiding this comment.
To check my understanding: the iterdomain in ForLoop is not Stream Parallel unlike the iterdomain of the output loop domain?
There was a problem hiding this comment.
Unlike kir::ForLoop, hir::ForLoop doesn't have the IterDomain as an input. Maybe it should.
There was a problem hiding this comment.
Is it as simple as addInput(id) as is done for kir::ForLoop?
There was a problem hiding this comment.
That's plausible. I'm debating with myself on how much host IR depend on nvFuser specifics like IterDomain. I could go with adding a ParallelType as an attribute. 🤷
Co-authored-by: Priya Mishra <[email protected]>
|
!test |
| // We should check that the loop is stream-parallel. This is not necessary | ||
| // at this moment because all loops are stream-parallel. This is also hard | ||
| // to do because hir::ForLoop doesn't point to the source IterDomain. |
There was a problem hiding this comment.
Do ALL hir::ForLoops stream-parallelize? Is there no case where we want to sequentially loop in hir? or is this pass triggered by some other condition I'm not seeing?
There was a problem hiding this comment.
Do ALL hir::ForLoops stream-parallelize?
Yes at this moment.
I'm considering separating ParallelType::Stream and ParallelType::HostSerial. The latter doesn't exist today. That's when we'll have to look at the parallel type of the loop index.
Priya2698
left a comment
Priya2698
left a comment
There was a problem hiding this comment.
LGTM. Let me know if you plan to test merging the for-loops in this PR.
| } | ||
|
|
||
| // We should check that the loop is stream-parallel. This is not necessary | ||
| // at this moment because all loops are stream-parallel. This is also hard |
There was a problem hiding this comment.
Not strictly for this PR, but similar to kir::ForLoop, hir::ForLoop can hold the source iterdomain for this check
| for_loop->body().insert(old_begin, set_stream); | ||
| for_loop->body().insert(old_begin, sync_main); | ||
|
|
||
| // After the loop: create a joining loop to synchronize all worker streams |
There was a problem hiding this comment.
Do you plan on merging this with the above for-loop?
There was a problem hiding this comment.
I haven't convinced myself that will work: http://nv/e-d
Co-authored-by: Priya Mishra <[email protected]>
Co-authored-by: Priya Mishra <[email protected]>
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
Implements stream parallelization for loops in nvFuser's host IR to enable overlapping computation and communication operations. Adds new AssignStreams optimization pass that transforms stream-parallel loops by capturing the main stream, assigning worker streams at loop iteration start with synchronization, and creating a joining loop afterward to synchronize all worker streams back. Includes comprehensive test coverage with benchmarks showing nvFuser slightly outperforms the PyTorch reference implementation (3.6ms vs 4.6ms mean).
Confidence Score: 4/5
- Safe to merge with minor improvement opportunities
- The implementation is well-structured with proper synchronization patterns, comprehensive tests, and correct integration into the host IR pipeline. The main concern is the acknowledged missing validation for loop stream-parallelization (line 31-33 in assign_streams.cpp), which could theoretically transform non-stream-parallel loops incorrectly, though the comment indicates all current loops are stream-parallel. The code follows established patterns, includes thorough test coverage, and demonstrates correct behavior via benchmarks and profiling.
csrc/host_ir/assign_streams.cpp- consider adding validation or assertion for stream-parallel loop check
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| csrc/host_ir/assign_streams.cpp | 4/5 | Implements AssignStreams pass to transform stream-parallel loops with proper synchronization; well-structured but lacks validation for loop stream-parallelization |
| csrc/host_ir/passes.cpp | 5/5 | Integrates AssignStreams pass into host IR pipeline after AllocateAndDeallocate; correct ordering |
| tests/python/direct/test_stream.py | 5/5 | Comprehensive tests for stream parallelization with matmul operations; validates correct kernel count and shapes |
Sequence Diagram
sequenceDiagram
participant Main as Main Stream
participant W0 as Worker Stream 0
participant W1 as Worker Stream 1
participant W2 as Worker Stream 2
Note over Main: GetCurrentStream (capture main)
Note over Main: FOR i=0 to 3
Main->>W0: SetCurrentStream(0)
W0->>Main: Synchronize(main_stream)
Note over W0: Execute loop body (matmul/allreduce)
Main->>W1: SetCurrentStream(1)
W1->>Main: Synchronize(main_stream)
Note over W1: Execute loop body (matmul/allreduce)
Main->>W2: SetCurrentStream(2)
W2->>Main: Synchronize(main_stream)
Note over W2: Execute loop body (matmul/allreduce)
Note over Main: SetCurrentStream(main_stream)
Note over Main: FOR i=0 to 3 (joining loop)
Main->>W0: Synchronize(worker_stream_0)
Main->>W1: Synchronize(worker_stream_1)
Main->>W2: Synchronize(worker_stream_2)
Note over Main: All workers synchronized back
| // We should check that the loop is stream-parallel. This is not necessary | ||
| // at this moment because all loops are stream-parallel. This is also hard | ||
| // to do because hir::ForLoop doesn't point to the source IterDomain. |
There was a problem hiding this comment.
The comment acknowledges this validation is skipped, but consider adding a TODO or assertion to track this technical debt. Without validation, non-stream-parallel loops could be incorrectly transformed, potentially leading to incorrect synchronization patterns. At minimum, add a NVF_CHECK that verifies the loop meets basic requirements (e.g., has a valid index, start, and stop).
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This PR implements stream parallelization for loops in nvFuser's host IR to enable overlapping of computation and communication operations. The core implementation adds a new AssignStreams optimization pass that transforms stream-parallel loops by assigning worker streams to each iteration and adding proper synchronization.
Key Changes:
- New
AssignStreamspass incsrc/host_ir/assign_streams.{cpp,h}that transforms loops to use worker streams - Integration into host IR pipeline via
csrc/host_ir/passes.cpp - Comprehensive test coverage in
tests/python/multidevice/test_overlap.pywith benchmarks - Code cleanup: removed unnecessary includes in
allocate_and_deallocate.h,ir.h, andinternal_nodes.h
Transformation Pattern:
For each loop, the pass:
- Captures the main stream before the loop
- At the start of each iteration: switches to a worker stream and synchronizes with the main stream
- After the loop: creates a joining loop that synchronizes all worker streams back to main
Issues Found:
- Copyright year is 2026 in both new files (should be 2025)
- Missing validation that loops are actually stream-parallel (acknowledged in code comment but not implemented)
The implementation correctly follows the stream synchronization pattern demonstrated in the PyTorch reference implementation. Benchmark results show nvFuser achieves slight performance improvements over the baseline.
Confidence Score: 4/5
- This PR is safe to merge with minor corrections needed for copyright years
- The implementation is technically sound with correct synchronization logic matching the reference implementation. The only actual errors are copyright year mistakes (2026 instead of 2025). The missing stream-parallel validation is acknowledged in comments and appears to be a known limitation rather than an oversight. Comprehensive tests provide good coverage.
- Pay attention to the copyright years in
csrc/host_ir/assign_streams.cppandcsrc/host_ir/assign_streams.hwhich need correction from 2026 to 2025
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| csrc/host_ir/assign_streams.cpp | 4/5 | New file implementing stream assignment pass for stream-parallel loops. Transforms loops to use worker streams with proper synchronization. Copyright year is incorrect (2026 instead of 2025). Logic is sound but lacks validation that loops are actually stream-parallel. |
| csrc/host_ir/assign_streams.h | 4/5 | New header file declaring AssignStreams optimization pass. Copyright year is incorrect (2026 instead of 2025). Clean interface design. |
| csrc/host_ir/passes.cpp | 5/5 | Added AssignStreams pass to the host IR pipeline after AllocateAndDeallocate. Simple integration, no issues found. |
| tests/python/multidevice/test_overlap.py | 5/5 | Added comprehensive tests for stream-parallelized row-parallel linear forward pass with benchmarks. Includes reference implementation matching the transformation pattern. Well-structured tests with proper validation. |
Sequence Diagram
sequenceDiagram
participant Main as Main Stream
participant Worker as Worker Stream i
participant Ops as Operations
Note over Main: Save main stream
Main->>Main: GetCurrentStream
Note over Main,Worker: Loop iteration starts
Main->>Worker: SetCurrentStream worker_stream i
Worker->>Main: Synchronize main_stream
Note over Worker: Wait for main work to complete
Worker->>Ops: Execute matmul
Worker->>Ops: Execute allreduce
Note over Worker: Work continues in parallel
Note over Main: After loop completes
Worker->>Main: SetCurrentStream main_stream
Note over Main: Join loop synchronizes workers
Main->>Worker: Synchronize worker_stream 0
Main->>Worker: Synchronize worker_stream 1
Main->>Worker: Synchronize worker_stream N
Note over Main: All workers synchronized
Note over Main: Continue on main stream
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
Overview
This PR implements stream parallelization for loops in nvFuser's host IR, enabling overlapping of computation (matmul) and communication (allreduce) operations to improve performance in distributed workloads.
Key Changes
New Stream Assignment Pass (csrc/host_ir/assign_streams.cpp):
- Transforms stream-parallel loops to use multiple CUDA streams
- Captures the main stream before the loop
- Sets worker streams at the beginning of each iteration with synchronization
- Creates a joining loop after the main loop to synchronize all worker streams back to the main stream
Integration:
- Added
AssignStreamspass to the host IR optimization pipeline (runs afterAllocateAndDeallocate) - Updated CMakeLists.txt to include the new source file
Tests & Benchmarks:
- Refactored
test_row_parallel_linear_forwardto extract reusable fusion definition - Added parametrized benchmark
test_row_parallel_linear_forward_benchmarkwith varying chunk counts (s=1,2,4) - Includes reference PyTorch implementation for validation
Code Cleanup:
- Removed unnecessary includes from header files
- Improved documentation for
Scope::insertmethod - Fixed test function signatures
Architecture
The stream assignment creates a pattern where:
- Main stream is captured before the loop
- Each loop iteration executes on a dedicated worker stream
- Worker streams synchronize with the main stream before starting work
- After the main loop, a joining loop makes the main stream wait for all workers
This enables concurrent execution of loop iterations on different streams, particularly beneficial for overlapping matmul operations with collective communication.
Performance Impact
Benchmark results show slight improvement with overlapping (3.6ms vs 3.7-3.9ms for s=1), with the implementation verified via nsys profiling to correctly assign different streams to different operations.
Confidence Score: 2/5
- This PR contains a critical logic error that will cause incorrect stream synchronization behavior at runtime.
- Score reflects a critical bug in csrc/host_ir/assign_streams.cpp (lines 49-53) where the insertion order of SetCurrentStream and join_loop is inverted, causing the join loop to execute before setting the current stream to main_stream. This means synchronizations will happen on the wrong stream (last worker stream instead of main stream), breaking the intended stream parallelization semantics. Additionally, the copyright year is incorrect (2026 instead of 2025).
- csrc/host_ir/assign_streams.cpp requires immediate attention to fix the insertion order bug on lines 49-53. The expected Host IR comment in tests/python/multidevice/test_overlap.py (line 74-75) should also be updated to include the missing SetCurrentStream statement.
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| csrc/host_ir/assign_streams.cpp | 2/5 | New file implementing stream assignment pass. Contains critical logic error in insertion order (lines 49-53) causing incorrect stream synchronization sequence. Also has incorrect copyright year (2026). |
| csrc/host_ir/assign_streams.h | 5/5 | New header file declaring AssignStreams optimization pass. Clean interface, no issues detected. |
| csrc/host_ir/passes.cpp | 5/5 | Integrates new AssignStreams pass into pipeline. Correct ordering after AllocateAndDeallocate pass. |
| tests/python/multidevice/test_overlap.py | 4/5 | Refactors test into reusable function, adds benchmark test, updates to use outer_split instead of split. Expected Host IR comment may be missing SetCurrentStream statement between loops. |
Sequence Diagram
sequenceDiagram
participant Main as Main Stream
participant W0 as Worker Stream 0
participant W1 as Worker Stream 1
participant W2 as Worker Stream 2
Note over Main: GetCurrentStream
rect rgb(200, 220, 255)
Note over Main,W2: Main Loop (stream-parallel iterations)
Main->>W0: SetCurrentStream(worker_0)
W0->>Main: Synchronize(main_stream)
Note over W0: Compute chunk 0<br/>(matmul + allreduce)
Main->>W1: SetCurrentStream(worker_1)
W1->>Main: Synchronize(main_stream)
Note over W1: Compute chunk 1<br/>(matmul + allreduce)
Main->>W2: SetCurrentStream(worker_2)
W2->>Main: Synchronize(main_stream)
Note over W2: Compute chunk 2<br/>(matmul + allreduce)
end
Note over Main: SetCurrentStream(main_stream)
rect rgb(255, 220, 200)
Note over Main,W2: Join Loop (synchronize workers)
W0->>Main: Synchronize(worker_0)
W1->>Main: Synchronize(worker_1)
W2->>Main: Synchronize(worker_2)
end
Note over Main: Continue on main stream
| @@ -0,0 +1,64 @@ | |||
| // clang-format off | |||
| /* | |||
| * SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES. | |||
There was a problem hiding this comment.
The copyright year is set to 2026, which appears to be incorrect. This should likely be 2025 or earlier.
| * SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES. | |
| * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES. |
| hic->topLevel().insert( | ||
| next_it, IrBuilder::create<SetCurrentStream>(main_stream)); | ||
| auto* join_loop = IrBuilder::create<ForLoop>( | ||
| for_loop->index(), for_loop->start(), for_loop->stop()); | ||
| hic->topLevel().insert(next_it, join_loop); |
There was a problem hiding this comment.
Critical logic error: The insertion order will result in incorrect code generation. Both statements use insert(next_it, ...) which inserts BEFORE the iterator position. Since both use the same next_it, the join_loop (inserted second) will be placed BEFORE the SetCurrentStream (inserted first), resulting in:
FOR loop (original)
join_loop ← wrong position
SetCurrentStream ← wrong position
The correct order should be SetCurrentStream THEN join_loop, because:
- After the first FOR loop exits, the current stream is the last worker stream (not the main stream)
- The join_loop synchronizes worker streams on the current stream
- Therefore, we must set the current stream to main_stream BEFORE the join_loop
This matches the reference implementation row_parallel_linear_forward_reference in test_overlap.py (lines 182-194) where main_stream.wait_stream(worker_stream) happens on the main_stream.
The fix is to save the iterator returned from the first insert and use it for the second insert:
| hic->topLevel().insert( | |
| next_it, IrBuilder::create<SetCurrentStream>(main_stream)); | |
| auto* join_loop = IrBuilder::create<ForLoop>( | |
| for_loop->index(), for_loop->start(), for_loop->stop()); | |
| hic->topLevel().insert(next_it, join_loop); | |
| // After the loop: create a joining loop to synchronize all worker streams | |
| auto set_main_it = hic->topLevel().insert( | |
| next_it, IrBuilder::create<SetCurrentStream>(main_stream)); | |
| auto* join_loop = IrBuilder::create<ForLoop>( | |
| for_loop->index(), for_loop->start(), for_loop->stop()); | |
| hic->topLevel().insert(std::next(set_main_it), join_loop); |
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This PR implements stream parallelization for loops in nvFuser's host IR to enable overlapping of computation (matmul) and communication (allreduce) operations. The core implementation adds an AssignStreams optimization pass that transforms stream-parallel loops by capturing the main stream, assigning worker streams to each iteration with proper synchronization, and creating a joining loop to synchronize all worker streams back to the main stream.
Key Changes:
- New
csrc/host_ir/assign_streams.{cpp,h}implementing the stream assignment pass - Integration of the pass into the host IR optimization pipeline
- Refactored test structure with new benchmark tests
- Code cleanup removing unnecessary includes
Issues Found:
- The joining loop reuses the main loop's index variable, which could cause bugs if code after the loops depends on the index value
- Copyright year typo (2026 instead of 2025) in the new header file
Confidence Score: 3/5
- This PR has one moderate logic issue that should be addressed before merging
- The implementation is generally sound with proper synchronization primitives and good test coverage. However, there's a potential bug where the joining loop reuses the main loop's index variable (line 51 in assign_streams.cpp), which could cause issues if any code depends on the loop index value after the loop completes. This needs verification or fixing. The copyright year typo is a minor style issue.
- Pay close attention to
csrc/host_ir/assign_streams.cpp- specifically the index variable reuse in the joining loop
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| csrc/host_ir/assign_streams.cpp | 3/5 | New file implementing stream parallelization pass. Contains potential bug with loop index reuse in joining loop. |
| csrc/host_ir/assign_streams.h | 4/5 | Header file with copyright year typo (2026 instead of 2025). |
| tests/python/multidevice/test_overlap.py | 4/5 | Refactored test to extract fusion definition, added benchmark test. Changed from torch.randint to torch.testing.make_tensor and torch.randn for test data generation. |
Sequence Diagram
sequenceDiagram
participant Main as Main Stream
participant Pass as AssignStreams Pass
participant ForLoop as Stream Parallel ForLoop
participant W0 as Worker Stream 0
participant W1 as Worker Stream 1
participant WN as Worker Stream N
Pass->>Main: GetCurrentStream
Note over Pass: Capture main stream before loop
Pass->>ForLoop: Transform loop structure
Note over ForLoop,WN: Main Loop Execution
ForLoop->>W0: SetCurrentStream worker 0
W0->>Main: Synchronize with main stream
W0->>W0: Execute iteration 0 body
ForLoop->>W1: SetCurrentStream worker 1
W1->>Main: Synchronize with main stream
W1->>W1: Execute iteration 1 body
ForLoop->>WN: SetCurrentStream worker N
WN->>Main: Synchronize with main stream
WN->>WN: Execute iteration N body
Pass->>Main: SetCurrentStream back to main
Note over ForLoop,WN: Joining Loop Execution
Main->>W0: Synchronize worker stream 0
Main->>W1: Synchronize worker stream 1
Main->>WN: Synchronize worker stream N
Note over Main: All workers joined to main stream
| @@ -0,0 +1,26 @@ | |||
| // clang-format off | |||
| /* | |||
| * SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES. | |||
There was a problem hiding this comment.
copyright year is 2026 (future year) - should be 2025
| auto* join_loop = IrBuilder::create<ForLoop>( | ||
| for_loop->index(), for_loop->start(), for_loop->stop()); |
There was a problem hiding this comment.
reusing for_loop->index() for the joining loop could cause issues. The index variable from the main loop will be overwritten during the join loop execution. If any code after the loops depends on the loop index value, this could lead to bugs. Consider creating a new index variable for the joining loop.
|
!test |
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
Implements stream parallelization for loops in nvFuser's host IR to enable overlapping of computation and communication operations. Adds an AssignStreams optimization pass that transforms loops by capturing the main stream, assigning worker streams to each iteration with proper synchronization, and creating a joining loop to synchronize all worker streams back to the main stream. Includes comprehensive test coverage with benchmarks showing slight performance improvements from overlapping.
Confidence Score: 4/5
- Safe to merge with minor style improvement suggested for iterator handling
- The implementation correctly achieves stream parallelization with proper synchronization semantics as verified by tests. The core logic is sound: capturing the main stream, assigning worker streams per iteration, and synchronizing back via a joining loop. One style issue noted regarding iterator invalidation (lines 44-46 in assign_streams.cpp) that works in practice but could be clearer. No functional bugs or security concerns identified.
- csrc/host_ir/assign_streams.cpp for iterator handling style improvement
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| csrc/host_ir/assign_streams.cpp | 4/5 | Adds stream assignment pass to parallelize loops across streams for overlapping computation and communication |
| csrc/host_ir/passes.cpp | 5/5 | Integrates AssignStreams pass into the host IR optimization pipeline |
| tests/python/multidevice/test_overlap.py | 4/5 | Refactors row parallel linear forward test, adds benchmark tests and updates expected host IR comments |
Sequence Diagram
sequenceDiagram
participant Main as Main Stream
participant W0 as Worker Stream 0
participant W1 as Worker Stream 1
participant W2 as Worker Stream 2
Note over Main: GetCurrentStream
rect rgb(200, 220, 255)
Note over Main,W2: Main Loop (3 iterations)
Main->>W0: SetCurrentStream(0)
W0->>Main: Synchronize(main)
Note over W0: Compute + AllReduce
par Parallel Execution
Main->>W1: SetCurrentStream(1)
W1->>Main: Synchronize(main)
Note over W1: Compute + AllReduce
and
Note over W0: Continue async work
end
par Parallel Execution
Main->>W2: SetCurrentStream(2)
W2->>Main: Synchronize(main)
Note over W2: Compute + AllReduce
and
Note over W0,W1: Continue async work
end
end
rect rgb(255, 220, 200)
Note over Main,W2: Joining Loop (synchronize workers)
Main->>W0: Synchronize(worker 0)
Main->>W1: Synchronize(worker 1)
Main->>W2: Synchronize(worker 2)
end
Main->>Main: SetCurrentStream(main)
| auto old_begin = for_loop->body().exprs().begin(); | ||
| for_loop->body().insert(old_begin, set_stream); | ||
| for_loop->body().insert(old_begin, sync_main); |
There was a problem hiding this comment.
The insertion order places sync_main before set_stream due to how insert works with old_begin. After the first insert at old_begin, set_stream is inserted. When the second insert happens at the same old_begin (which still points to the original beginning), sync_main gets inserted before set_stream.
This creates the sequence: sync_main → set_stream → rest of body. However, the correct order should be set_stream → sync_main because you need to set the worker stream first before synchronizing with the main stream on that worker stream.
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| auto* worker_stream = IrBuilder::create<Stream>(for_loop->index()); | |
| auto* set_stream = IrBuilder::create<SetCurrentStream>(worker_stream); | |
| auto* sync_main = IrBuilder::create<Synchronize>(main_stream); | |
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| auto new_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(std::next(new_begin), sync_main); |
| auto old_begin = for_loop->body().exprs().begin(); | ||
| for_loop->body().insert(old_begin, set_stream); | ||
| for_loop->body().insert(old_begin, sync_main); |
There was a problem hiding this comment.
Using old_begin after the first insert relies on iterator invalidation rules that may lead to undefined behavior. While this works in practice (confirmed by tests), consider capturing the return value for clarity:
| auto old_begin = for_loop->body().exprs().begin(); | |
| for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(old_begin, sync_main); | |
| auto old_begin = for_loop->body().exprs().begin(); | |
| auto it = for_loop->body().insert(old_begin, set_stream); | |
| for_loop->body().insert(std::next(it), sync_main); |
This makes the intent clearer and avoids potential issues with iterator invalidation.
|
!test |
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This PR implements stream parallelization for loops in nvFuser's host IR to enable overlapping of computation and communication operations. The core implementation adds a new AssignStreams optimization pass that transforms stream-parallel loops by:
- Capturing the main stream before the loop begins
- Setting worker streams at the start of each loop iteration and synchronizing with the main stream to ensure proper ordering
- Creating a join loop after the main loop that synchronizes all worker streams back to the main stream
The implementation correctly handles:
- Stream assignment and synchronization order (SetCurrentStream → Synchronize → work)
- Reusing the loop index variable for both the main and join loops (as shown in expected IR)
- Integration into the host IR pass pipeline after allocation/deallocation
Additional changes include:
- Test refactoring to use
outer_split()API instead ofsplit(axis, factor, inner_split=False) - New benchmark test for performance measurement with different chunk sizes
- Documentation updates for profiling multi-process applications
- Cleanup of unused includes and test fixture parameters
Benchmark results show slight performance improvements from overlapping, with the implementation verified via nsys profiling showing proper stream assignment and overlapping of matmul and allreduce operations.
Confidence Score: 5/5
- Safe to merge - implementation is correct with proper stream synchronization semantics
- The implementation correctly follows CUDA stream semantics for parallelization. The insertion order of SetCurrentStream followed by Synchronize is correct (verified by tracing through std::vector::insert semantics). The reuse of loop index variables between main and join loops is intentional and matches the expected IR output documented in tests. All changes are well-tested with both functional tests and benchmarks.
- No files require special attention
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| csrc/host_ir/assign_streams.cpp | 5/5 | New pass implementing stream parallelization for loops - correctly inserts stream setup/synchronization and creates join loop for cleanup |
| csrc/host_ir/assign_streams.h | 5/5 | Header for AssignStreams optimization pass - standard pass declaration with no issues |
| csrc/host_ir/passes.cpp | 5/5 | Integrates AssignStreams pass into host IR pipeline after AllocateAndDeallocate |
| tests/python/multidevice/test_overlap.py | 5/5 | Refactors test into reusable function, uses outer_split API, adds benchmark test, and updates expected host IR documentation |
Sequence Diagram
sequenceDiagram
participant Main as Main Stream
participant Pass as AssignStreams Pass
participant W0 as Worker Stream 0
participant W1 as Worker Stream 1
participant WN as Worker Stream N
Pass->>Main: GetCurrentStream()
Note over Pass: Transform loop body
Note over Pass,WN: Original Loop Iterations
Pass->>W0: SetCurrentStream(Stream 0)
W0->>Main: Synchronize(main_stream)
W0->>W0: Execute iteration 0 work (matmul + allreduce)
Pass->>W1: SetCurrentStream(Stream 1)
W1->>Main: Synchronize(main_stream)
W1->>W1: Execute iteration 1 work (matmul + allreduce)
Pass->>WN: SetCurrentStream(Stream N)
WN->>Main: Synchronize(main_stream)
WN->>WN: Execute iteration N work (matmul + allreduce)
Pass->>Main: SetCurrentStream(main_stream)
Note over Pass,WN: Join Loop - Synchronize All Workers
Main->>W0: Synchronize(Stream 0)
Main->>W1: Synchronize(Stream 1)
Main->>WN: Synchronize(Stream N)
Fixes #5308
Overlapping improves the wall time slightly.
Stream assignment and overlapping are verified by the following:
The performance is suboptimal for two reasons: