[Feature] PTO DSL 缺失的 SIMT 标量运算-影响 RNG 算子的完整实现

[bingmeiyou]

Summary

背景

在 TileLang 的 PTO 后端 codegen 中，我们需要将 T.rng_init / T.rng_rand / T.rng_rand_float 三个算子 lowering 到 PTO DSL。这三个算子实现的是 Philox 4x32-10 随机数生成算法，要求在 @pto.simt 函数内以标量运算逐线程执行。

AscendC 后端已完整实现该功能（通过 C++ 头文件直接调用标量数学函数），但 PTO DSL 缺少若干关键 API，导致我们只能用近似方法实现，与 AscendC 的实现结果存在差异。

问题一：缺少标量三角函数（sin / cos）

现象

T.rng_rand_float(dist="normal") 需要将均匀分布随机数转换为标准正态分布。标准做法是 Box-Muller 变换：

z0 = sqrt(-2 * ln(u1)) * cos(2π * u2)
z1 = sqrt(-2 * ln(u1)) * sin(2π * u2)

PTO 现状

• pto 命名空间下没有 sin / cos 函数
• scalar 模块也没有 sin / cos
• 全仓库搜索（包括测试文件）未找到任何三角函数的使用

我们的临时方案

用 Irwin-Hall 近似代替：生成 12 个 [0,1) 均匀随机数求和后减 6，近似服从 N(0,1)。

与 AscendC 的差异

方面	AscendC（Box-Muller）	PTO（Irwin-Hall 近似）
精度	精确	尾部有偏差（±3σ 以外误差较大）
每次输出	2 个正态随机数	1 个
计算量	2 次 uniform + 三角/对数/开方运算	12 次 uniform + 11 次加法
可复现性	与 CUDA curand 的正态分布序列不同	与 AscendC 的正态分布序列不同

期望

PTO DSL 在 @pto.simt 上下文中提供标量 sin / cos 运算（或提供 Box-Muller 的内置封装），使 PTO 后端能生成与 AscendC 相同的正态分布随机数序列。

问题二：缺少 SIMT 级别的随机数生成原语

现象

T.rng_init / T.rng_rand / T.rng_rand_float 的语义是 per-thread 的（每个 SIMT 线程独立维护 PRNG 状态，逐次调用返回随机数），对应 CUDA curand 的 device API。

PTO 现状

PTO DSL 有 pto.trandom 操作，但它是 vector engine 的 tile 级操作：

• 一次调用产出一整块 ui32 随机数（如 4×256）
• 无状态管理（每次调用独立计算，不支持连续抽取）
• 只输出 uint32，不提供 uniform/normal 转换
• 运行在 vector engine（lane 并行），不是 SIMT（thread 并行）

这与我们的需求不匹配：我们需要每个线程维护独立的 Philox 状态，支持多次调用 rand() 并推进内部计数器。

我们的临时方案

用 pto.mulhi + 标量运算（*、^、+）手动实现完整的 Philox 4x32-10 算法，封装为 Python 类 PhiloxRNG（放在 tilelang/contrib/ptodsl/rng.py）。

期望

PTO DSL 提供 SIMT 级别的随机数生成原语，或提供状态管理机制使 pto.trandom 能支持 per-thread 的连续抽取语义。

总结

#	问题	严重程度	当前临时方案
1	缺少标量 sin/cos	中	Irwin-Hall 近似（精度有损）
2	缺少 SIMT 级 RNG 原语	中	手动实现 Philox 算法

这两个问题导致 PTO 后端的 RNG 实现与 AscendC 不完全一致（正态分布精度有损、随机数序列不同）。

Motivation / use case

--

Proposed API / behavior

No response

Alternatives considered

No response

Additional context

No response

[Zhendong404]

sin/cos的支持方案需要端到端完整设计 @mouliangyu

[WenboCodes]

Ascend C的接口实现是通过库吗？如果是的话，建议需求方采用相同的方式支持。