[Pipeline] Phaseout fragment and double buffer info from pipeline pass #711
Conversation
…ndition handling - Introduced a new function to replace IfThenElse nodes with their then_case while preserving attributes. - Streamlined the PipelineBodyRewriter to improve buffer access rewriting and async state management. - Enhanced the handling of pipeline loop conditions and added support for predicate conditions in the pipeline body. - Removed obsolete code and improved overall code clarity and maintainability.
|
👋 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! 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @LeiWang1999, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly refactors the pipeline pass by phasing out the use of fragment and double buffer information. The changes involve removing specialized rewriters and associated data structures, simplifying constructors, and integrating necessary buffer access logic directly into core components. This leads to a cleaner and more streamlined pipeline transformation process, reducing complexity and dependencies on outdated concepts.
Highlights
- Removal of specialized rewriter: The
PipelineOpaqueAccessRewriterclass, which handled specific buffer access rewrites for tensor core fragments, has been completely removed. This simplifies the pipeline pass by eliminating a specialized rewriter. - Simplification of
PipelineBodyRewriter: Thefragment_infoparameter and theopaque_access_rewriter_member were removed fromPipelineBodyRewriter, streamlining its constructor and dependencies. - Relocation of buffer access logic: The general
RewriteBufferAccessmethod was moved from the removedPipelineOpaqueAccessRewriterdirectly intoPipelineBodyRewriter, indicating that general buffer access rewriting is still necessary but is now integrated more directly into the main body rewriting logic. - Streamlined
PipelineRewriterinterface: Parameters related todouble_buffersandfragment_infowere removed from thePipelineRewriterconstructor and its staticRewritemethod, simplifying its interface and internal state. - Phase-out of double-buffering logic: The logic for handling
double_bufferswas removed, including a specific check fornum_versions == 1 && double_buffers_.count(buffer), indicating that this particular double-buffering mechanism is no longer managed within this pass. - Refined asynchronous state management: The internal structs and logic for managing asynchronous states (
AsyncStateLocal) were simplified, removing fields likeseen,predicate, andcommit_groups, and updating how wait counts are populated and loop statements are completed. - Simplified loop emission: The
EmitImplmethod, responsible for emitting loop statements, was simplified by removing theextra_loop_lower_boundparameter and adding aneed_bound_checkboolean, making the loop emission logic more direct. - Cleanup in
PipelineInjector: Thefragment_info_anddouble_buffersmembers were removed fromPipelineInjector, aligning with the overall phase-out of these concepts across the pipeline pass.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request refactors the software pipeline injection pass to remove dependencies on fragment_info and double_buffer_info, simplifying the overall logic. It also introduces significant improvements to the handling of asynchronous pipeline stages, making the logic more robust and clear.
My review found one issue: the error messages for missing pipeline annotations in src/transform/inject_pipeline.cc have been swapped, which could be misleading. I've provided a suggestion to correct this. Otherwise, the changes look good and are a positive step towards a cleaner pipeline implementation.
| if (has_stage) { | ||
| LOG(FATAL) | ||
| << "ValueError: Order of the software pipeline is not defined."; | ||
| << "ValueError: Stage of the software pipeline is not defined."; | ||
| } | ||
| if (has_order) { | ||
| LOG(FATAL) | ||
| << "ValueError: Stage of the software pipeline is not defined."; | ||
| << "ValueError: Order of the software pipeline is not defined."; | ||
| } |
There was a problem hiding this comment.
The error messages for missing pipeline annotations seem to have been swapped in this change, which could be misleading for developers.
- If
has_stageis true buthas_orderis false, the error should indicate that theorderannotation is missing. - If
has_orderis true buthas_stageis false, the error should indicate that thestageannotation is missing.
The previous implementation had the correct error messages.
| if (has_stage) { | |
| LOG(FATAL) | |
| << "ValueError: Order of the software pipeline is not defined."; | |
| << "ValueError: Stage of the software pipeline is not defined."; | |
| } | |
| if (has_order) { | |
| LOG(FATAL) | |
| << "ValueError: Stage of the software pipeline is not defined."; | |
| << "ValueError: Order of the software pipeline is not defined."; | |
| } | |
| if (has_stage) { | |
| LOG(FATAL) | |
| << "ValueError: Order of the software pipeline is not defined."; | |
| } | |
| if (has_order) { | |
| LOG(FATAL) | |
| << "ValueError: Stage of the software pipeline is not defined."; | |
| } |
…y std::move calls - Updated return statements in multiple methods to return objects directly instead of using std::move, improving code clarity and potentially avoiding unnecessary moves. - Ensured consistent handling of BufferStore and BufferLoad nodes during pipeline transformations.
tile-ai#711) * Update submodule 'tvm' to commit e11521e6936a827efa334588d29571fbb4620107 * Refactor inject_pipeline.cc to enhance pipeline body rewriting and condition handling - Introduced a new function to replace IfThenElse nodes with their then_case while preserving attributes. - Streamlined the PipelineBodyRewriter to improve buffer access rewriting and async state management. - Enhanced the handling of pipeline loop conditions and added support for predicate conditions in the pipeline body. - Removed obsolete code and improved overall code clarity and maintainability. * lint fix * Refactor return statements in inject_pipeline.cc to remove unnecessary std::move calls - Updated return statements in multiple methods to return objects directly instead of using std::move, improving code clarity and potentially avoiding unnecessary moves. - Ensured consistent handling of BufferStore and BufferLoad nodes during pipeline transformations. * test fix
1 participant
as title.