[RELAX][LAYOUT] Support multiple axis paching#18869
[RELAX][LAYOUT] Support multiple axis paching#18869srkreddy1238 wants to merge 1 commit intoapache:mainfrom
Conversation
Summary of ChangesHello, 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 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++ Highlights
🧠 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 to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe 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 by creating a comment using either
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for packed axes in data layouts, like OIHW[4o4i], which is a significant feature for handling complex tensor layouts. No security vulnerabilities were found. However, there are a few issues that need to be addressed: a potential bug in the Python __contains__ implementation for Layout, a likely bug in the C++ GetStoreRule function due to incorrect vector manipulation, an unused variable in GetStoreRule, and an inconsistency between the documentation and implementation of PackIterVar. Additionally, there's a typo in the PR title: "paching" should be "packing".
| for (size_t j = 0; j < src_unpacked_axes.size(); j++) { | ||
| index_divs[j] = value; | ||
| const auto* extent = src_unpacked_axes[j]->dom->extent.as<IntImmNode>(); | ||
| TVM_FFI_ICHECK(extent) << "Expected Integer Extents for Offset Calculation"; | ||
| index_divs.push_back(value); | ||
| value = value * extent->value; | ||
| } |
There was a problem hiding this comment.
There appears to be a bug in this loop. The vector index_divs is initialized with a specific size at line 330. Inside the loop, both index_divs[j] = value; and index_divs.push_back(value); are called. This will result in a vector of double the intended size with incorrect contents, which will likely cause incorrect behavior in layout transformations. The push_back call should probably be removed.
for (size_t j = 0; j < src_unpacked_axes.size(); j++) {
index_divs[j] = value;
const auto* extent = src_unpacked_axes[j]->dom->extent.as<IntImmNode>();
TVM_FFI_ICHECK(extent) << "Expected Integer Extents for Offset Calculation";
value = value * extent->value;
}| def __contains__(self, axis): | ||
| return len(axis) == 1 and axis[0].isalpha() and axis[0] in self.name | ||
| # Note: We do a weaker check for packed axis assuming layout is valid | ||
| return not any(bkt in axis for bkt in "[]") and axis in self.name |
There was a problem hiding this comment.
The current implementation of __contains__ using a substring check (axis in self.name) can lead to incorrect results and is inconsistent with index_of. For example, for a layout NCHW16c, __contains__("16") would return True, which is incorrect as "16" is not an axis. Similarly, for OIHW[4o4i], "i" in layout is True while layout.index_of("i") is -1.
A more robust implementation should check against the actual axes of the layout. The current substring check is too broad and gives false positives.
| for (size_t j = 0; j < src_unpacked_axes.size(); j++) { | ||
| const int extent = src_unpacked_axes[j]->dom->extent.as<IntImmNode>()->value; | ||
| const LayoutAxis& store_axis_impl = LayoutAxis::Get(src_unpacked_axes[j]); | ||
| const LayoutAxis& sub_axis = store_axis_impl.ToSubordinate(); /* Not Needed */ |
There was a problem hiding this comment.
The comment /* Not Needed */ suggests this variable sub_axis is unused. If it's indeed not needed, it should be removed to improve code clarity and avoid confusion.
| * \brief Packs the Given Array of IterVars into a Single IterVar. Each IterVar in the Array | ||
| * should represent either a single primal axis or one or more subordinate axis | ||
| * \param iters Array of iter vars to be packed | ||
| * \return A packed iter var | ||
| */ |
There was a problem hiding this comment.
The documentation for PackIterVar states that it can pack a "single primal axis or one or more subordinate axis". However, the implementation in src/s_tir/data_layout.cc only supports packing subordinate axes and includes a check Packed Axis can contain only Subordinate Axes. This discrepancy should be resolved. Please update the documentation to reflect the implementation's behavior.
srkreddy1238
changed the title
Like OIHW[4o4i] where we can pack multiple axis. Helpful while handling complex target layouts. This PR covers layout representation and transforms for these.
srkreddy1238
force-pushed
the
layout-packed-rebase
branch
from
0b54493 to
84e03e0
Compare
1 participant
Like OIHW[4o4i] where we can pack multiple axis.
Helpful while handling complex target layouts.
This PR covers layout representation and transforms for these.