[Feat] Support explicit workspace arenas for Ascend tile APIs 🤖 - #1496
Draft
platelett wants to merge 2 commits into
Draft
[Feat] Support explicit workspace arenas for Ascend tile APIs 🤖#1496platelett wants to merge 2 commits into
platelett wants to merge 2 commits into
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
- Expose keyword-only tmp arenas across public workspace APIs - Centralize target-specific sizing, typed views, and zero-workspace policies - Preserve explicit BufferRegion byte geometry in AscendC and PTO lowering
- Keep focused frontend, IR, codegen, and dav-2201 runtime regressions - Document explicit arena geometry and target-specific workspace policy - Record the current AscendC LocalTensor region-boundary limitation
platelett
force-pushed
the
feat/restore-explicit-tmp
branch
from
July 29, 2026 19:43
fb8a3be to
ff165f7
Compare
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
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.
如何审阅
tmp=API、统一 workspace policy,以及 AscendC/PTO lowering、metadata 与 codegen。ascendc_pto基线经 Ruff 0.15.21 与 clang-format 18 检查均不产生 reformat diff,因此没有 formatter-only commit。问题
公开 API 过去无法一致地传入显式 UB workspace;隐式容量、backend view dtype、零需求路径和 BufferRegion 地址处理又分散在 collector、allocation 与 codegen 中。这会导致:
改动
公开 API 契约
以下接口统一增加 keyword-only
tmp: Buffer | BufferRegion | None = None:T.reduce_sum/max/minT.tile.broadcast/sort/merge_sort/topkgather_mask/select/gathersigmoid/sin/cos/pow/bitwise_xorclamp/clamp_max/clamp_min/roundbilinear_interpolationreduce_sum_experiment/reduce_sum_mask_experiment显式 arena 必须是一维、静态、连续、任意定宽标量 dtype 的
shared.ubBuffer,或同类 BufferRegion;起始 byte address 必须 32B 对齐。arena dtype 只描述 byte address 和容量,lowering 会建立 target 所需的 typed view。非零显式 arena 不与隐式 sizing heuristic 比较,容量由调用者负责。真实零需求路径会删除 operand,因此允许传入零 extent arena。Reduce 保留旧 positional
clear/real_shape行为,tmp仅允许关键字传入。统一 lowering policy
InjectTmpBuffer以单一 per-callWorkspaceSpec统一描述是否需要 workspace、typed view dtype、隐式 bytes 与 access mask:uint8主 arena,再为每次调用建立 typed view;clear=False行归约将一个显式 arena 切成不重叠的 main/output views,列归约只建立 output view;隐式路径继续使用两个 allocation。固定 tmp slot 和 target support 只有一张配置表。PipelinePlanning 与两个 SyncInsert consumer 根据 lowering 后真实的
tvm_access_ptr和 access mask 恢复 buffer metadata,从而不再为 optional operand 维护第二套易漂移的 positional ABI。dav-2201 AscendC
Mins/Maxs,不消费 workspace;float Round 直接生成CAST_RINT。uint8workspace policy。这些数值只服务 compiler-managed allocation,不作为显式 arena 的公开最小容量。
PTO 与地址几何
uint8。范围说明
InjectTmpBuffer;未新增 pass、wrapper 或配置,也未改变LowerAndLegalize中 pass 的数量、位置或执行顺序。--npu-arch=dav-2201实现,不在本 PR 全局拒绝 A5。LocalTensor::SetSize约束 region extent,因此本 PR 不宣称 AscendC region extent 是严格上界。set_env.sh、flag、row-expand 以及后续 bounded-LocalTensor 计划均不在提交范围内。验证
永久测试保持为必要的契约、policy、metadata、地址几何和三个设备 smoke;更大的组合矩阵只在本地执行:
68 passed7 passed3 passed138 passed, 2 xfailed154/154 passedcmake --build build --target tilelang tilelang_module -j4:通过bash install_ascend.sh --enable-incremental:通过git diff --check:通过