triton_ascend入门学习二

昇腾与GPU的开发差异，分为三方面：多核任务并行、单核数据搬运、编译优化（AscendNPU IR优化）（略）。

NPU是物理核强绑定模式，与GPU逻辑维度并行+硬件自动物理映射的模式不同

GPU:可绑定多个维度轴（三维grid=[n,m,l] 等同于乘积n×m×l个并行线程），每个线程仅对应一次kernel执行，且仅执行一次。
NPU:Vector核，Cube核属于多个物理核，不同代际硬件核数不同，每个核仅执行一次Block,且支持对该Block重复调度执行。

写Triton内核函数时，合理的数据切分策略对性能优化至关重要。常见的切分参数包括：

ncore：使用的核数（跨核切分）
xblock：核间数据块大小（核间切分）
xblock_sub：核内切分粒度（核内细粒度划分）

以GELU算子开发示例，standard_unary 为标准Torch计算。

def standard_unary(x0):
    res = x0 * 0.5 * (1.0 + torch.erf(x0 / torch.sqrt(torch.tensor(2.0))))
    return res

triton_easy_kernel 为简单Triton实现。所有输入数据一次性被加载到内存中进行计算，只适合于小规模张量的计算，用于理解 Triton 内核的基本写法和调用方式。

# 定义triton_kernel核函数
@triton.jit
def triton_easy_kernel(in_ptr0, out_ptr0, NUMEL: tl.constexpr):
    idx_block = tl.arange(0, NUMEL)
    x = tl.load(in_ptr0 + idx_block)
    ret = x * 0.5 * (1.0 + tl.erf(x / tl.sqrt(2.0)))
    tl.store(out_ptr0 + idx_block, ret)

triton_better_kernel为更高效的Triton实现。

# 定义triton_kernel核函数
@triton.jit
def triton_better_kernel(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr):
    # 计算当前核处理数据块的起始偏移地址，实现核间切分。每个核仅负责 XBLOCK 大小的数据范围。
    xoffset = tl.program_id(0) * XBLOCK
    # 在单个核内部进一步细分数据块，每次处理 XBLOCK_SUB 大小的数据，实现核内切分。
    for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB):
        # 构造当前迭代的数据索引数组，用于访问输入和输出张量。
        x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:]
        # 设置掩码以防止越界访问，确保只处理合法范围内的数据。
        xmask = x_index < xnumel
        # load&store 分别用于从全局内存加载数据到片上内存，以及将计算结果写回全局内存。
        x = tl.load(in_ptr0 + x_index, xmask)
        ret = x * 0.5 * (1.0 + tl.erf(x / tl.sqrt(2.0)))
        tl.store(out_ptr0 + x_index, ret, xmask)

# 调用triton_kernel核函数
ncore = 32
xblock = 32768
xblock_sub = 8192
triton_better_kernel[ncore, 1, 1](x0, out1, x0.numel(), xblock, xblock_sub)

Triton算子迁移。移除 GPU 专属同步 API（如cuda中控制线程 / 流 / kernel 同步的专用接口）；grid 优先用 1D，2D NPU适配写法也会合并为1D, 实际grid值应对齐芯片物理核数，比如：(20,) 与 (4, 5) 的效果是一样的。完整迁移示例（向量加法）

import torch
+ import torch_npu  # 【新增】导入昇腾NPU PyTorch适配库，提供NPU设备支持
import triton
import triton.language as tl

- DEVICE = triton.runtime.driver.active.get_active_torch_device()  # 【删除】GPU设备自动获取，NPU无需此逻辑

@triton.jit
def add_kernel(x_ptr, # Pointer to first input vector.
    y_ptr, # Pointer to second input vector.
    output_ptr, # Pointer to output vector.
    n_elements, # Size of the vector.
    BLOCK_SIZE: tl.constexpr, # Number of elements each program should process.
):
    pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0.
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
-   assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE  # 【删除】GPU设备一致性校验，NPU无需显式断言
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output

torch.manual_seed(0)
size = 98432
- x = torch.rand(size, device='cuda')  # 【删除】GPU设备指定
+ x = torch.rand(size, device='npu')  # 【修改】指定为昇腾NPU设备
- y = torch.rand(size, device='cuda')  # 【删除】GPU设备指定
+ y = torch.rand(size, device='npu')  # 【修改】指定为昇腾NPU设备
output_torch = x + y
output_triton = add(x, y)
print(output_torch)
print(output_triton)
print(f'The maximum difference between torch and triton is '
f'{torch.max(torch.abs(output_torch - output_triton))}')

物理AI Core个数通过driver.active.utils.get_device_properties接口获取，建议将并发任务个数配置为AI Core个数。对于仅涉及Vector计算的Triton算子，并发任务个数应等于Vector Core的个数；其他类型的Triton算子（即Triton算子内使用了tl.dot），并发任务个数应等于AI Core的个数。
补充：TRITON_ALL_BLOCKS_PARALLEL ：启用或禁用自动根据物理核数优化逻辑核数，仅当逻辑核间可并行时方可启动。当逻辑核数大于物理核数时，启动该优化，则编译器自动调整逻辑核数量为物理核数，减少调度开销。

完成迁移基础步骤后，可能会遇到新的问题，可归纳为以下两类：

1.coreDim限制问题
当网格维度超过NPU硬件限制时触发。
典型错误信息：coreDim=xxxx can’t be greater than UINT16_MAX（65535）

2.UB空间溢出
内存使用超出NPU缓存容量。
典型错误信息：ub overflow, requires xxxx bits while 1572684 bits available!

解决思路1：
昇腾编译器有对应的解决方案，只需将环境变量’TRITON_ALL_BLOCKS_PARALLEL’设为1。设置命令如下：
export TRITON_ALL_BLOCKS_PARALLEL=1

解决思路2：
通过增大 BLOCK_SIZE 来减少所需的核心数量，确保 coreDim 不超过限制。但又可能引发新的UB溢出问题，因为单个线程块需要处理的数据量超出了NPU的UB缓存容量。解决思路是引入 BLOCK_SIZE_SUB 参数，将大块进一步细分，在保持合理 coreDim 的同时控制内存使用。

为什么会出现UBSIZE超出内存的错误，因为切分不合理,存在过多的非对齐访存或者运算。但举的例子不明不白，就不说了。