流压缩(从给定序列中选择满足条件的元素并保序输出)是 MoE 分发、top-k 筛选、稀疏矩阵处理等场景中的常见计算。该算子的特点在于每个线程无法提前知道自己的元素在输出中的位置,输出位置必须由线程协作计算得出。
本需求开发一个 SIMD 与 SIMT 混合编程最佳实践样例 simd_simt_stream_compaction,以稳定流压缩算子为载体,呈现确定性(稳定)压缩的调优路径。
SIMT 样例开发需求
算子语义:给定输入序列 input[](int32)与判断条件,输出命中元素并保持原相对顺序。出参:output[](int32,GM,长度与 input 一致,多余位置补 INT32_MIN)、selected_count(uint32,GM,命中总数)。示例:挑选偶数,input [4,3,2,7,9,8,2] → output [4,2,8,2,INT32_MIN,INT32_MIN,INT32_MIN],selected_count=4。
input[]
output[]
selected_count
[4,3,2,7,9,8,2]
[4,2,8,2,INT32_MIN,INT32_MIN,INT32_MIN]
5 个场景的调优路径:
thread_selected_flags
warp_selected_counts
lane_offset
warp_selected_offsets
asc_ballot
__popc
lanemask_lt()
你好,感谢你的建议,我们将安排需求增加此样例。
您好,相关代码已合入,我们先关闭此issue,后续如果有其他问题,欢迎再提issue交流~
Backgroud(背景信息)
流压缩(从给定序列中选择满足条件的元素并保序输出)是 MoE 分发、top-k 筛选、稀疏矩阵处理等场景中的常见计算。该算子的特点在于每个线程无法提前知道自己的元素在输出中的位置,输出位置必须由线程协作计算得出。
本需求开发一个 SIMD 与 SIMT 混合编程最佳实践样例 simd_simt_stream_compaction,以稳定流压缩算子为载体,呈现确定性(稳定)压缩的调优路径。
Origin(信息来源)
SIMT 样例开发需求
Benefit / Necessity (价值/作用)
Design(设计方案)
算子语义:给定输入序列
input[](int32)与判断条件,输出命中元素并保持原相对顺序。出参:output[](int32,GM,长度与 input 一致,多余位置补 INT32_MIN)、selected_count(uint32,GM,命中总数)。示例:挑选偶数,input[4,3,2,7,9,8,2]→ output[4,2,8,2,INT32_MIN,INT32_MIN,INT32_MIN],selected_count=4。5 个场景的调优路径:
thread_selected_flags(1024),再遍历求前缀;以 UB 共享替代 GM 重复读取warp_selected_counts与 warp 内顺序lane_offset,再算warp_selected_offsets;UB 读取从约 5.36 亿次降至 33264 次asc_ballot+__popc+lanemask_lt()在寄存器层交换数据,不经 UB