feat(language): add Tile type and block operations to Language DSL #112
Conversation
Adds comprehensive Tile type support for block-level programming: - New `pl.Tile[[shape], dtype]` type annotation mirroring Tensor pattern - Block operations wrapper (load, store, add, mul, sqrt, sum, etc.) - Type resolver support for Tile with 2D shape validation - Updated tests to use modern @pl.program decorator pattern - Documentation updates for operator organization and syntax
Summary of ChangesHello @Hzfengsy, 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 PyPTO Language DSL by integrating a new 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. 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 comprehensive support for a Tile type and block-level operations in the language DSL. The changes are well-structured, including updates to the language API, parser, and documentation. The refactoring of the test_basic_memory_reuse.py test to use the new @pl.program decorator is a significant improvement, making the tests more readable and modern. My review includes a couple of suggestions to improve documentation consistency and reduce code duplication in the new block_ops.py module.
| def add(lhs: Tile, rhs: Tile) -> Tile: | ||
| """Element-wise addition of two tiles. | ||
|
|
||
| Args: | ||
| lhs: Left-hand side tile | ||
| rhs: Right-hand side tile | ||
|
|
||
| Returns: | ||
| Tile wrapping the add operation | ||
| """ | ||
| call_expr = _ir_ops.add(lhs.unwrap(), rhs.unwrap()) | ||
| return Tile(expr=call_expr) | ||
|
|
||
|
|
||
| def sub(lhs: Tile, rhs: Tile) -> Tile: | ||
| """Element-wise subtraction of two tiles. | ||
|
|
||
| Args: | ||
| lhs: Left-hand side tile | ||
| rhs: Right-hand side tile | ||
|
|
||
| Returns: | ||
| Tile wrapping the sub operation | ||
| """ | ||
| call_expr = _ir_ops.sub(lhs.unwrap(), rhs.unwrap()) | ||
| return Tile(expr=call_expr) |
There was a problem hiding this comment.
There is a lot of repetitive code for creating the block operations. You can reduce this boilerplate by using helper functions for common patterns like unary and binary operations. This will improve maintainability.
For example, for binary operations like add, sub, mul, div, and maximum, you could use a helper like this:
from collections.abc import Callable
def _create_binary_op(op_func: Callable[[Expr, Expr], Expr], lhs: Tile, rhs: Tile) -> Tile:
"""Helper to create a binary operation on two tiles."""
call_expr = op_func(lhs.unwrap(), rhs.unwrap())
return Tile(expr=call_expr)
def add(lhs: Tile, rhs: Tile) -> Tile:
"""Element-wise addition of two tiles."""
return _create_binary_op(_ir_ops.add, lhs, rhs)
def sub(lhs: Tile, rhs: Tile) -> Tile:
"""Element-wise subtraction of two tiles."""
return _create_binary_op(_ir_ops.sub, lhs, rhs)Similar helpers can be created for unary operations and operations with scalars.
There was a problem hiding this comment.
Pull request overview
Adds first-class Tile support to the PyPTO Language DSL, enabling type-annotated block/tile programming alongside existing Tensor patterns.
Changes:
- Introduces
pl.Tile[[shape], dtype]wrapper type and exports it frompypto.language. - Adds
pl.op.block.*language-level wrappers for block/tile ops (load/store/elementwise/unary/reduction/matmul). - Extends the language type resolver to parse
Tileannotations (subscript and legacy call forms), updates docs, and modernizes the BasicMemoryReuse tests to the@pl.programstyle.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/ut/ir/transforms/test_basic_memory_reuse.py | Refactors memory reuse tests to the @pl.program DSL style and uses pl.Tile annotations + pl.op.block calls. |
| python/pypto/language/tile.py | Adds the Tile language wrapper type (annotation + runtime wrapper). |
| python/pypto/language/parser/type_resolver.py | Adds Tile type annotation parsing and rank validation (<=2 dims). |
| python/pypto/language/op/block_ops.py | Adds language-level, type-safe wrappers around IR block_ops that return/accept Tile/Tensor wrappers. |
| python/pypto/language/op/init.py | Exposes pl.op.block alongside pl.op.tensor. |
| python/pypto/language/init.py | Re-exports Tile and documents block/tile usage in module docstring. |
| docs/dev/07-python_syntax.md | Updates syntax docs and adds a Tile/block example section. |
| docs/dev/06-operator_organization.md | Adds a high-level Language DSL example for block ops. |
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 participant
Adds comprehensive Tile type support for block-level programming:
pl.Tile[[shape], dtype]type annotation mirroring Tensor pattern