gemini-code-assist[bot] commented on PR #18869: URL: https://github.com/apache/tvm/pull/18869#issuecomment-3995134966
## Summary of Changes Hello, I'm Gemini Code Assist[^1]! 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 enhances the data layout system by introducing support for multi-axis packing, allowing for more flexible and expressive tensor layouts. This capability is crucial for effectively mapping and optimizing computations on diverse hardware architectures. The changes involve fundamental modifications to the C++ `Layout` class, its parsing and transformation logic, and corresponding updates to the Python API and test suite, ensuring robust handling of these new packed dimensions. ### Highlights * **Multi-Axis Packing Support**: Introduced the ability to pack multiple axes into a single logical dimension within data layouts, such as `OIHW[4o4i]`, to support more complex target layouts. * **Layout Class Enhancements**: Added `PackIterVar` and `UnpackIterVar` static methods to the `Layout` class, and updated existing methods like `ndim_primal`, `IndexOf`, `Contains`, and `PackedAxisAt` to correctly process these new packed axes. * **Layout String Parsing and Validation**: Modified the `Layout` constructor to parse layout strings containing packed axes (e.g., `[4o4i]`) and added comprehensive validation checks for their structure. * **Transformation Rule Refactoring**: Refactored the `GetStoreRule` function to accurately handle index and shape transformations when dealing with layouts that include packed axes, ensuring correct data mapping. * **Python API Updates and Testing**: Updated Python bindings for `Layout` to reflect the new packed axis behavior and added extensive test cases to validate the functionality, including error handling for invalid packed layout strings. 🧠**New Feature in Public Preview:** You can now enable **Memory** to help **Gemini Code Assist** learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. **Click [here](https://codeassist.google/code-review/login) to enable Memory in your admin console.** <details> <summary><b>Changelog</b></summary> * **include/tvm/s_tir/data_layout.h** * Added `tvm/tir/var.h` include for IterVar definitions. * Introduced `static IterVar PackIterVar(ffi::Array<IterVar> iters)` to combine multiple IterVars into a single packed IterVar. * Introduced `static ffi::Array<IterVar> UnpackIterVar(IterVar packed_iter)` to decompose a packed IterVar into its constituent IterVars. * Modified `ndim_primal()` to iterate through unpacked IterVars when counting primal axes. * Overloaded `IndexOf` to accept `std::string` and `tir::IterVar` for more flexible axis lookup, and updated the `LayoutAxis` version to use the string-based one. * Updated `Contains()` to check for axis presence within the unpacked IterVars of each packed axis. * Added `PackedAxisAt(int32_t i)` to retrieve the packed IterVar at a specific index. * **python/tvm/s_tir/data_layout.py** * Updated `__contains__` method to perform a weaker check for packed axes, assuming layout validity. * Modified `index_of` docstring to indicate support for packed axis names. * **src/s_tir/data_layout.cc** * Added necessary includes for `tvm/ir/expr.h`, `tvm/runtime/data_type.h`, `tvm/runtime/logging.h`, `tvm/runtime/object.h`, `tvm/tir/analysis.h`, `tvm/tir/expr.h`, `tvm/tir/var.h`, and `algorithm`. * Refactored `Layout` constructor to parse layout strings, including support for `[ ]` syntax to define packed axes, and added validation for nested packing and variable-sized axes within packed dimensions. * Implemented the logic for `Layout::UnpackIterVar` to parse the name hint of a packed IterVar and reconstruct its constituent IterVars. * Implemented the logic for `Layout::PackIterVar` to concatenate names and multiply extents of multiple IterVars into a single packed IterVar. * Modified `Layout::FactorOf` to correctly calculate the factor of a subordinate axis by iterating through unpacked IterVars. * Rewrote `GetStoreRule` to handle index and shape transformations for layouts with packed axes, introducing `norm_indexes` and `index_divs` for complex offset calculations. * Updated `TransformShape` to check if an axis is primal by first unpacking it and then checking its constituent parts. * Adjusted FFI bindings for `s_tir.LayoutIndexOf` and `s_tir.LayoutGetItem` to align with the new packed axis representation and string-based lookup. * **tests/python/s_tir/base/test_tir_data_layout.py** * Updated imports to use `tvm.testing` and `InternalError`. * Modified `test_layout` to include assertions for `index_of` and `__contains__` with packed axis names like '16c'. * Added new test cases for `OIHW[4o4i]` layout, verifying `factor_of`, `index_of`, `__contains__`, and `__getitem__` behavior. * Introduced `pytest.raises(InternalError)` tests for various invalid packed layout string formats. * Added new `bijective_layout` tests for conversions between `OIHW` and `OIHW[4o4i]`, including `forward_shape`, `backward_shape`, `forward_index`, and `backward_index` checks. * Changed the test execution block to use `tvm.testing.main()`. </details> <details> <summary><b>Activity</b></summary> * The pull request introduces significant new functionality for data layout management. * It involves changes across C++ headers, implementation files, Python bindings, and test cases. </details> <details> <summary><b>Using Gemini Code Assist</b></summary> <br> The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips. <b>Invoking Gemini</b> You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page. 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 pull request comments and review comments. Help | `/gemini help` | Displays a list of available commands. <b>Customization</b> 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](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). <b>Limitations & Feedback</b> Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: 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](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG). <b>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](https://cloud.google.com/products/gemini/code-assist).</b> </details> [^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
