Add TGATHERB#932
Conversation
![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 a new TileLang DSL template for the pto.tgatherb operation along with corresponding test cases in test/tilelib-st/a5/tgatherb/case.py and minor .gitignore updates. Feedback on the changes includes correcting the dtypes specification in the @pto.vkernel decorator to prevent type-matching failures for non-i8 destinations, removing unused imports and variables in the template, and adding an assertion or dynamic masking in the test kernel to handle column sizes that are not multiples of 16 safely.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| dtypes=[ | ||
| (pto.i8, pto.f32, pto.f32, pto.f32, pto.f32), | ||
| pto.AnyInt, | ||
| (pto.i8, pto.i8, pto.i8, pto.i8, pto.i8), | ||
| ], |
There was a problem hiding this comment.
The dtypes specification in the @pto.vkernel decorator contains redundant duplicate types and restricts the destination tile (dst) to only pto.i8 (repeated 5 times). Since the template is generic and designed to support other types like pto.f32 (as demonstrated in the test cases), this restriction will cause type-matching failures when using pto.tgatherb with non-i8 destinations. Please update the allowed types to correctly reflect the supported data types (e.g., pto.i8, pto.f16, pto.bf16, pto.f32) for both src and dst tiles.
| dtypes=[ | |
| (pto.i8, pto.f32, pto.f32, pto.f32, pto.f32), | |
| pto.AnyInt, | |
| (pto.i8, pto.i8, pto.i8, pto.i8, pto.i8), | |
| ], | |
| dtypes=[ | |
| (pto.i8, pto.f16, pto.bf16, pto.f32), | |
| pto.AnyInt, | |
| (pto.i8, pto.f16, pto.bf16, pto.f32), | |
| ], |
| import sys | ||
| from pathlib import Path |
There was a problem hiding this comment.
The imports sys and Path are not used in this file and should be removed to keep the code clean.
| lanes = pto.get_lanes(dtype) | ||
| repeat_times = (valid_cols + lanes - 1) // lanes + 1 |
There was a problem hiding this comment.
The variable repeat_times is calculated but never used in the function. It should be removed to keep the code clean and maintainable.
| lanes = pto.get_lanes(dtype) | |
| repeat_times = (valid_cols + lanes - 1) // lanes + 1 | |
| lanes = pto.get_lanes(dtype) |
| def _tgatherb_body(src_ptr, offset_ptr, dst_ptr, *, src_rows, src_cols, dst_rows, dst_cols): | ||
| """Shared kernel body for the tgatherb cases. | ||
|
|
||
| Loads *src* and *offset* tiles from GM, performs block-gather using | ||
| ``pto.vgatherb``, and stores *dst* back to GM. | ||
| """ |
There was a problem hiding this comment.
The _tgatherb_body function uses a hardcoded loop step of 16 and an all-enabled mask (pto.MaskPattern.ALL). If a test case with a column size that is not a multiple of 16 is added in the future, this will lead to out-of-bounds memory accesses or writing garbage data. Please add an assertion to ensure dst_cols is a multiple of 16, or implement dynamic masking using pto.make_mask as done in the template.
| def _tgatherb_body(src_ptr, offset_ptr, dst_ptr, *, src_rows, src_cols, dst_rows, dst_cols): | |
| """Shared kernel body for the tgatherb cases. | |
| Loads *src* and *offset* tiles from GM, performs block-gather using | |
| ``pto.vgatherb``, and stores *dst* back to GM. | |
| """ | |
| def _tgatherb_body(src_ptr, offset_ptr, dst_ptr, *, src_rows, src_cols, dst_rows, dst_cols): | |
| """Shared kernel body for the tgatherb cases. | |
| Loads *src* and *offset* tiles from GM, performs block-gather using | |
| ``pto.vgatherb``, and stores *dst* back to GM. | |
| """ | |
| assert dst_cols % 16 == 0, "dst_cols must be a multiple of 16 for the hardcoded mask" |
Codex Review该评论由 review 机器人自动更新。
SummaryReview failed at stage Findings未生成结构化 findings，因为 review 过程提前失败。 Log Tail |
f58bde3 to
29233e4
Compare
| valid_rows, valid_cols = dst.valid_shape | ||
| lanes = pto.elements_per_vreg(dtype) | ||
| repeat_times = (valid_cols + lanes - 1) // lanes + 1 | ||
| with pto.for_(0, valid_rows, step=1) as row: |
| def tgatherb(src, offsets, dst): | ||
| """``pto.tgatherb`` – tile gather using byte offsets (DPS).""" | ||
| _pto.tgatherb( | ||
| unwrap_surface_value(src), | ||
| unwrap_surface_value(offsets), | ||
| unwrap_surface_value(dst), | ||
| ) |
| build_plain/ | ||
| build_plan/ | ||
| install/ | ||
| /build-sim |
There was a problem hiding this comment.
非普适性的本地的ignore目录不要提上来
| src_ptr = src.as_ptr() | ||
| valid_rows, valid_cols = dst.valid_shape | ||
| lanes = pto.elements_per_vreg(dtype) | ||
| repeat_times = (valid_cols + lanes - 1) // lanes + 1 |
| mask, remained = pto.make_mask(dtype, col_loop.remained) | ||
| dst_reg = pto.vgatherb(src_ptr, offset_reg, mask) | ||
| pto.vsts(dst_reg, dst[row, col:], mask) | ||
| col_loop.update(remained=remained) |
Zhendong404
left a comment
Zhendong404
left a comment
There was a problem hiding this comment.
Review
总体方向没问题：为 pto.tgatherb 补上 ptodsl 高层封装（_ops.py / tile namespace）、a5 TileLib 模板和 ST 测试，另带一个 PTOViewToMemref.cpp 的属性保留修复。但新增的 TileLib 模板目前跑不起来，且测试完全没有覆盖到它，建议修好后再合入。
🔴 Blocker 1：模板中 col_loop 未定义，渲染时必然 NameError
ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py:39-43：
for row in range(0, valid_rows, step=1):
for col in range(0, valid_cols, step=lanes):
...
mask, remained = pto.make_mask(dtype, col_loop.remained) # col_loop 不存在
...
col_loop.update(remained=remained)证据：
- 对该函数做 symtable 分析，
col_loop是自由变量（global ref），函数内从未赋值，模块中也没有这个全局量； - 模板追踪机制（
ptodsl/ptodsl/_ast_rewrite.py:1129）重写for x in range(...)时使用内部新鲜名字（self._fresh("loop")），不会注入col_loop；全仓库col_loop只出现在模板文件里； - 其他用到
col_loop的模板都是显式定义的，如templates/a5/tcvt.py:96的col_loop = pto.for_(0, valid_cols, step=lanes).carry(...)；而使用裸range循环的模板（templates/a5/_elementwise.py:47-60）用的是普通局部变量remained。
按项目惯例建议改成：
for row in range(0, valid_rows, step=1):
remained = valid_cols
for col in range(0, valid_cols, step=lanes):
offset_reg = pto.vlds(offset[row, col:])
mask, remained = pto.make_mask(dtype, remained)
dst_reg = pto.vgatherb(src_ptr, offset_reg, mask)
pto.vsts(dst_reg, dst[row, col:], mask)🔴 Blocker 2：新模板没有任何测试覆盖
新增的 ST 测试走 explicit 模式 pto.tile.gatherb（直接 lowering 成 op），完全不经过 templates/a5/tgatherb.py 和 templates/__init__.py 里的注册项——这正是上面的 NameError 能混进来的原因。建议至少补一个类似 ptodsl/tests/test_tilelib_render.py（tadd 的那个）的渲染测试，specialize(src, offset, dst).mlir_text() 一次即可挡住这类问题。
🟡 应修：测试文件头注释与事实不符
test/tilelib-st/a5/tgatherb/case.py:13-21：
- 声称 "there is no high-level
pto.tile.gatherbwrapper" —— 本 PR 恰恰新增了这个 wrapper，且测试正文就在调用它； - 声称内核 "call
pto.vgatherbdirectly" —— 实际调的是pto.tile.gatherb； - 引用
lib/TileOps/tgatherb_template.py—— 该文件不存在。
🟡 应修：测试覆盖偏弱
- 9 个 case 全部
src_shape == dst_shape，且 offset 生成的是恒等映射（offset[i] = i * 32，期望输出恒等于 src）——索引/步长搞反了也测不出来。建议至少加一个乱序 offset 的 case，以及一个dst小于src的 case； _make_expected最后dst.reshape(src.shape)用的是 src 的 shape 而非 dst_shape，目前只因为两者总相等才没出错；_make_inputs里rng_seed/rng是死代码（rng创建后从未使用），且hash(name)跨进程不确定（PYTHONHASHSEED 随机化）——直接删掉即可；- offset 用
np.int32生成而内核声明pto.ui32，正数下字节相同无碍，建议统一为np.uint32。
🟢 小问题
ptodsl/ptodsl/_ops.py：tgatherb未加入__all__（同类tgather在 5922 行的列表里）；.gitignore：混入与本 PR 无关的个人工具目录（.opencode/、.arts/），且# OpenCode 配置中文注释与该文件全英文惯例不符，/build-sim写法也与其他目录条目（xxx/）不一致——建议拆出本 PR 或统一风格；lib/PTO/Transforms/PTOViewToMemref.cpp:3717的setAttrs(op->getAttrs())：与紧挨着的TLogOp处理模式一致，本身没问题；但TGatherBOp在PTOOps.td:4980并未声明任何 attribute，这行保留的只是上游 pass 挂上的元数据属性——加一句注释说明要保留什么属性会更易读。
做得好的地方
- C++ 修复与文件内既有模式（
TLogOp）一致； _ops.py/_tile_namespace.py的封装严格跟随tgather的现有写法；- ST 测试结构符合 tilelib-st 的
golden_output_case/auto_main约定，dtype 覆盖较全（含ui16_257x128这种非 2 的幂行数）。
结论：请先解决模板 NameError 和模板零覆盖这两个问题，其余为打磨项。修好后这个 PR 就可以了。
100ef65 to
5815798
Compare
3 participants
No description provided.