SIMD与SIMT混合编程HashTable MTE队列优化样例

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

概述

本样例比较HashTableinsert_or_assign的两种value写入路径：SIMT Warp直接写value，以及单kernel内使用MTE任务队列写value。样例实现一个固定容量的C++风格HashTable，仅用于展示本次优化路径，保留本次对比需要的语义和数据结构。

优化路径：

两种Case的数据处理流水如下图。

本样例支持的产品及CANN软件版本

目录结构介绍

├── simd_simt_hash_table_mte_queue │ ├── CMakeLists.txt // 编译工程文件 │ ├── figures // README中的图片资源 │ ├── hash_table_mte_queue.asc // Ascend C样例实现 │ ├── README.md // 样例说明文档 │ └── scripts │ └── run.sh // 编译并使用msOpProf采集两种Case性能数据

样例描述

样例功能：

样例实现一个固定容量HashTable，支持与C++17std::unordered_map::insert_or_assign类似的插入或更新语义：key已存在时更新value，key不存在时插入key/value。Hash冲突通过线性探测处理。

HashTable由两段GM存储组成：table[bucket_count]保存bucket元数据，其中每个bucket只保存key和状态；table_values[bucket_count, dim]保存value数据，其中第bucket_idx行保存table[bucket_idx]对应key的value向量。key和value不在同一结构体内连续存放，而是通过相同的bucket_idx关联；每条value向量长度为dim，dim为运行时参数。

bucket的key字段同时承载占用信息：

插入或更新过程中，kernel先将目标bucket的key从EMPTY_KEY或真实key CAS为LOCKED_KEY，再写入value，最后发布真实key。find只匹配真实key，因此不会观察到“key已经可见但value尚未写完”的中间结果。

HashTable插入流程如下：

$$ bucket = hash(key) & (capacity - 1) $$

$$ bucket = (bucket + 1) & (capacity - 1), \quad \text{when collision occurs} $$

样例规格：

样例参数如下：

样例实现

性能对比以msOpProf输出的Task Duration(μs)为准。

Case 0：SIMT Warp直接写value

实现方式：参考insert_or_assign_warp_store_vf()函数实现。

该实现中，每个Warp处理一条key。lane 0执行hash、线性探测和CAS锁定bucket，Warp内线程通过asc_shfl获取目标bucket index，并按WARP_SIZE分片写入value向量。value搬运由SIMT线程直接访问GM完成。

写入顺序为：CAS锁定目标bucket，Warp内线程写value，执行asc_threadfence()，最后由lane 0发布真实key。

关键代码：

__simt_vf__ inline void insert_or_assign_warp_store_vf(__gm__ Bucket* table, __gm__ float* table_values, __gm__ const int64_t* keys, __gm__ const float* values, uint32_t key_num, uint32_t capacity, uint32_t dim, uint32_t block_index, uint32_t num_blocks) { const uint32_t lane_id = static_cast<uint32_t>(threadIdx.x) % WARP_SIZE; ... if (lane_id == 0) { int64_t old_key = asc_atomic_cas(key_addr, EMPTY_KEY, LOCKED_KEY); if (old_key == LOCKED_KEY) { old_key = wait_until_unlocked(key_addr); } ... } ... for (uint32_t j = lane_id; j < dim; j += WARP_SIZE) { table_values[dst_base + j] = values[src_base + j]; } asc_threadfence(); if (lane_id == 0) { (void)asc_atomic_exch(&table[bucket_idx].key, key); } }

性能数据：

性能数据分析：

• 随着dim增大，SIMT线程逐元素写value的耗时明显上升。dim=256时，value搬运已经成为该路径的主要开销。

原理说明：

HashTable插入包含hash、线性探测、CAS和冲突处理等离散控制流，适合用SIMT实现。value向量是连续内存，SIMT线程逐元素访问GM时，通常不如MTE连续搬运高效。

优化方向：

Case 0中，控制流和值写入均由同一组SIMT线程完成。Hash/probe/CAS属于离散控制流，适合SIMT线程处理；但value向量按GM连续地址存放，随着dim增大，Warp内线程需要执行更多逐元素GM写入，SIMT线程会被value搬运占用，后续key的探测也要等待当前value写完后才能继续推进。

Case 1保留SIMT处理hash/probe/CAS的职责，只把连续value搬运拆给MTE执行。SIMT线程锁定bucket后将value源地址、目的地址和待发布key写入UB任务队列，然后继续生产后续任务；MTE按任务队列批量执行GM->UB->GM搬运。这样，后续key的SIMT探测可以与前序value搬运在单kernel内并行推进，从而缓解Case 0中SIMT线程被连续value写入占用的问题。

Case 1：MTE task queue搬运value

实现方式：参考mte_task_queue和insert_or_assign_mte_task_queue_kernel()函数实现。

该实现仍由SIMT线程执行hash/probe/CAS，但不再把完整任务列表写入GM。每个Vector Core维护一个UB task-id ring queue，SIMT producer把成功锁定的任务写入ring，任务中记录输入value源地址、HashTable value目的地址、待发布key和key地址；MTE consumer按ring中的任务批量搬运value，value写完后由SIMT drain路径发布真实key。global_id、simt_assign_id、mte_finish_id、key_assign_id和key_finish_id分别跟踪任务分配、MTE可见范围、value搬运完成进度、key发布认领范围和最终发布进度。bucket key在value写完前保持LOCKED_KEY，真实key只在完整value向量写回GM后发布。

下图给出Case 1中单个任务的可见性时序。global_id、simt_assign_id、key_assign_id和key_finish_id由SIMT路径推进，mte_finish_id由MTE路径推进。

关键代码：

class mte_task_queue { public: __simt_callee__ inline void assign_warp_task( uint64_t src_addr, uint64_t dst_addr, int64_t key, uint64_t key_addr, bool need_assign) { uint32_t write_mask = asc_ballot(need_assign ? 1 : 0); uint32_t warp_write_count = __popc(write_mask); uint32_t task_start_id = (laneid() == 0) ? apply_id(warp_write_count) : 0; task_start_id = asc_shfl(task_start_id, 0); drain(task_start_id + warp_write_count); if (need_assign) { uint32_t lane_offset = __popc(write_mask & lanemask_lt()); fill_task_info(task_start_id + lane_offset, src_addr, dst_addr, key, key_addr); } asc_threadfence(); if (laneid() == 0) { assign_task(task_start_id, task_start_id + warp_write_count); } } __aicore__ inline void run_task() const { while (true) { uint32_t simt_assign_id = *reinterpret_cast<__ubuf__ volatile uint32_t*>(simt_assign_id_); uint32_t proc_count = min(max_proc_batch, simt_assign_id - mte_finish_id); ... asc_copy_gm2ub_align(value_local_, src_addr, 1, value_bytes, 0, 0, true, values_cache_mode_, 0, 0); asc_sync_notify(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); asc_sync_wait(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); asc_copy_ub2gm_align(dst_addr, value_local_, 1, value_bytes, table_values_cache_mode, 0, 0); asc_sync_notify(PIPE_MTE3, PIPE_S, EVENT_ID0); asc_sync_wait(PIPE_MTE3, PIPE_S, EVENT_ID0); *reinterpret_cast<__ubuf__ volatile uint32_t*>(mte_finish_id_) = mte_finish_id + proc_count; } } };

性能数据：

性能数据分析：

• Case 1把长度为dim的value向量搬运从SIMT线程中拆出，在单kernel内由SIMT producer和MTE consumer并行推进。
• 相比Case 0，Case 1在五组dim上均取得加速；dim=128时加速比最高，约2.79x。

性能对比总结

Ascend 950PR性能对比

各dim配置下的性能数据如下：

优化要点总结

编译运行

在本样例根目录下执行如下步骤，编译并执行样例。

• 切换Case

  在cmake编译时通过-DSCENARIO_NUM=N指定要编译的Case，各Case说明：

  • 0: SIMT Warp直接写value
  • 1: MTE task queue搬运value

  示例：

  cmake .. -DCMAKE_ASC_ARCHITECTURES=dav-3510 -DSCENARIO_NUM=1

• 配置环境变量

  请根据当前环境上CANN开发套件包的安装方式，配置环境变量。

  source ${install_path}/cann/set_env.sh

  说明：${install_path}为CANN包安装目录，未指定安装目录时默认安装至/usr/local/Ascend下。

• 样例执行

  在样例目录下执行如下命令。

  SCENARIO_NUM=1 # 选择执行场景 mkdir -p build cd build cmake -DSCENARIO_NUM=$SCENARIO_NUM -DCMAKE_ASC_ARCHITECTURES=dav-3510 .. make -j ./hash_table_mte_queue 128

  也可以直接运行对比脚本，编译两种Case并采集16/32/64/128/256五组dim的性能数据。

  ./scripts/run.sh

• 编译选项说明

  选项	可选值	说明
  CMAKE_ASC_ARCHITECTURES	dav-3510	NPU架构，对应 Ascend 950PR/Ascend 950DT
  SCENARIO_NUM	0、1	Case编号：0=SIMT Warp直接写value，1=MTE task queue搬运value
  PROFILE_REPEAT_TIMES	正整数	单次进程中重复执行insert_or_assign的次数。对比脚本设置为1，使msOpProf采集一次完整insert路径。

• 执行结果

  执行结果如下，表示HashTable插入结果和值校验通过。

  Verification PASSED

性能分析

使用msOpProf获取性能数据。对比脚本会先执行一次样例并检查Verification PASSED，确认结果正确后再通过msprof op命令采集目标kernel性能数据。

Case 0测试方式如下：

cmake -DSCENARIO_NUM=0 -DPROFILE_REPEAT_TIMES=1 -DCMAKE_ASC_ARCHITECTURES=dav-3510 .. make -j ./hash_table_mte_queue 128 msprof op ./hash_table_mte_queue 128

Case 1测试方式如下：

cmake -DSCENARIO_NUM=1 -DPROFILE_REPEAT_TIMES=1 -DCMAKE_ASC_ARCHITECTURES=dav-3510 .. make -j ./hash_table_mte_queue 128 msprof op ./hash_table_mte_queue 128

当前目录下会生成OPPROF_前缀的文件夹，保存当前kernel的性能分析数据。

OPPROF_xxxx_XXXXXXXX ├── ArithmeticUtilization.csv ├── L2Cache.csv ├── Memory.csv ├── MemoryL0.csv ├── MemoryUB.csv ├── OpBasicInfo.csv ├── PipeUtilization.csv └── ResourceConflictRatio.csv

查看具体的性能分析结果：

# 查看Task Duration等基础信息 cat ./OPPROF_*/OpBasicInfo.csv # 查看Vector/MTE等流水耗时和占比 cat ./OPPROF_*/PipeUtilization.csv

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit