关于ub_size ub_block_num tile_num 的一点小心得

最新推荐文章于 2024-10-04 22:36:54 发布

william_myq

最新推荐文章于 2024-10-04 22:36:54 发布

阅读量609

点赞数 11

文章标签：人工智能

本文链接：https://blog.csdn.net/xiujiti6871/article/details/139367784

版权

正文前感谢昇腾各位工作人员，没有你们辛勤付出就没有我们的进步
这篇文章的立意在于gitee中算子Addcdiv中关于tiling切分的部分代码解析上
相信大家是从cuda，Nvidia 迁移到昇腾国产平台进行学习的，在达芬奇架构上面可能会有点迷
因为在架构概念上会有不一样的构思点，首先我们复习下cuda的架构体，这里不得不提下小白哥的陈年老贴
- https://www.hiascend.com/forum/thread-0232107179558829103-1-1.html

很熟悉是吧，三大天王,thread,block,grid
再来提下达芬奇架构的昇腾 ai_core
- 当单一ai_core中cube unit,vector unit都为1的情况下ai_core就是类似于 cuda 的thread，block_idx其实就是ai_core_idx
- 当单一ai_core中cube unit,vector unit不为1的情况下cube unit或者vector unit 就是cuda 的 thread,ai_core 相当于是cuda的block
至于grid，昇腾暂时没有这种概念
cuda 的block跟代码里面的block是不一样的概念，要注意区分，容易混淆
个人理解，讲的不对还望不吝赐教

引入痛点代码来分析
有很多变量，重点来分析ub_size ub_block_num tile_num
首先是ub_size，获取unified buffer的大小，这个是内部存储空间之一

uint64_t ub_size;//对应类型的存储空间大小，单位：字节
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);//获取硬件平台存储空间的内存大小
auto aivNum = ascendcPlatform.GetCoreNumAiv();//获取当前硬件平台AI Core中Vector Core的核数。
auto aicNum = ascendcPlatform.GetCoreNumAic();//获取当前硬件平台AI Core中Cube Core的核数复制

在200DK A2 也就是Ascend310B上运行命令
ub_size 为253952 ,也就是248KB、
vector core 为1
cube core为1

ub_size aivNum aicNum log is 253952,1,1复制

数据需要做对齐处理，如下
sizeofdatatype根据具体数据类型来定义
因为unified buffer都要求32字节对齐，所以最小的BLOCK_SIZE 定义为32字节
fp16 占用两个字节，一个BLOCK就有16个fp16
这里的BLOCK是数据块的意思，跟CUDA的不一样

  uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;//32/2=16保证block 以fp16对齐目的是为了确保每个数据块中的数据是按照指定的数据类型对齐的复制

获取的输入totalLength/ALIGN_NUM
就是总共有多少个fp16输入一个最小BLOCK的16个fp16,也就是32字节对齐

  // 1.输入向量满足32字节对齐 block_size 对齐
  if (totalLength % ALIGN_NUM != 0) {  //不对齐，先32位对齐
    totalLengthAligned =
        ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
  } else {
    totalLengthAligned = totalLength;
  }复制

接下来代码就很神奇了，应该是很让人迷惑的
前面有了block块，现在又有了ub_block_num ，ub_block_num 从哪里来的？
很不幸地讲，ub_block_num是自定义的，为什么引入这个变量，就要从后面tile_num结合分析

  uint32_t ub_block_num = 5;  //为测试方便，验证代码流程
  uint32_t tile_num;

  if (ub_block_num % 2 != 0) {//分配的内存块个数
    ub_block_num = ub_block_num - 1;//5-1=4
  }复制

tile_num的计算如下
个人理解，首先将全部数据进行block_dim核均分，可类比与thread均分
然后每个核承担的数据，再进行更小块切分，为什么要继续切分呢？不能直接计算吗？
- 每次最多读取连续8个block进行计算
- 个人理解这个ub_block_num最大就是8
- 为了后面乒乓操作好做处理，所以ub_block_num设为偶数
- 所以tile_num应该理解为tile_num个最小BLOCK（32字节）
总结
- 先对输入数据做block_dim核切分
- 再对单核切分完成数据进行tile_num个最小BLOCK切分
- 最后剩下的数据就是lasttileLength

if ((totalLengthAligned / ALIGN_NUM) % block_dim == 0) {  //核间可均分
    blockLength = totalLengthAligned / block_dim;//核均分
    tile_num = blockLength / ALIGN_NUM / ub_block_num;//当blockLength/ALIGN_NUM < ub_block_num
    if ((totalLengthAligned / block_dim / ALIGN_NUM) % ub_block_num == 0 ||
        tile_num == 0) {  
          //满足32字节对齐，可以核内均分
      if (tile_num == 0) {
        tile_num = 1;
      } 
      if (blockLength < ub_block_num* ALIGN_NUM) {
        tileLength = ((blockLength / ALIGN_NUM) + 1) / 2 * 2 * ALIGN_NUM;
        lasttileLength = tileLength;
      } else {
        tileLength = ub_block_num * ALIGN_NUM;
        lasttileLength = tileLength;
      }
    } else {  //满足32字节对齐，核内不能均分
      tile_num = tile_num + 1;
      tileLength = ub_block_num * ALIGN_NUM;
      lasttileLength = blockLength - (tile_num - 1) * tileLength;
    }