Skip to content

[Feat] Support explicit workspace arenas for Ascend tile APIs 🤖 - #1496

Draft
platelett wants to merge 2 commits into
tile-ai:ascendc_ptofrom
platelett:feat/restore-explicit-tmp
Draft

[Feat] Support explicit workspace arenas for Ascend tile APIs 🤖#1496
platelett wants to merge 2 commits into
tile-ai:ascendc_ptofrom
platelett:feat/restore-explicit-tmp

Conversation

@platelett

@platelett platelett commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

如何审阅

  • Commit 1 — 生产代码:公开 tmp= API、统一 workspace policy,以及 AscendC/PTO lowering、metadata 与 codegen。
  • Commit 2 — 测试和文档:最小永久回归集、dav-2201 sizing 说明及使用契约。
  • 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 调用提供完整的显式 arena;
  • PTO 的 workspace slice 可能丢失非零 byte offset,或按错误 dtype 解释 extent;
  • target 不使用 workspace 时仍可能保留 operand 或产生零尺寸 allocation;
  • dav-2201 的部分隐式 sizing、Clamp、Round 与 experimental ReduceSum 路径不符合实际后端契约。

改动

公开 API 契约

以下接口统一增加 keyword-only tmp: Buffer | BufferRegion | None = None

  • T.reduce_sum/max/min
  • T.tile.broadcast/sort/merge_sort/topk
  • gather_mask/select/gather
  • sigmoid/sin/cos/pow/bitwise_xor
  • clamp/clamp_max/clamp_min/round
  • deprecated bilinear_interpolation
  • reduce_sum_experiment/reduce_sum_mask_experiment

显式 arena 必须是一维、静态、连续、任意定宽标量 dtype 的 shared.ub Buffer,或同类 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-call WorkspaceSpec 统一描述是否需要 workspace、typed view dtype、隐式 bytes 与 access mask:

  • 所有隐式调用复用一个最大尺寸的 uint8 主 arena,再为每次调用建立 typed view;
  • 显式调用不参与隐藏 arena 的最大值统计;
  • target 不使用 workspace 时统一删除 operand,不创建零尺寸 allocation;
  • pass 重入保持 no-op;
  • 隐式 buffer 使用冲突安全的 fresh name;
  • PTO 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

  • 更新 reduce、broadcast、sort/topk、bilinear、sin/cos、tensor-tensor pow、xor、round half、sigmoid 和 experimental ReduceSum 的隐式 sizing 与 view dtype。
  • Clamp 系列直接使用带正确模板 dtype 的 basic Mins/Maxs,不消费 workspace;float Round 直接生成 CAST_RINT
  • 恢复 tensor-tensor int32 Pow 的 uint8 workspace policy。
  • experimental ReduceSum workspace 修正为源 dtype。
  • merge-sort、select、gather、gather-mask 删除未使用的可选 operand。
  • Broadcast 区分 b8/b16/b32,以及 equal/scalar、axis 0、axis 1 aligned/unaligned staging。

这些数值只服务 compiler-managed allocation,不作为显式 arena 的公开最小容量。

PTO 与地址几何

  • XOR 使用源 dtype view;Sort/TopK/MergeSort 使用数据 dtype;Gather/custom GatherMask 使用 indices dtype;Select 与 Reduce 使用 uint8
  • BufferRegion 的 byte offset 与换算后的 typed extent 同时保留。
  • Select、Reduce、Sort/TopK、Gather/GatherMask 不再退回只打印 backing base 的路径。
  • tail-mask reduce/broadcast 同时支持带或不带 workspace operand 的调用布局。

范围说明

  • 本 PR 增量修改现有 InjectTmpBuffer;未新增 pass、wrapper 或配置,也未改变 LowerAndLegalize 中 pass 的数量、位置或执行顺序。
  • 当前 AscendC policy 只按现有固定 --npu-arch=dav-2201 实现,不在本 PR 全局拒绝 A5。
  • AscendC 已保留 BufferRegion 起始 byte address 与 workspace dtype,但尚未统一调用 LocalTensor::SetSize 约束 region extent,因此本 PR 不宣称 AscendC region extent 是严格上界。
  • 不增加公开 size-query API。
  • zsh 兼容的本地 set_env.sh、flag、row-expand 以及后续 bounded-LocalTensor 计划均不在提交范围内。

验证

永久测试保持为必要的契约、policy、metadata、地址几何和三个设备 smoke;更大的组合矩阵只在本地执行:

  • 专项 frontend/IR/codegen:68 passed
  • 既有 Reduce positional/keyword 兼容:7 passed
  • dav-2201 NPU smoke(显式 PTO reduce、显式 AscendC broadcast、隐式 sigmoid/round):3 passed
  • 相关真实后端/数值矩阵:138 passed, 2 xfailed
  • 未提交的生成式 workspace 矩阵:154/154 passed
  • int32 Pow 64 元素、768B 隐式 arena:设备数值通过
  • cmake --build build --target tilelang tilelang_module -j4:通过
  • bash install_ascend.sh --enable-incremental:通过
  • Ruff 0.15.21、clang-format 18.1.8、git diff --check:通过

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

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
platelett force-pushed the feat/restore-explicit-tmp branch from fb8a3be to ff165f7 Compare July 29, 2026 19:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant