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

[platelett]

Summary

一些 Tile API（例如 Reduce、Sort）执行时需要额外的 UB 临时空间。默认情况下，这些临时 buffer 由编译器自动插入。

这对普通 kernel 很方便，但会影响手动内存规划。用户使用 T.annotate_address 安排 UB 地址时，编译器后来插入的 buffer 不在原有规划中，可能与用户指定的区域冲突。这正是 #1164 希望解决的问题。

此前，如果用户想手动控制这些临时空间，只能对整个 kernel 关闭自动注入，并依赖 tmp_ub、tmp_ub_reduce_out 等编译器内部名称。这样无法表达“这个调用自动分配，另一个调用使用我提供的 arena”。

本 PR 改为逐调用控制：

# 这个调用继续由编译器管理
T.reduce_sum(src0, dst0)

# 这个调用复用用户规划的 UB
tmp_storage = T.alloc_ub((4096,), "float32")
T.annotate_address({tmp_storage: 32768})
T.reduce_sum(src1, dst1, tmp=tmp_storage[8:])

上例中的显式 arena 是 float32，region 从第 8 个元素开始，也就是 byte 32。Reduce 后端实际需要的是 uint8 workspace，因此 lowering 会在同一个 byte 32 地址上建立 uint8 view。

这里不会进行数据转换或额外拷贝，只是改变后端解释这段存储的 dtype。

是否使用显式临时空间，只取决于当前调用有没有传入 tmp=；不需要新的 kernel 级开关，也不需要用户了解任何 tmp_ub* 内部名称。

User-facing contract

tmp= 是 keyword-only 参数，接受：

• 一维、静态、连续的 shared.ub Buffer；
• 或上述 Buffer 中的一段静态、连续 BufferRegion。

Buffer 和 BufferRegion 可以使用任意定宽标量 dtype，起始 byte address 必须 32B 对齐。

用户始终只提供一个 arena。AscendC 和 PTO 需要的 workspace 数量、dtype 和对齐方式可能不同，但这些后端差异由 lowering 处理，不暴露到前端。

调用方式	行为
未传 tmp	编译器继续自动分配 workspace
传入非空 tmp	只有当前调用使用该 arena，不再为它分配隐藏 workspace
当前 target 不需要 workspace	删除 tmp operand，不生成零尺寸 allocation
当前 target 需要 workspace，但显式 arena 为空	在 lowering 阶段报告错误

同一 kernel 可以混用自动和显式调用。

显式非空 arena 的容量由调用者负责。本 PR 会检查 rank、scope、静态连续性、region 边界和起始地址对齐，但不会把内部启发式容量作为公开的最小值检查。

Supported public APIs

• reduce_sum/max/min
• 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

How one arena can provide multiple internal workspaces

前端只接收一个 tmp=arena。如果后端内部需要多段空间，compiler pass 会从这个 arena 中自动建立互不重叠的 view。

以 PTO clear=False Reduce 为例：

T.reduce_sum(src, dst, clear=False, tmp=arena)

clear=False 表示需要保留 dst 中的已有值。PTO 会先把本次 Reduce 的结果写入临时输出，再与 dst 合并：

tmp_result = reduce(src, main_workspace)
dst = dst + tmp_result

因此，一个 frontend arena 在 PTO 内部会被解释为：

arena
├── Reduce 算法使用的主 workspace
├── 32B 对齐 padding
└── 本次 Reduce 的临时输出

用户仍然只传一个 tmp=arena，不需要分别声明两个 buffer。

对于 first-axis Reduce，PTO 不需要主 workspace，因此只建立临时输出 view。

Compiler changes

本 PR 继续使用现有的 InjectTmpBuffer。它现在统一决定：

1. 当前 backend 是否需要 workspace；
2. 未传 tmp 时需要分配多少字节；
3. workspace 应当使用什么 dtype；
4. 一个 frontend arena 是否需要切成多个内部 view；
5. 当前 backend 不使用 workspace 时是否应删除 operand。

同一 kernel 中所有隐式调用仍共享一个最大尺寸的 uint8 allocation，每个调用在这段存储上建立自己需要的 typed view。显式调用不参与这个最大值计算，因此不会放大隐藏 allocation。

显式 workspace 也必须参与流水线依赖分析。例如，在一次调用完成前，后续 DMA 不能覆盖同一段 arena。

以前，相关分析会根据 API 名称和固定参数下标判断“哪个参数是 workspace”。tmp= 变为可选参数后，不同调用可能没有 workspace、只有一个 workspace，或者在 lowering 后产生多个 view，固定下标不再可靠。

现在 InjectTmpBuffer 会直接标记 lowering 后实际存在的 workspace view；后续的流水线规划和同步插入读取这些标记，而不再重新猜测参数位置。

Backend-specific behavior

PTO

此前 PTO 的 workspace 都由编译器创建，是从 offset 0 开始的完整 Buffer，因此部分 codegen 路径只使用 backing buffer 的起始地址，没有完整处理 BufferRegion。

加入 tmp=arena[offset:end] 后，继续使用旧路径会丢失 slice offset。例如，用户传入：

arena = T.alloc_ub((4096,), "float32")
T.tile.select(dst, mask, src0, src1, tmp=arena[8:])

这里 workspace 应从 byte 32 开始。如果 codegen 退回 arena 的起点，Select 会从 byte 0 写入，可能覆盖 arena 中属于其他调用的数据。

此外，不同 PTO API 对 workspace dtype 的要求不同。假如某个 view 从 float32 元素 8 开始，而后端需要 uint8 workspace，不能把元素 offset 8 原样套到 uint8 上，否则地址会从 byte 32 错误地变成 byte 8。

本 PR 因此统一按 byte address 处理中间几何：

frontend region offset
→ 换算为 byte offset
→ 保留同一个 byte address
→ 再换算为 backend workspace dtype 的 element offset

当前公开 PTO workspace 路径都使用这一逻辑：

操作	PTO workspace view
Sort、TopK、MergeSort	输入数据的 dtype
Gather、custom GatherMask	index 的 dtype
Select、Reduce	uint8

PTO clear=False Reduce 还会在保留 arena 起始 offset 的基础上，将一个 frontend arena 切分为互不重叠的主 workspace 和临时输出。

AscendC（dav-2201）

将 tmp= 变成公开接口后，compiler policy 必须与 CANN 函数实际是否接收 workspace 保持一致。逐项核对 dav-2201 实现时发现了三类旧的不一致：

1. 实现不需要 workspace，compiler 却仍然分配或传入。 Clamp、float32 Round、MergeSort、Select、Gather 和 GatherMask 属于这一类。本 PR 会删除这些调用的 workspace operand；half Clamp 直接生成对应 dtype 的 Mins/Maxs，float32 Round 直接生成 CAST_RINT。
2. 实现需要 workspace，compiler 却没有保留。 tensor-tensor int32 Pow 属于这一类，现在会建立并传入 uint8 workspace view。
3. workspace 存在，但 dtype 与 CANN 接口不一致。 experimental ReduceSum 以前错误地建立 uint8 view；其接口实际要求与输入相同的 dtype，本 PR 已按源 dtype 建立 view。

这些规则同时用于显式 arena 和隐式 allocation，避免两条路径再次产生不同的函数签名或 workspace 类型。

Scope

内容	本 PR
逐调用 tmp=	支持
自动和显式调用混用	支持
前端始终只传一个 arena	支持
静态 BufferRegion	支持
zero-workspace 路径接受空 arena	支持
get_tmp_buffer_size()	暂不提供
动态 offset/extent BufferRegion	暂不支持
对非空显式 arena 做精确容量验证	暂不支持；容量由调用者负责
AscendC A5 workspace policy	不在本 PR 范围内

其他边界：

• 没有新增 pass、wrapper 或配置，也没有改变 LowerAndLegalize 的 pass 顺序。
• AscendC workspace sizing 只按仓库当前固定的 --npu-arch=dav-2201 实现。
• AscendC 已保留 BufferRegion 的起始地址和 workspace dtype，但尚未统一使用 LocalTensor::SetSize 将 region extent 变成严格访问上界；这部分留给后续 PR。
• PTO 仍不支持 bilinear interpolation、sin/cos 和两个 experimental ReduceSum API；这些调用会在 codegen 前得到明确诊断。

Validation

仓库中只保留覆盖公开契约和主要回归风险的最小测试；更大的 API、dtype、workspace 组合矩阵仅在本地运行，避免增加后续 CI 的长期负担。

验证层次	主要覆盖内容	结果
Frontend 与 IR	所有公开 API 的 keyword-only tmp=；Buffer/BufferRegion 检查；旧 Reduce positional 调用兼容	68 passed，兼容测试另有 7 passed
Workspace lowering	显式与隐式调用混用、zero-workspace、typed view、非零 byte offset、PTO Reduce arena 切分、pass 重入	包含在上述专项测试中
Codegen 与依赖分析	AscendC LocalTensor 和 PTO tile 的实际地址、dtype、extent，以及显式 arena 的流水线读写依赖	包含在上述专项测试中
dav-2201 device smoke	代表性的显式、隐式和 workspace dtype reinterpret 路径	3 passed
本地生成式矩阵	跨 API、target、Buffer/BufferRegion、显式/隐式和 zero-workspace 的组合	154/154 passed
真实后端数值测试	Reduce、Broadcast、Sort/TopK、Sin/Cos/Pow/Xor、Clamp/Round/Sigmoid，以及可执行的 Gather/Select 路径	138 passed, 2 xfailed
工程门禁	增量构建、Python/C++ 格式和 diff 检查	全部通过

Commands

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

Review order

1. 4d852b36：公开 tmp= 接口、workspace lowering、PTO/AscendC codegen 和依赖分析。
2. ff165f7c：最小永久测试集和文档。

[gemini-code-assist]

Caution

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

[github-actions]

👋 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!

🚀

[platelett]

/re-test

[github-actions]

🔄 Re-running failed jobs

Original workflow run: View details

Only the failed jobs will be re-executed.

[gemini-code-assist]

Caution

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

[LLMZhangYC]

lgtm