CANNBot Triton-Ascend Amin归约原子操作优化案例

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsname: triton-ascend-case-reduction-amin-atomic description: 原子操作归约amin优化非reduce轴很小时将reduce轴映射多核提供循环内/外两种原子操作方案减少存储vs减少竞争通过二次切分计算重组提升性能适用于MN如16×262144的极端shape比例场景 category: example version: 1.0.0 metadata: backend: ascend dsl: triton-ascend hardware: Atlas A2, Atlas A3Amin 归约原子操作优化案例任务特征数据尺寸(16, 262144)非reduce轴很小reduce轴很大策略将reduce轴映射到多核通过原子操作实现跨线程块归约优化 1切分策略调整# 简单方式非reduce轴映射多核 grid lambda meta: (triton.cdiv(M, meta[BLOCK_SIZE_M]),) # 错误优化方式reduce轴映射多核 grid lambda meta: (triton.cdiv(N, meta[BLOCK_SIZE_N]),) # Kernel内对列进行二次切分 for n_start in range(0, BLOCK_SIZE_N, SUB_BLOCK_SIZE_N): n_offsets pid * BLOCK_SIZE_N n_start tl.arange(0, SUB_BLOCK_SIZE_N)优化内容调整切分策略由非reduce轴映射多核调整为reduce轴映射多核为了不超过硬件缓存kernel内对列进行二次切分优化 2计算重组# 简单方式循环内多次归约 row_min float(inf) for n_start in range(0, BLOCK_SIZE_N, SUB_BLOCK_SIZE_N): 错误curr_min tl.min(data_block, 1) row_min tl.minimum(curr_min, row_min) # 正确优化方式维护矩阵结构 curr_min tl.full((BLOCK_SIZE_M, SUB_BLOCK_SIZE_N), float(inf), dtypetl.float32) for n_start in range(0, BLOCK_SIZE_N, SUB_BLOCK_SIZE_N): curr_min tl.minimum(data_block, curr_min) row_min tl.min(curr_min, 1)优化内容利用curr_min保持矩阵结构维护中间结果将多次归约合并为一次归约减少归约次数优化 3原子操作两种方案方案一循环内进行原子操作for m_start in range(0, M, BLOCK_SIZE_M): m_offsets m_start tl.arange(0, BLOCK_SIZE_M) mmask m_offsets M curr_min tl.full((BLOCK_SIZE_M, SUB_BLOCK_SIZE_N), float(inf), dtypetl.float32) for n_start in range(0, BLOCK_SIZE_N, SUB_BLOCK_SIZE_N): n_offsets pid * BLOCK_SIZE_N n_start tl.arange(0, SUB_BLOCK_SIZE_N) nmask n_offsets N mask (mmask[:, None]) (nmask[None, :]) block_ptrs in_ptr0 m_offsets[:,None] * in_stride0 n_offsets[None,:] * in_stride1 data_block tl.load(block_ptrs, maskmask, otherfloat(inf)) curr_min tl.minimum(data_block, curr_min) row_min tl.min(curr_min, 1) output_ptrs out_ptr0 m_offsets * out_stride0 tl.atomic_min(output_ptrs, row_min, maskmmask) # 每块立即原子操作特点减少了中间存储但增加了原子操作频率方案二循环外进行原子操作all_row_min tl.full((M,), float(inf), dtypetl.float32) # 预分配完整数组 for m_start in range(0, M, BLOCK_SIZE_M): m_offsets m_start tl.arange(0, BLOCK_SIZE_M) mmask m_offsets M curr_min tl.full((BLOCK_SIZE_M, SUB_BLOCK_SIZE_N), float(inf), dtypetl.float32) for n_start in range(0, BLOCK_SIZE_N, SUB_BLOCK_SIZE_N): n_offsets pid * BLOCK_SIZE_N n_start tl.arange(0, SUB_BLOCK_SIZE_N) nmask n_offsets N mask (mmask[:, None]) (nmask[None, :]) block_ptrs in_ptr0 m_offsets[:,None] * in_stride0 n_offsets[None,:] * in_stride1 data_block tl.load(block_ptrs, maskmask, otherfloat(inf)) curr_min tl.minimum(data_block, curr_min) row_min tl.min(curr_min, 1) curr_block_size_m tl.minimum(BLOCK_SIZE_M, M - m_start) all_row_min tl.insert_slice(all_row_min, row_min, [m_start], [curr_block_size_m], [1]) # 暂存中间结果 output_ptrs out_ptr0 tl.arange(0, M) * out_stride0 tl.atomic_min(output_ptrs, all_row_min) # 最后统一原子操作特点通过集中执行原子操作减少了竞争需要额外存储空间适合大规模数据优化 4配置# AI core40 # grid3240, UB用满 triton.Config({BLOCK_SIZE_M: 8, BLOCK_SIZE_N: 8192, SUB_BLOCK_SIZE_N: 1024}) # M切分较小UB用满 triton.Config({BLOCK_SIZE_M: 16, BLOCK_SIZE_N: 8192, SUB_BLOCK_SIZE_N: 512}) # M切分调大至16N的SUB切分调小至512防止超UB优化内容切分能被对应shape整除Grid数尽量大但不超过AI coreBLOCK_SIZE_N8192使grid32在UB用满的前提下进行kernel内切分大小调整总结非reduce轴很小、reduce轴很大时将reduce轴映射到多核并结合原子操作两种原子操作方案各有优劣方案一减少存储但原子操作频繁方案二集中原子操作但需额外空间确定核数后若超过硬件缓存可以考虑二次切分调整切分和核数配置尽量保证在不超出UB的前提下尽量用满UB【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考