CANN/shmem原理与架构详解

SHMEM原理概述【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmemSHMEM初始化流程介绍shmem的初始化接口int aclshmemx_init_attr(aclshmemx_init_attr_t *attributes)会根据传入的参数完成shmem功能所需资源的初始化。其中包含多进程间信息的同步以及建链、虚拟内存及device上物理内存的分配和映射、host和device间state信息同步以及初始化共享内存管理、team管理、同步管理功能所需要的资源。这些资源信息都会记录在一个aclshmem_device_host_state_t的state里。初始化流程如下图所示参数设置及参数校验相关内容不做具体介绍。多进程间的建链以uid初始化流程为例。一般会在初始化接口调用前通过torch的能力将rank0的uid信息广播到所有进程内部包含rank0节点的ip、port、magic通信域标识符等信息。多进程间的建链是基于TCP socket实现的。所有进程会先基于uid内部的ip::port信息和rank0建立连接连接时会根据uid内magic信息是否一致判断是否为同一通信域不同则断开连接。连接成功后所有rank与rank0通信基于此实现host侧的allgather和barrier等能力。该部分能力并非aclshmem_barrier相关能力而是类似于MPI_Barrier内存堆的初始化当前基于驱动能力实现分配虚拟内存、再分配物理内存然后实现虚拟内存和物理内存间的映射可以在连续的虚拟地址空间管理访问共享内存地址。实际效果和流程可以参考acl接口 ACL虚拟内存管理接口该部分会分配出两块内存用户申请的共享内存空间每个rank申请的内存大小和初始化传入attributes的local_mem_size大小一致后续可以通过aclshmem_malloc、aclshmem_free等共享内存接口管理。用于在device保存shmem的state、team等元数据空间当state和team等信息发生变化进程将自动同步到元数据空间元数据空间32M不对外提供访问接口host和device间state信息同步将host的state信息同步。即当host的state发生变化时将其复制一份覆盖到上个阶段分配的在device侧存储state信息的内存中后续在device侧操作时可以直接拿到更新后的state信息。共享内存管理初始化基于heap_base和heap_size初始化一个memory_manager用于后续共享内存管理。heap_base为当前rank的共内存首地址heap_size为用户申请的当前rank的共享内存大小。aclshmem_malloc内部采用first fit找到第一个满足要求的chunk从该chunk中分离出对应大小的chunk剩余空间如果有则作为新的空闲chunk。aclshmem_free释放后补充空闲chunk 。若新添加chunk与其他空闲chunk连续则合并。aclshmem_malloc和aclshmem_free接口使用时需要在所有进程同步调用且需要分配或释放相同大小的内存。SHMEM内部默认aclshmem_malloc分配的内存在虚拟地址空间对称即当前rank malloc的地址heap_size下一个rank malloc的地址。如果进程间分配内存不同会导致后续分配的内存首地址不对称无法通过地址偏移访问到正确的数据。team管理初始化创建一个ACLSHMEM_MAX_TEAMS大小的全局team_pool内部信息全部初始化为-1。 初始化全局team即ACLSHMEM_TEAM_WORLDstart 0 ; stride 1 ; size npes ; mype mype 即从第0个rank开始步长为1全局的pe数量为npes当前pe为mype(team相关具体介绍会在team部分做详细梳理此处仅简单展示全局team的属性信息)。在device上分配并初始化后续team级别同步需要的资源sync_pool、sync_counter、core_sync_pool、core_sync_counter。前俩个成员用于team内所有的卡间同步、后两个成员用于卡内核间同步。team_pool、sync_pool、sync_counter、core_sync_pool、core_sync_counter这些信息都会在state里储存。同步管理初始化该部分仅获取一个ffts地址后续可通过shmemx_get_ffts_config获得在算子内通过shmemx_set_ffts_config设置用于runtime同步AscendC层面会影响的SyncAll、CrossCoreSetFlag和CrossCoreWaitFlag接口shmem层面会影响到barrier相关接口。shmemx_set_ffts_config开销较小建议算子内都调用一次该接口。SHMEM通信域Team介绍Team是shmem的通信域概念在相关接口中可以通过team_id访问初始化后会有一个默认的全局通信域其team_id是ACLSHMEM_TEAM_WORLD 0。team的信息存储在在state的team_pools里team_pools是一个aclshmemx_team_t的数组。aclshmemx_team_t内会存储当前team的idteam相关接口使用的索引、当前进程在该team内的rank id、该team内的起始rank、rank间步长、rank数量等相关信息。aclshmemx_team_t内存储的mype和size是team内部的信息aclshmem_device_host_state_t里存储的mype和npes是全局的信息。例如对4个rank初始化shmem前两卡和后两卡各为一个team。此时这四个rank的state里的mype分别是0,1,2,3。npes则都是4。但team_pools里的mype则分别是0101。size为2。子Team切分shmem提供了专门的接口进行子Team的切分。int aclshmem_team_split_strided(aclshmem_team_t parent_team, int pe_start, int pe_stride, int pe_size, aclshmem_team_t *new_team);parent_team为父teampe_start为起始pepe_stride为每次划分的步长pe_size为划分的新team里pe的个数new_team是出参是切分得到的新team的team_id。以初始化好8个rank的场景为例以如下方式调用切分接口。aclshmem_team_t new_team; // 从全局通信域中idx为1的pe开始以步长为2切分出3个pe后停止。 auto ret aclshmem_team_split_strided(ACLSHMEM_TEAM_WORLD, 1, 2, 3, new_team);此时new_team中有3个pe在state中的mype信息分别为1,3,5他们的team_pools里的mype信息则分别为0,1,2。Team的使用当算子只需要在部分rank运行时就会需要用到team相关接口,如aclshmem_team_my_pe(my_team)可以返回当前rank在my_team的mype信息aclshmem_my_pe()可以返回当前rank在全局的mype信息。通常我们可以用team级别的mype信息作为算子内部资源数组的索引使用全局的mype作为全局共享内存地址信息的索引。同步接口也会使用到team_id如shmem内部提供的team内的同步接口aclshmem_barrier(aclshmem_team_t tid)。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考