feat(codegen): Add PTOAS-aware type annotations and update PTO naming#267
feat(codegen): Add PTOAS-aware type annotations and update PTO naming#267Hzfengsy merged 1 commit intohw-native-sys:mainfrom
Conversation
Summary of ChangesHello @zhangqi-chen, 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 PTO Codegen by introducing robust type awareness and aligning its output with the latest PTO-ISA specification. The changes ensure that generated MLIR code includes precise type annotations for tensor views and tile buffers, improving the correctness and readability of the intermediate representation. By updating naming conventions and dynamically deriving type information, the codegen becomes more accurate and maintainable, facilitating better integration with the PTO backend. Highlights
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
|
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughPTOCodegen class refactored to inherit from CodegenBase with relocated header file and expanded public API. Introduces type-annotated code generation with partition_view semantics for load/store operations. Helper methods added for type annotation and tile buffer type generation. MemRef-to-TileType mapping supports richer type information. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/codegen/pto/pto_codegen.cpp (1)
500-521: Consider handling 0D and 3+D tile shapes.The implementation handles 1D and 2D shapes but falls back to defaults (32x32) for 0D or 3+D shapes. If 3D or higher-dimensional tiles are valid in this context, consider adding explicit handling or a log warning.
💡 Suggested enhancement
if (tile_type->shape_.size() >= 2) { if (auto c0 = As<ir::ConstInt>(tile_type->shape_[0])) rows = c0->value_; if (auto c1 = As<ir::ConstInt>(tile_type->shape_[1])) cols = c1->value_; } else if (tile_type->shape_.size() == 1) { if (auto c0 = As<ir::ConstInt>(tile_type->shape_[0])) { rows = 1; cols = c0->value_; } + } else if (tile_type->shape_.empty()) { + // Scalar tile - use 1x1 + rows = 1; + cols = 1; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/codegen/pto/pto_codegen.cpp` around lines 500 - 521, PTOCodegen::GetTileBufTypeString currently only handles 1D and 2D tile_type->shape_ and leaves rows/cols as default 32 for 0D or >=3D shapes; update the logic in GetTileBufTypeString (and use memref_to_tile_type_/tile_type->shape_) to explicitly handle 0D by setting rows=1, cols=1, and handle 3D+ by either flattening higher dims (e.g., set rows = product(shape[0..n-2]) and cols = shape[n-1]) or emit a diagnostic/log warning and choose a deterministic mapping, and ensure dtype_str is still set via GetTypeString(tile_type->dtype_); keep behavior consistent and add a clear comment explaining the chosen mapping.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@src/codegen/pto/pto_codegen.cpp`:
- Around line 500-521: PTOCodegen::GetTileBufTypeString currently only handles
1D and 2D tile_type->shape_ and leaves rows/cols as default 32 for 0D or >=3D
shapes; update the logic in GetTileBufTypeString (and use
memref_to_tile_type_/tile_type->shape_) to explicitly handle 0D by setting
rows=1, cols=1, and handle 3D+ by either flattening higher dims (e.g., set rows
= product(shape[0..n-2]) and cols = shape[n-1]) or emit a diagnostic/log warning
and choose a deterministic mapping, and ensure dtype_str is still set via
GetTypeString(tile_type->dtype_); keep behavior consistent and add a clear
comment explaining the chosen mapping.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
docs/dev/codegen/00-pto_codegen.mdinclude/pypto/codegen/pto/pto_codegen.hsrc/backend/910B_PTO/backend_910b_pto_ops.cppsrc/codegen/pto/pto_codegen.cpptests/ut/codegen/test_pto_codegen.pytests/ut/codegen/test_pto_codegen_ops.py
![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 significantly enhances the PTO codegen by making it more aware of TileType metadata, leading to more accurate MLIR type generation. The renaming of pto.subview to pto.partition_view and the update of memory space names align the codegen with the latest PTO-ISA specification. The addition of type annotations to ins/outs clauses is also a valuable improvement for MLIR readability and correctness. The changes are well-implemented across the documentation, C++ source, and tests. I've included a couple of suggestions to further improve documentation accuracy and make the tile_buf type generation even more metadata-driven.
Emit accurate tile_buf/tensor_view type strings derived from TileType metadata instead of hardcoded defaults. Extract blayout, slayout, fractal, and pad from TileView when available, falling back to sensible defaults. Rename pto.subview to pto.partition_view and update memory space names (ub->vec, ddr->gm, l1->mat, etc.) to match current PTO-ISA specification. Add ins/outs type annotation support to operator codegen callbacks. Update codegen documentation accordingly.
2 participants
Emit accurate tile_buf/tensor_view type strings derived from TileType metadata instead of hardcoded defaults. Rename pto.subview to pto.partition_view and update memory space names (ub->vec, ddr->gm, l1->mat, etc.) to match current PTO-ISA specification. Add ins/outs type annotation support to operator codegen callbacks.