Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Host Ir] stream lowering, first milestone #4148

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from

Conversation

samnordmann
Copy link
Collaborator

@samnordmann samnordmann commented Mar 26, 2025

Copy link

Description

  • Added support for ParallelType::Stream lowering in Host IR for single device fusions.

  • Introduced new tests for stream parallelization in test_host_ir_stream_lowering.cpp.

  • Enhanced HostIrEvaluator to handle pre-allocated outputs and stream management.

  • Updated HostIrLower to support standalone host operations and custom merging logic.


Changes walkthrough 📝

Relevant files
Tests
4 files
test_host_ir_stream_lowering.cpp
Added comprehensive tests for stream parallelization         
+823/-0 
test_host_irs.cpp
Added tests for pre-allocated outputs and binary operations
+290/-4 
test_multidevice_host_ir.cpp
Disabled StreamParallelType pass temporarily in some tests
+10/-0   
test_resharding.cpp
Updated SegmentCandidateFinderOptions for custom merging 
+1/-1     
Enhancement
16 files
executor.cpp
Enhanced HostIrEvaluator for stream management and pre-allocated
outputs
+207/-106
stream_parallel_type.cpp
Implemented StreamParallelType pass for stream parallelization
+386/-0 
lower.cpp
Added logic to handle standalone host operations and custom merging
+72/-26 
utils.cpp
Updated newOutputDomain and newValLike to support reduction axes
+19/-8   
fusion_segmenter.cpp
Updated codeGenSupportedMerge and deriveSchedulerType for custom
merging
+4/-11   
indexing.cpp
Updated select to support reduction axes                                 
+8/-2     
container.cpp
Added print method for aliases in HostIrContainer               
+5/-0     
executor.h
Added methods for alias handling and stream management in
HostIrEvaluator
+39/-1   
stream_parallel_type.h
Added StreamParallelType pass declaration                               
+35/-0   
utils.h
Updated function signatures for newOutputDomain and newValLike
+11/-3   
container.h
Added methods for alias management and resetting top-level expressions
+16/-0   
lower.h
Added methods for standalone host operations and custom merging logic
+7/-0     
indexing.h
Updated select function signature to support reduction axes
+5/-1     
executor.h
Added method to access HostIrEvaluator                                     
+4/-0     
fusion_segmenter.h
Updated SegmentCandidateFinderOptions for custom merging 
+2/-1     
internal_nodes.h
Added iterDomain method to ForLoop                                             
+4/-0     
Configuration changes
1 files
CMakeLists.txt
Added new source and test files                                                   
+2/-0     
Additional files
1 files
test_multidevice_pipeline.cpp +0/-131 

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests
⚡ Recommended focus areas for review


Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant