Add merge-axis TileLang element-wise templates#456
Open
liuzidi wants to merge 1 commit into
Open
Conversation
Zhendong404
requested changes
May 29, 2026
| total_elems = valid_rows * valid_cols | ||
| lanes = pto.get_lanes(dtype) | ||
|
|
||
| with pto.strict_vecscope(dst, src0, src1, total_elems, 0, total_elems, lanes) as ( |
Collaborator
There was a problem hiding this comment.
PTOAS后端会自动推导vecscope，这里不需要显式写
| priority=100, | ||
| advanced=True, | ||
| ) | ||
| def template_tabs_merge_axis(src: pto.Tile, dst: pto.Tile): |
Collaborator
There was a problem hiding this comment.
binary op能否共用1个模板，看一下手册里有template的写法
| op="pto.tabs", | ||
| constraints=[full_axis_constraint], | ||
| priority=100, | ||
| advanced=True, |
Collaborator
There was a problem hiding this comment.
跟ptr无关的op不需要advanced模式
Zhendong404
requested changes
May 29, 2026
|
|
||
|
|
||
| // Merge-axis case: f32 full-axis 16x72 (1152 elements) | ||
| func.func @TABS_f32_merge_axis_16x72(%a_ptr: !pto.ptr<f32>, %b_ptr: !pto.ptr<f32>) attributes {pto.aicore} { |
Collaborator
There was a problem hiding this comment.
测试用例现在推荐用PTODSL写，放到test/dsl-st目录下面
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
标题
TileLang element-wise OP 增加满轴合轴版本并补齐 ST 覆盖
描述
本次主要为 lib/TileOps 中的 element-wise 类 OP 增加“合轴/满轴”版本实现，并同步补齐对应的 TileLang ST 用例。
修改内容
新增 lib/TileOps/merge_axis.py，抽出满轴场景的公共约束与合轴遍历辅助逻辑。
为 element-wise OP 增加 merge_axis 版本模板，满轴时优先选择 1D 合轴实现，非满轴仍保持原有 2D 实现兜底。
更新 test/tilelang_st 下相关 testcase，补充合轴场景 .pto、cases.py、launch.cpp、main.cpp。
调整 run_ptoas_to_file.cmake，确保 ST 编译时使用当前仓库内的 TileOps / tilelang-dsl 实现，避免误用外部安装版本。
修正部分合轴 case 的 shape / valid_shape / stride 配置，使其满足满轴展开条件。
验证
已通过 test/tilelang_st 相关 element-wise 合轴 ST 回归。