[Enhancement] Refactor inflight computing to support dynamic pipeline extents #1399
Conversation
…lation - Adjusted output directories for the tilelang_cython_wrapper to ensure that development builds place the extension in build/lib. - Updated installation paths to place the extension in tilelang/lib within the wheel, improving organization and avoiding potential conflicts with other modules. - Modified the internal library path exposure in env.py to prevent shadowing of common module names, enhancing compatibility and usability in user projects.
- Set output directories for both tilelang and tilelang_module libraries to "${CMAKE_BINARY_DIR}/lib" for consistency in development builds.
- This change enhances organization and ensures that all build artifacts are located in a unified directory structure.
- Updated the TVM subproject to commit 90581fe9e5287bbcf1844ad14255a1e1e8cdf7f0. - Added new fields to `PipelineAnnotation` and `RewrittenBlockInfo` structures to track original statement indices and improve async state management. - Refactored `EmitImpl` and `PopulateWaitCounts` methods to enhance clarity and functionality, including better handling of commit groups and wait counts. - Simplified access index calculations and strengthened analyzer constraints for loop bounds.
…ne.cc - Eliminated the Apache license block from the top of the file to streamline the code. - Removed unused include directives for memory and stringstream to enhance code clarity and reduce unnecessary dependencies.
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
WalkthroughUpdates a TVM submodule reference and extends the software-pipeline injection: adds per-block original index and start/end iteration bounds, makes emission and wait-count logic epilogue-aware, records per-async-stage commit predicates, and adjusts tests and phase pipeline to run extra Simplify passes. Changes
Sequence DiagramsequenceDiagram
participant Collector as BlockCollector
participant Annotator as PipelineAnnotator
participant Emitter as Emitter
participant Waiter as AsyncWaitComputer
Collector->>Annotator: collect blocks (index i)
Annotator->>Annotator: create PipelineAnnotation(original_idx = i)
Annotator->>Annotator: record block start/end in RewrittenBlockInfo
rect rgb(220,240,255)
Note over Emitter: Prologue / Body emission (is_epilogue=false)
Collector->>Emitter: EmitImpl(start,end,..., is_epilogue=false)
Emitter->>Emitter: bind loop var domain\napply bound guards
end
rect rgb(255,240,220)
Note over Emitter,Waiter: Epilogue emission (is_epilogue=true)
Collector->>Emitter: EmitImpl(start,end,..., is_epilogue=true)
Emitter->>Waiter: PopulateWaitCounts(new_blocks, async_states, is_epilogue=true)
Waiter->>Waiter: build order->access index map\nidentify dependent commit groups
Waiter->>Waiter: compute waits (epilogue-aware)\nrecord commit_predicate in AsyncStateLocal
Waiter-->>Emitter: store pending_waits per block
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
🧰 Additional context used🧬 Code graph analysis (1)tilelang/engine/phase.py (2)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
🔇 Additional comments (2)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
- Added an additional Simplify transformation in the InjectSoftwarePipeline to improve optimization. - Updated the test file to call `test_trival_pipeline()` directly, commenting out the previous main execution for better test isolation.
1 participant
This pull request updates the TVM submodule and refactors the pipeline injection logic in
inject_pipeline.ccto improve correctness and robustness of async pipeline scheduling. The main changes include more accurate wait count computation for async dependencies, enhanced symbolic analysis for loop bounds, and richer metadata for pipeline blocks.Pipeline scheduling and async dependency handling:
PopulateWaitCountsto precisely track async producer-consumer relationships using commit groups, buffer dependencies, and symbolic analysis. This ensures correct synchronization and reduces unnecessary waits.commit_predicateto async state tracking and propagate it through the pipeline, improving the handling of conditional execution of commit blocks. [1] [2]Loop and block metadata improvements:
PipelineAnnotationandRewrittenBlockInfostructs to includeoriginal_idx,start, andendfields, enabling more accurate mapping and analysis of block order and loop bounds. [1] [2] [3]General codebase updates:
<memory>and<sstream>to support new features in the refactored code.Summary by CodeRabbit
Refactor
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.