寒武纪显卡实现softmax算子

本文介绍了在寒武纪GPU上实现softmax算法的详细步骤,包括unionMaxKernel、blockMaxKernel、unionSumKernel、blockSumKernel和softmaxKernel五个核心步骤。利用__bang_atomic_add和__bang_atomic_max进行融合优化,同时讲解了如何在计算过程中同时获取全局最大值和求和。文章还提到了数值精度问题以及如何利用__bang_active_exp_less_0提高计算精度。

寒武纪显卡实现softmax基本逻辑

寒武纪实现softmax包括下面5个步骤,我们也采取5个kernel来实现softmax:
unionMaxKernel(float* middle, float* source1, int num),这个kernel使用的任务类型是union1,其中middle的长度为taskDim,达到的目的是middle对应元素存储的是不同taskId处理的那部分数据的局部max

blockMaxKernel(float* dstMax, float* middle, int taskNum),这个kernel使用的任务类型是BLOCK,任务总数也是1,其中dstMax的长度为1,达到的目的是将长度为taskDim的middle进一步做max规约得到全局max

unionSumKernel(float* middle, float* source1, int num, float globalMax),这个kernel使用的任务类型是union1,其中middle的长度为taskDim,达到的目的是middle对应元素存储的是不同taskId处理的那部分数据的局部数值和,其中globalMax就是上面两个kernel计算出来的全局max
这个函数特别值得注意的是,由于寒武纪显卡无法直接针对向量source1的特定元素进行操作,为此我们引入了一个中间向量tmp[maxNum],其中tmp每个元素都是-globalMax,借助于__bangd_add来实现source1 = source1 + tmp。除此之外,寒武纪有3个实现exp的函数,但是经过查证,这三个exp函数的数值精度差距很大,由于经过线性变换以后source1的元素都在负半轴,因此这里我们使用__bang_active_exp_less_0函数来提高精度,详细内容可以参考官网链接说明

blockSumKernel(float *dstSum, float *middle, int taskNum),这个kernel使用的任务类型是BLOCK,任务总数也是1,其中dstSum的长度为1,达到的目的是将长度为taskDim的middle进一步做sum规约得到全局sum

softmaxKernel(float *dst, float *source1, float globalMax, float globalSum, int num),这个kernel使用的dst,source1长度都是num,globalMax,globalSum在上面已经计算出来,最终希望dst存储最终的softmax结果。
编译命令为:

cncc softmax.mlu -o softmax --bang-mlu-arch=mtp_372 -O3

#include <bang.h>
#include <bang_device_functions.h>
#define EPS 1e-7
const int NRAM_MAX_SIZE = 1024 * 128;//后续树状求和必须保证NRAM_MAX_SIZE为2的幂次
const int maxNum = NRAM_MAX_SIZE/sizeof(float); //NRAM上最多存储maxNum个float元素
const int warpSize = 32;//__bang_reduce_sum每次从src取128字节数据相加,对应32个float元素,并且0-31的结果保存在索引0,32-63的结果保存在索引1

__nram__ float src1[maxNum];//每次搬运maxNum数据到NRAM
__nram__ float destMax;//方便后面比较最大值
__nram__ float destSum[maxNum];//后面数值求和
__nram__ float destSumFinal[warpSize];//将destSum规约到destFinal[0]
__mlu_entry__ void unionMaxKernel(float* middle, float* source1, int num) {
  
  int remain = num%taskDim;//如果不能整除,则让前部分taskId多处理一个元素
  int stepEasy = (num - remain)/taskDim;
  int stepHard = stepEasy + 1;
  int step = (taskId < remain ? stepHard : stepEasy);//前部分tsakId多处理一个元素
  int indStart = (taskId < remain ? taskId * stepHard : remain * stepHard + (taskId - remain) * stepEasy);
  int remainNram = step%maxNum;
  int repeat = step/maxNum;//如果一个task处理元素个数超出NRAM最大内存,则需要for循环
  //maxNum尽量取大一些,免得repeat过大导致求和过程累加过于严重,使得大数吃小数
  source1 = source1 + indStart;//设定起始偏移量

  
  destMax = -INFINITY;//初始化为负无穷
  __nram__ float srcMax[2];
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destMax < srcMax[0]){
      destMax = srcMax[0];
    }
  }
 
  if(remainNram){
    __bang_write_value(src1, maxNum, -INFINITY);//必须要初始化src1全部元素为负无穷
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destMax < srcMax[0]){
      destMax = srcMax[0];
    }
  }//结束以后向量destMax保存了source1[indSart:indStart+step]这部分数据的全局最大值
  
  __memcpy(middle + taskId, &destMax, sizeof(float), NRAM2GDRAM);//middle长度为taskDim
}
//----------------------

__mlu_entry__ void blockMaxKernel(float* dstMax, float* middle, int taskNum) {//将长度为taskDim的middle继续做Max规约
  
  int remain = taskNum%warpSize;
  int repeat = (taskNum - remain)/warpSize;//如果taskDim太大,超过warpSize,使用for循环规约
  __nram__ float srcMid[warpSize];
  __nram__ float srcMax[2];
  destMax = -INFINITY;
  for(int i = 0; i < repeat; i++){
    __memcpy(srcMid, middle + i * warpSize, warpSize * sizeof(float), GDRAM2NRAM);//每次迁移32个float数据到NRAM
    __bang_argmax(srcMax, srcMid, warpSize);
    if(destMax < srcMax[0]){
      destMax = srcMax[0];
    }
  }
  if(remain){
    __bang_write_value(srcMid, warpSize, -INFINITY);//初始化srcMid全部元素为负无穷
    __memcpy(srcMid, middle + repeat * warpSize, remain * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, srcMid, warpSize);
    if(destMax < srcMax[0]){
      destMax = srcMax[0];
    }
  }
 
  __memcpy(dstMax, &destMax, sizeof(float), NRAM2GDRAM);//这个kernel只能使用Block类型,1个任务
}

__mlu_entry__ void unionSumKernel(float* middle, float* source1, int num, float globalMax) {
  
  int remain = num%taskDim;//如果不能整除,则让前部分taskId多处理一个元素
  int stepEasy = (num - remain)/taskDim;
  int stepHard = stepEasy + 1;
  int step = (taskId < remain ? stepHard : stepEasy);//前部分tsakId多处理一个元素
  int indStart = (taskId < remain ? taskId * stepHard : remain * stepHard + (taskId - remain) * stepEasy);
  int remainNram = step%maxNum;
  int repeat = step/maxNum;//如果一个task处理元素个数超出NRAM最大内存,则需要for循环
  //maxNum尽量取大一些,免得repeat过大导致求和过程累加过于严重,使得大数吃小数
  source1 = source1 + indStart;//设定起始偏移量

  __nram__ float tmp[maxNum];
  __nram__ float srcMax[2];
  __bang_write_value(tmp, maxNum, -globalMax);//初始化tmp全部元素为-globalMax
  __bang_write_zero(destSum, maxNum);
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, tmp, src1, maxNum);//src1 = src1 - globalMax 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    //__bang_active_exphp(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    __bang_add(destSum, destSum, src1, maxNum);//destSum += exp(src1 - globalMax)
  }
  if(remainNram){
    __bang_write_value(src1, maxNum, globalMax);
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, tmp, src1, maxNum);//src1 = src1 - globalMax ,后面maxNum-remainNram这部分直接为0
    
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    //__bang_active_exphp(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    
    __bang_add(destSum, destSum, src1, maxNum);//destSum在原来基础上又多加了(maxNum - remainNram)
    
  }//结束以后长度为maxNum的向量destSum保存了source1[indSart:indStart+step]这部分数据的数值和+(maxNum - remainNram)
  //__bang_printf("destSum[%d]:%.6f, src1:%.6f\n",remainNram, destSum[remainNram], src1[remainNram]);
  //下面开始针对destSum做规约
  __bang_write_zero(destSumFinal, warpSize);//初始化destSumFinal全部元素为0
  int segNum = maxNum / warpSize;//将destSum分成segNum段,每段向量长度为warpSize,分段进行树状求和,segNum要求是2的幂次
  for(int strip = segNum/2; strip > 0; strip = strip / 2){//segNum要求是2的幂次即maxNum必须选取2的幂次
    for(int i = 0; i < strip ; i++){
      __bang_add(destSum + i * warpSize, destSum + i * warpSize, destSum + (i + strip) * warpSize, warpSize);
    } 
  }
  __bang_reduce_sum(destSumFinal, destSum, warpSize);
  destSumFinal[0] = destSumFinal[0] - (maxNum - remainNram);//把上面多加的(maxNum - remainNram)减掉
  //__bang_printf("taskId:%d,maxNum - remainNram:%d,but get sum:%.6f\n",taskId, maxNum - remainNram, destSumFinal[0]);
  __memcpy(middle + taskId, destSumFinal, sizeof(float), NRAM2GDRAM);
}
//----------------------

__mlu_entry__ void blockSumKernel(float *dstSum, float *middle, int taskNum)//将长度为taskDim的middle继续做Sum规约
{
    
    int remain = taskNum % warpSize;
    int repeat = (taskNum - remain) / warpSize;
    __nram__ float srcMid[warpSize];
    __bang_write_zero(destSumFinal, warpSize); // 初始化destSumFinal全部元素为0
    //__bang_printf("sum:%.6f\n",destSumFinal[0]);
    for (int i = 0; i < repeat; i++)
    {
        __memcpy(srcMid, middle + i * warpSize, warpSize * sizeof(float), GDRAM2NRAM); // 每次迁移32个float数据到NRAM
        __bang_add(destSumFinal, destSumFinal, srcMid, warpSize);                      // destSumFinal存储add结果
    }
    if (remain)
    {
        __bang_write_zero(srcMid, warpSize); // 初始化destSumFinal全部元素为0
        __memcpy(srcMid, middle + repeat * warpSize, remain * sizeof(float), GDRAM2NRAM);
        __bang_add(destSumFinal, destSumFinal, srcMid, warpSize); // destSumFinal存储add结果
    }
    __bang_reduce_sum(destSumFinal, destSumFinal, warpSize); // 针对destSumFinal规约即可把结果保存到destSumFinal[0]
    //__bang_printf("xiao,taskId:%d,sum:%.6f\n", taskId, destSumFinal[0]);
    __memcpy(dstSum, destSumFinal, sizeof(float), NRAM2GDRAM); // 这个kernel只能使用Block类型,1个任务
}
__mlu_entry__ void softmaxKernel(float *dst, float *source1, float globalMax, float globalSum, int num){
  int remain = num%taskDim;//如果不能整除,则让前部分taskId多处理一个元素
  int stepEasy = (num - remain)/taskDim;
  int stepHard = stepEasy + 1;
  int step = (taskId < remain ? stepHard : stepEasy);//前部分tsakId多处理一个元素
  int indStart = (taskId < remain ? taskId * stepHard : remain * stepHard + (taskId - remain) * stepEasy);
  int remainNram = step%maxNum;
  int repeat = step/maxNum;//如果一个task处理元素个数超出NRAM最大内存,则需要for循环
  //maxNum尽量取大一些,免得repeat过大导致求和过程累加过于严重,使得大数吃小数
  source1 = source1 + indStart;//设定起始偏移量
  dst = dst + indStart;//设定起始偏移量
  float globalSumInv = 1.0/globalSum;
  
  for(int i = 0; i < repeat; i++){
    __bang_write_value(destSum, maxNum, -globalMax);//初始化destSum全部元素为-globalMax
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, destSum, src1, maxNum);//src1 = src1 - globalMax 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    //__bang_active_exphp(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    __bang_write_value(destSum, maxNum, globalSumInv);//初始化destSum全部元素为globalSumInv,使用1.0/globalSum编译报错
    __bang_mul(src1, src1, destSum, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + i * maxNum, src1, maxNum * sizeof(float), NRAM2GDRAM);
  }
  
  if(remainNram){ 
    __bang_write_value(src1, maxNum, -globalMax);
    __bang_write_value(destSum, maxNum, -globalMax);//初始化destSum全部元素为-globalMax
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, destSum, src1, maxNum);//src1 = src1 - globalMax 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    //__bang_active_exphp(src1, src1, maxNum);//src1 = exp(src1 - globalMax)
    __bang_write_value(destSum, maxNum, globalSumInv);//初始化destSum全部元素为globalSumInv,使用1.0/globalSum编译报错
    __bang_mul(src1, src1, destSum, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + repeat * maxNum, src1, remainNram * sizeof(float), NRAM2GDRAM);
  }
  //__bang_printf("Inv:%.6f\n",globalSumInv);
}
int main(void)
{
  int num = 102001010;
  //int num = 10;
  cnrtQueue_t queue;
  CNRT_CHECK(cnrtSetDevice(0));
  CNRT_CHECK(cnrtQueueCreate(&queue));

  cnrtDim3_t dim = {4, 1, 1};
  int taskNum = dim.x * dim.y * dim.z;
  cnrtFunctionType_t ktype = CNRT_FUNC_TYPE_UNION1;

  cnrtNotifier_t start, end;
  CNRT_CHECK(cnrtNotifierCreate(&start));
  CNRT_CHECK(cnrtNotifierCreate(&end));

  float* host_dst = (float*)malloc(num * sizeof(float));
  float* host_src1 = (float*)malloc(num * sizeof(float));
  

  for (int i = 0; i < num; i++) {
    host_src1[i] = i%4;
  }

  float* mlu_middle;
  float* mlu_dstMax;
  float* mlu_dstSum;
  float* mlu_dst;
  float* mlu_src1;
  
  CNRT_CHECK(cnrtMalloc((void**)&mlu_middle, taskNum * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_dstMax, sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_dstSum, sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_dst, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_src1, num * sizeof(float)));
 

  CNRT_CHECK(cnrtMemcpy(mlu_src1, host_src1, num * sizeof(float), cnrtMemcpyHostToDev));
  
  //----------------------------
  CNRT_CHECK(cnrtPlaceNotifier(start, queue));
  unionMaxKernel<<<dim, ktype, queue>>>(mlu_middle, mlu_src1, num);
  
  cnrtQueueSync(queue);
  //---------------------------
  cnrtDim3_t dimBlock = {1, 1, 1};
  
  blockMaxKernel<<<dimBlock, CNRT_FUNC_TYPE_BLOCK, queue>>>(mlu_dstMax, mlu_middle, taskNum);
  
  cnrtQueueSync(queue);
  float globalMax;
  CNRT_CHECK(cnrtMemcpy(&globalMax, mlu_dstMax, sizeof(float), cnrtMemcpyDevToHost));
  printf("max:%.6f\n",globalMax);
  //----------------------------
  
  unionSumKernel<<<dim, ktype, queue>>>(mlu_middle, mlu_src1, num, globalMax);
  
  cnrtQueueSync(queue);
  //---------------------------
  
  blockSumKernel<<<dimBlock, CNRT_FUNC_TYPE_BLOCK, queue>>>(mlu_dstSum, mlu_middle, taskNum); 
  cnrtQueueSync(queue);
  float globalSum;
  CNRT_CHECK(cnrtMemcpy(&globalSum, mlu_dstSum, sizeof(float), cnrtMemcpyDevToHost));
  printf("sum:%.6f\n",globalSum);
  //----------------------------
  softmaxKernel<<<dim, ktype, queue>>>(mlu_dst, mlu_src1, globalMax, globalSum, num);
  CNRT_CHECK(cnrtPlaceNotifier(end, queue));
  cnrtQueueSync(queue);
  //printf("max:%.6f,sum:%.6f\n", globalMax ,globalSum);
  //---------------------------
  CNRT_CHECK(cnrtMemcpy(host_dst, mlu_dst, num * sizeof(float), cnrtMemcpyDevToHost));
  for(int i = 0; i < 10; i++){
    printf("softmax[%d]:%.6e,origin:%.6f\n", i, host_dst[i], host_src1[i]);
  }
  float timeTotal;
  CNRT_CHECK(cnrtNotifierDuration(start, end, &timeTotal));
  printf("Total Time: %.3f ms\n", timeTotal / 1000.0);

  CNRT_CHECK(cnrtQueueDestroy(queue));

  cnrtFree(mlu_middle);
  cnrtFree(mlu_src1);
  
  free(host_dst);
  free(host_src1);
  

  return 0;
}











softmax融合

在这里插入图片描述
借助于__bang_atomic_add和__bang_atomic_max两个函数,我们可以把上面的5个kernel融合成一个kernel,代码框架逻辑和上面完全一样,我们以__bang_atomic_add函数来重点介绍一下寒武纪显卡的原子操作,该函数涉及到三个向量的原子操作,其中dst和src2都是NRAM上的向量,但是src1是__mlu_device__上的向量,对于我们的代码,比如说上面的sum规约例子,不同的taskId最后都会得到一个结果保存在destSumFinal[0],这个结果保存的是taskId对应的这部分数据的数值和,而我们需要获得全局数值和。之前的做法是,我们把不同taskId的结果保存在全局内存变量middle里面,middle的长度为taskDim,然后借助于另一个kernel对middle规约。现在我们有了原子操作,我们可以借助于原子操作把不同taskId的结果直接累加到全局变量里面。
在这里插入图片描述
上面三行代码结束以后,不管是哪个taskId,计算出来的globalSum都是一样的结果,也就是不同taskId对应的destSumFinal[0]相加的结果,这里特别值得注意的就是调用原子操作之前一定要先做同步,保证不同taskId都计算结束,其次就是globalSum必须设置为指针型向量,不能使用常数。详细代码如下所示:

#include <bang.h>
#include <bang_device_functions.h>
#define EPS 1e-7
const int NRAM_MAX_SIZE = 1024 * 128;//后续树状求和必须保证NRAM_MAX_SIZE为2的幂次
const int maxNum = NRAM_MAX_SIZE/sizeof(float); //NRAM上最多存储maxNum个float元素
const int warpSize = 32;//__bang_reduce_sum每次从src取128字节数据相加,对应32个float元素,并且0-31的结果保存在索引0,32-63的结果保存在索引1

__nram__ float src1[maxNum];//每次搬运maxNum数据到NRAM
__nram__ float destMax;//方便后面比较最大值
__nram__ float destSum[maxNum];//后面数值求和
__nram__ float destSumFinal[warpSize];//将destSum规约到destFinal[0]
__nram__ float srcMax[2];//存储最大值

__mlu_entry__ void softmaxKernel(float* dst, float* source1, float* globalMax, float* globalSum, int num) {
  
  int remain = num%taskDim;//如果不能整除,则让前部分taskId多处理一个元素
  int stepEasy = (num - remain)/taskDim;
  int stepHard = stepEasy + 1;
  int step = (taskId < remain ? stepHard : stepEasy);//前部分tsakId多处理一个元素
  int indStart = (taskId < remain ? taskId * stepHard : remain * stepHard + (taskId - remain) * stepEasy);
  int remainNram = step%maxNum;
  int repeat = step/maxNum;//如果一个task处理元素个数超出NRAM最大内存,则需要for循环
  //maxNum尽量取大一些,免得repeat过大导致求和过程累加过于严重,使得大数吃小数
  source1 = source1 + indStart;//设定起始偏移量

  //------------------------------------下面开始计算max
  destMax = -INFINITY;//初始化为负无穷
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destMax < srcMax[0]){
      destMax = srcMax[0];
    }
  }
  if(remainNram){
    __bang_write_value(src1, maxNum, -INFINITY);//必须要初始化src1全部元素为负无穷
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);
    if(destMax < srcMax[0]){
      destMax = srcMax[0];
    }
  }//结束以后向量destMax保存了source1[indSart:indStart+step]这部分数据的全局最大值
  globalMax[0] = -INFINITY;
  __sync_all();
  
  __bang_atomic_max(&destMax, globalMax, &destMax, 1);//globalMax[0]必须初始化为负无穷
  
  //------------------------------------下面开始计算sum
  __nram__ float tmp[maxNum];
  __bang_write_value(tmp, maxNum, -globalMax[0]);//初始化tmp全部元素为-globalMax[0]
  __bang_write_zero(destSum, maxNum);
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, tmp, src1, maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_add(destSum, destSum, src1, maxNum);//destSum += exp(src1 - globalMax[0])
  }
  if(remainNram){
    __bang_write_value(src1, maxNum, globalMax[0]);
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, tmp, src1, maxNum);//src1 = src1 - globalMax[0] ,后面maxNum-remainNram这部分直接为0
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_add(destSum, destSum, src1, maxNum);//destSum在原来基础上又多加了(maxNum - remainNram)
    
  }//结束以后长度为maxNum的向量destSum保存了source1[indSart:indStart+step]这部分数据的数值和+(maxNum - remainNram)
  
  __bang_write_zero(destSumFinal, warpSize);//初始化destSumFinal全部元素为0
  int segNum = maxNum / warpSize;//将destSum分成segNum段,每段向量长度为warpSize,分段进行树状求和,segNum要求是2的幂次
  for(int strip = segNum/2; strip > 0; strip = strip / 2){//segNum要求是2的幂次即maxNum必须选取2的幂次
    for(int i = 0; i < strip ; i++){
      __bang_add(destSum + i * warpSize, destSum + i * warpSize, destSum + (i + strip) * warpSize, warpSize);
    } 
  }
  __bang_reduce_sum(destSumFinal, destSum, warpSize);
  destSumFinal[0] = destSumFinal[0] - (maxNum - remainNram);//把上面多加的(maxNum - remainNram)减掉
  __sync_all();
  globalSum[0] = 0.0;
  __bang_atomic_add(destSumFinal, globalSum, destSumFinal, 1);//globalSum[0]必须初始化为0
  //------------------------------------下面开始计算softmax
  dst = dst + indStart;//设定起始偏移量
  float globalSumInv = 1.0/globalSum[0];
  
  for(int i = 0; i < repeat; i++){
    __bang_write_value(destSum, maxNum, -globalMax[0]);//初始化destSum全部元素为-globalMax[0]
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, destSum, src1, maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_write_value(destSum, maxNum, globalSumInv);//初始化destSum全部元素为globalSum[0]Inv,使用1.0/globalSum[0]编译报错
    __bang_mul(src1, src1, destSum, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + i * maxNum, src1, maxNum * sizeof(float), NRAM2GDRAM);
  }
  
  if(remainNram){ 
    __bang_write_value(src1, maxNum, globalMax[0]);
    __bang_write_value(destSum, maxNum, -globalMax[0]);//初始化destSum全部元素为-globalMax[0]
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, destSum, src1, maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_write_value(destSum, maxNum, globalSumInv);//初始化destSum全部元素为globalSum[0]Inv,使用1.0/globalSum[0]编译报错
    __bang_mul(src1, src1, destSum, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + repeat * maxNum, src1, remainNram * sizeof(float), NRAM2GDRAM);
  }
  
  __bang_printf("taskId:%d,max:%.6f, sum:%.6f\n",taskId, globalMax[0], globalSum[0]);
}


int main(void)
{
  //int num = 102001010;
  int num = 10;
  cnrtQueue_t queue;
  CNRT_CHECK(cnrtSetDevice(0));
  CNRT_CHECK(cnrtQueueCreate(&queue));

  cnrtDim3_t dim = {4, 1, 1};
  int taskNum = dim.x * dim.y * dim.z;
  cnrtFunctionType_t ktype = CNRT_FUNC_TYPE_UNION1;

  cnrtNotifier_t start, end;
  CNRT_CHECK(cnrtNotifierCreate(&start));
  CNRT_CHECK(cnrtNotifierCreate(&end));

  float* host_dst = (float*)malloc(num * sizeof(float));
  float* host_src1 = (float*)malloc(num * sizeof(float));
  

  for (int i = 0; i < num; i++) {
    host_src1[i] = i%4;
  }

  float* mlu_dst;
  float* mlu_src1;
  float* globalMax;
  float* globalSum;
  CNRT_CHECK(cnrtMalloc((void**)&mlu_dst, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_src1, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&globalMax, sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&globalSum, sizeof(float)));

  CNRT_CHECK(cnrtMemcpy(mlu_src1, host_src1, num * sizeof(float), cnrtMemcpyHostToDev));
  
  //----------------------------
  CNRT_CHECK(cnrtPlaceNotifier(start, queue));
  softmaxKernel<<<dim, ktype, queue>>>(mlu_dst, mlu_src1, globalMax, globalSum, num);
  CNRT_CHECK(cnrtPlaceNotifier(end, queue));
  cnrtQueueSync(queue);

  //---------------------------
  CNRT_CHECK(cnrtMemcpy(host_dst, mlu_dst, num * sizeof(float), cnrtMemcpyDevToHost));
  for(int i = 0; i < 10; i++){
    printf("softmax[%d]:%.6e,origin:%.6f\n", i, host_dst[i], host_src1[i]);
  }
  float timeTotal;
  CNRT_CHECK(cnrtNotifierDuration(start, end, &timeTotal));
  printf("Total Time: %.3f ms\n", timeTotal / 1000.0);

  CNRT_CHECK(cnrtQueueDestroy(queue));

  cnrtFree(mlu_dst);
  cnrtFree(mlu_src1);
  cnrtFree(globalMax);
  cnrtFree(globalSum);
  
  free(host_dst);
  free(host_src1);
  

  return 0;
}









softmax同时计算max和sum

回忆之前英伟达平台编写softmax的经验,当时我们引入了一个结构体来存储计算过程中的max和数值和,
在这里插入图片描述
这里我们也采取类似的思路来减少从GDRM到NRAM的数据访问频率,我们以下面这段代码为例子重点介绍如何运用寒武纪框架一边计算全局max,一边计算数值和。
代码29行初始化一个变量destNewMax为负无穷,方便后面比较最大值
进入循环以后,代码31行开始根据循环体不断读取全局变量不同位置的数据,将数值转移到NRAM数组src1上
代码32行借助于函数__bang_argmax计算出src1的最大值和对应的索引,注意__bang_argmax的两个向量参数长度必须相同,但是第一个向量计算以后索引为0的位置存储的是value,索引为1的位置存储的是index
代码33行——35行开始判断是否要更新最大值destNewMax
代码36行——38行其实就是在做一个src1 = exp(src1-destNewMax),只不过本人一时半会没有找到一个函数可以直接让向量减去一个常数,为此这里使用__bang_add来达成目的
代码39行——44行,这里必须要判断当前循环数是否大于0,i = 0的时候destSum = src1即可,但是一旦i大于0,destSum = destSum×exp(destOldMax - destNewMax) + src1,这一步就是保证数值和正确的关键,每次循环都会更新最大值,所以每次循环都要对上一次的数值和destSum做一个缩放,并且每次循环都需要更新destOldMax的数值
代码46行——61行这部分代码的逻辑和上面这个循环类似,只不过是针对无法整除多余的那部分做了一点特殊处理:
代码47行这里必须要对src1重新初始化为负无穷,因为经过上面的循环以后,src1都已经有值了,如果这里步重新初始化覆盖,就会导致前remainNram个元素是source1末尾的这部分元素,但是后面maxNum-remainNram这部分元素还是上一轮循环的结果,此时盲目去比较最大值,是无法得到前remainNram个元素的最大值的
代码53行也是类似,必须要重新初始化src1,要不然src1会有一部分负无穷
在这里插入图片描述

#include <bang.h>
#include <bang_device_functions.h>
#define EPS 1e-7
const int NRAM_MAX_SIZE = 1024 * 128;//后续树状求和必须保证NRAM_MAX_SIZE为2的幂次
const int maxNum = NRAM_MAX_SIZE/sizeof(float); //NRAM上最多存储maxNum个float元素
const int warpSize = 32;//__bang_reduce_sum每次从src取128字节数据相加,对应32个float元素,并且0-31的结果保存在索引0,32-63的结果保存在索引1

__nram__ float src1[maxNum];//每次搬运maxNum数据到NRAM
__nram__ float destSum[maxNum];//后面数值求和
__nram__ float destSumFinal[warpSize];//将destSum规约到destFinal[0]
__nram__ float srcMax[2];

__mlu_entry__ void softmaxKernel(float* dst, float* source1, float* globalMax, float* globalSum, int num) {
  
  int remain = num%taskDim;//如果不能整除,则让前部分taskId多处理一个元素
  int stepEasy = (num - remain)/taskDim;
  int stepHard = stepEasy + 1;
  int step = (taskId < remain ? stepHard : stepEasy);//前部分tsakId多处理一个元素
  int indStart = (taskId < remain ? taskId * stepHard : remain * stepHard + (taskId - remain) * stepEasy);
  int remainNram = step%maxNum;
  int repeat = step/maxNum;//如果一个task处理元素个数超出NRAM最大内存,则需要for循环
  //maxNum尽量取大一些,免得repeat过大导致求和过程累加过于严重,使得大数吃小数
  source1 = source1 + indStart;//设定起始偏移量

  //------------------------------------下面开始计算max
  __nram__ float tmp[maxNum];
  __nram__ float destOldMax;
  __nram__ float destNewMax;
  __bang_write_zero(destSum, maxNum);
  destNewMax = -INFINITY;//初始化为负无穷
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destNewMax < srcMax[0]){
      destNewMax = srcMax[0];//更新最大值
    }
    __bang_write_value(tmp, maxNum, -destNewMax);
    __bang_add(src1, tmp, src1, maxNum);//src1 = src1 - 最大值
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - 最大值)
    if(i > 0){
      __bang_write_value(tmp, maxNum, exp(destOldMax - destNewMax));
      __bang_mul(destSum, destSum, tmp, maxNum);//destSum = destSum * exp(destOldMax - destNewMax)
    }
    __bang_add(destSum, destSum, src1, maxNum);//destSum = destSum + exp(src1 - destNewMax)
    destOldMax = destNewMax;
  }
  if(remainNram){
    __bang_write_value(src1, maxNum, -INFINITY);//必须要初始化src1全部元素为负无穷
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destNewMax < srcMax[0]){
      destNewMax = srcMax[0];
    }
    __bang_write_value(src1, maxNum, destNewMax);//必须重新初始化为destNewMax
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);//必须再次读取
    __bang_write_value(tmp, maxNum, -destNewMax);
    __bang_add(src1, tmp, src1, maxNum);//后面maxNum-remainNram部分为0
    __bang_active_exp_less_0(src1, src1, maxNum);//相当于多加了maxNum-remainNram
    if(repeat > 0){
      __bang_write_value(tmp, maxNum, exp(destOldMax - destNewMax));
      __bang_mul(destSum, destSum, tmp, maxNum);
    }
    __bang_add(destSum, destSum, src1, maxNum);
    destOldMax = destNewMax;
  }//结束以后向量destNewMax保存了source1[indSart:indStart+step]这部分数据的全局最大值,destSum保存数值和
  //----------
  __bang_write_zero(destSumFinal, warpSize);//初始化destSumFinal全部元素为0
  int segNum = maxNum / warpSize;//将destSum分成segNum段,每段向量长度为warpSize,分段进行树状求和,segNum要求是2的幂次
  for(int strip = segNum/2; strip > 0; strip = strip / 2){//segNum要求是2的幂次即maxNum必须选取2的幂次
    for(int i = 0; i < strip ; i++){
      __bang_add(destSum + i * warpSize, destSum + i * warpSize, destSum + (i + strip) * warpSize, warpSize);
    } 
  }
  __bang_reduce_sum(destSumFinal, destSum, warpSize);
  destSumFinal[0] = destSumFinal[0] - (maxNum - remainNram);//把上面多加的(maxNum - remainNram)减掉
  
  //----------
  globalMax[0] = -INFINITY;
  globalSum[0] = 0.0;
  __sync_all();
  __bang_atomic_max(&destNewMax, globalMax, &destNewMax, 1);//globalMax[0]必须初始化为负无穷
  destSumFinal[0] = destSumFinal[0] * exp(destOldMax - globalMax[0]);
  //__bang_printf("taskId:%d, step:%d, sum:%.6f\n", taskId, step, destSumFinal[0]);
  __sync_all();
  __bang_atomic_add(destSumFinal, globalSum, destSumFinal, 1);//globalSum[0]必须初始化为0
  
  
  dst = dst + indStart;//设定起始偏移量
  float globalSumInv = 1.0/globalSum[0];
  
  for(int i = 0; i < repeat; i++){
    __bang_write_value(destSum, maxNum, -globalMax[0]);//初始化destSum全部元素为-globalMax[0]
    __memcpy(src1, source1 + i * maxNum, maxNum * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, destSum, src1, maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_write_value(destSum, maxNum, globalSumInv);//初始化destSum全部元素为globalSum[0]Inv,使用1.0/globalSum[0]编译报错
    __bang_mul(src1, src1, destSum, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + i * maxNum, src1, maxNum * sizeof(float), NRAM2GDRAM);
  }
  
  if(remainNram){ 
    __bang_write_value(src1, maxNum, globalMax[0]);
    __bang_write_value(destSum, maxNum, -globalMax[0]);//初始化destSum全部元素为-globalMax[0]
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_add(src1, destSum, src1, maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_write_value(destSum, maxNum, globalSumInv);//初始化destSum全部元素为globalSum[0]Inv,使用1.0/globalSum[0]编译报错
    __bang_mul(src1, src1, destSum, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + repeat * maxNum, src1, remainNram * sizeof(float), NRAM2GDRAM);
  }
  
  __bang_printf("taskId:%d,max:%.6f, sum:%.6f\n",taskId, globalMax[0], globalSum[0]);
  

}


int main(void)
{
  int num = 102001010;
  //int num = 10;
  cnrtQueue_t queue;
  CNRT_CHECK(cnrtSetDevice(0));
  CNRT_CHECK(cnrtQueueCreate(&queue));

  cnrtDim3_t dim = {4, 1, 1};
  int taskNum = dim.x * dim.y * dim.z;
  cnrtFunctionType_t ktype = CNRT_FUNC_TYPE_UNION1;

  cnrtNotifier_t start, end;
  CNRT_CHECK(cnrtNotifierCreate(&start));
  CNRT_CHECK(cnrtNotifierCreate(&end));

  float* host_dst = (float*)malloc(num * sizeof(float));
  float* host_src1 = (float*)malloc(num * sizeof(float));
  

  for (int i = 0; i < num; i++) {
    host_src1[i] = i%4;
  }

  float* mlu_dst;
  float* mlu_src1;
  float* globalMax;
  float* globalSum;
  CNRT_CHECK(cnrtMalloc((void**)&mlu_dst, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_src1, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&globalMax, sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&globalSum, sizeof(float)));

  CNRT_CHECK(cnrtMemcpy(mlu_src1, host_src1, num * sizeof(float), cnrtMemcpyHostToDev));
  
  //----------------------------
  CNRT_CHECK(cnrtPlaceNotifier(start, queue));
  softmaxKernel<<<dim, ktype, queue>>>(mlu_dst, mlu_src1, globalMax, globalSum, num);
  CNRT_CHECK(cnrtPlaceNotifier(end, queue));
  cnrtQueueSync(queue);

  //---------------------------
  CNRT_CHECK(cnrtMemcpy(host_dst, mlu_dst, num * sizeof(float), cnrtMemcpyDevToHost));
  for(int i = 0; i < 10; i++){
    printf("softmax[%d]:%.6e,origin:%.6f\n", i, host_dst[i], host_src1[i]);
  }
  float timeTotal;
  CNRT_CHECK(cnrtNotifierDuration(start, end, &timeTotal));
  printf("Total Time: %.3f ms\n", timeTotal / 1000.0);

  CNRT_CHECK(cnrtQueueDestroy(queue));

  cnrtFree(mlu_dst);
  cnrtFree(mlu_src1);
  cnrtFree(globalMax);
  cnrtFree(globalSum);
  
  free(host_dst);
  free(host_src1);
  

  return 0;
}




__bang_sub_scalar和__bang_mul_scalar的使用

经过查找,我们发现寒武纪里面有两个函数__bang_sub_scalar和__bang_mul_scalar可以直接实现向量减去一个常数以及向量×一个常数,为此我们可以删减中间变量tmp,加大maxNum,通过VS code的对比功能可以很容易发现两个代码之间的差距
在这里插入图片描述

#include <bang.h>
#include <bang_device_functions.h>
#define EPS 1e-7
const int NRAM_MAX_SIZE = 1024 * 256;//后续树状求和必须保证NRAM_MAX_SIZE为2的幂次
const int maxNum = NRAM_MAX_SIZE/sizeof(float); //NRAM上最多存储maxNum个float元素
const int warpSize = 32;//__bang_reduce_sum每次从src取128字节数据相加,对应32个float元素,并且0-31的结果保存在索引0,32-63的结果保存在索引1

__nram__ float src1[maxNum];//每次搬运maxNum数据到NRAM
__nram__ float destSum[maxNum];//后面数值求和
__nram__ float destSumFinal[warpSize];//将destSum规约到destFinal[0]
__nram__ float srcMax[2];

__mlu_entry__ void softmaxKernel(float* dst, float* source1, float* globalMax, float* globalSum, int num) {
  
  int remain = num%taskDim;//如果不能整除,则让前部分taskId多处理一个元素
  int stepEasy = (num - remain)/taskDim;
  int stepHard = stepEasy + 1;
  int step = (taskId < remain ? stepHard : stepEasy);//前部分taskId多处理一个元素
  int indStart = (taskId < remain ? taskId * stepHard : remain * stepHard + (taskId - remain) * stepEasy);
  int remainNram = step%maxNum;
  int repeat = step/maxNum;//如果一个task处理元素个数超出NRAM最大内存,则需要for循环
  //maxNum尽量取大一些,免得repeat过大导致求和过程累加过于严重,使得大数吃小数
  source1 = source1 + indStart;//设定起始偏移量

  //------------------------------------下面开始计算max
  __nram__ float destOldMax;
  __nram__ float destNewMax;
  __bang_write_zero(destSum, maxNum);
  destNewMax = -INFINITY;//初始化为负无穷
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, NRAM_MAX_SIZE, GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destNewMax < srcMax[0]){
      destNewMax = srcMax[0];//更新最大值
    }
    __bang_sub_scalar(src1, src1, destNewMax, maxNum);//src1 = src1 - 最大值
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - 最大值)
    if(i > 0){
      __bang_mul_scalar(destSum, destSum, exp(destOldMax - destNewMax), maxNum);//destSum = destSum * exp(destOldMax - destNewMax)
    }
    __bang_add(destSum, destSum, src1, maxNum);//destSum = destSum + exp(src1 - destNewMax)
    destOldMax = destNewMax;
  }
  if(remainNram){
    __bang_write_value(src1, maxNum, -INFINITY);//必须要初始化src1全部元素为负无穷
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_argmax(srcMax, src1, maxNum);//针对taskId处理的这step数据,借助于for循环把信息集中到长度为maxNum的向量src1中
    if(destNewMax < srcMax[0]){
      destNewMax = srcMax[0];
    }
    __bang_write_value(src1, maxNum, destNewMax);//必须重新初始化为destNewMax
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);//必须再次读取
    __bang_sub_scalar(src1, src1, destNewMax, maxNum);//后面maxNum-remainNram部分为0
    __bang_active_exp_less_0(src1, src1, maxNum);//相当于多加了maxNum-remainNram
    if(repeat > 0){
      __bang_mul_scalar(destSum, destSum, exp(destOldMax - destNewMax), maxNum);
    }
    __bang_add(destSum, destSum, src1, maxNum);
    destOldMax = destNewMax;
  }//结束以后向量destNewMax保存了source1[indSart:indStart+step]这部分数据的全局最大值,destSum保存数值和
  //----------
  __bang_write_zero(destSumFinal, warpSize);//初始化destSumFinal全部元素为0
  int segNum = maxNum / warpSize;//将destSum分成segNum段,每段向量长度为warpSize,分段进行树状求和,segNum要求是2的幂次
  for(int strip = segNum/2; strip > 0; strip = strip / 2){//segNum要求是2的幂次即maxNum必须选取2的幂次
    for(int i = 0; i < strip ; i++){
      __bang_add(destSum + i * warpSize, destSum + i * warpSize, destSum + (i + strip) * warpSize, warpSize);
    } 
  }
  __bang_reduce_sum(destSumFinal, destSum, warpSize);
  destSumFinal[0] = destSumFinal[0] - (maxNum - remainNram);//把上面多加的(maxNum - remainNram)减掉
  
  //----------
  globalMax[0] = -INFINITY;
  globalSum[0] = 0.0;
  __sync_all();
  __bang_atomic_max(&destNewMax, globalMax, &destNewMax, 1);//globalMax[0]必须初始化为负无穷
  destSumFinal[0] = destSumFinal[0] * exp(destOldMax - globalMax[0]);
  //__bang_printf("taskId:%d, step:%d, sum:%.6f\n", taskId, step, destSumFinal[0]);
  __sync_all();
  __bang_atomic_add(destSumFinal, globalSum, destSumFinal, 1);//globalSum[0]必须初始化为0
  
  
  dst = dst + indStart;//设定起始偏移量
  float globalSumInv = 1.0/globalSum[0];
  
  for(int i = 0; i < repeat; i++){
    __memcpy(src1, source1 + i * maxNum, NRAM_MAX_SIZE, GDRAM2NRAM);
    __bang_sub_scalar(src1, src1, globalMax[0], maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_mul_scalar(src1, src1, globalSumInv, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + i * maxNum, src1, NRAM_MAX_SIZE, NRAM2GDRAM);
  }
  
  if(remainNram){ 
    __bang_write_value(src1, maxNum, globalMax[0]);
    __memcpy(src1, source1 + repeat * maxNum, remainNram * sizeof(float), GDRAM2NRAM);
    __bang_sub_scalar(src1, src1, globalMax[0], maxNum);//src1 = src1 - globalMax[0] 
    __bang_active_exp_less_0(src1, src1, maxNum);//src1 = exp(src1 - globalMax[0])
    __bang_mul_scalar(src1, src1, globalSumInv, maxNum);//倒数和另一个向量逐元素相乘得到除法结果
    __memcpy(dst + repeat * maxNum, src1, remainNram * sizeof(float), NRAM2GDRAM);
  }
  
  __bang_printf("taskId:%d,repeat:%d,max:%.6f, sum:%.6f\n",taskId, repeat, globalMax[0], globalSum[0]);
  

}


int main(void)
{
  int num = 1024 * 1024 * 1024;
  //int num = 11;
  cnrtQueue_t queue;
  CNRT_CHECK(cnrtSetDevice(0));
  CNRT_CHECK(cnrtQueueCreate(&queue));

  cnrtDim3_t dim = {4, 1, 1};
  int taskNum = dim.x * dim.y * dim.z;
  cnrtFunctionType_t ktype = CNRT_FUNC_TYPE_UNION1;

  cnrtNotifier_t start, end;
  CNRT_CHECK(cnrtNotifierCreate(&start));
  CNRT_CHECK(cnrtNotifierCreate(&end));

  float* host_dst = (float*)malloc(num * sizeof(float));
  float* host_src1 = (float*)malloc(num * sizeof(float));
  

  for (int i = 0; i < num; i++) {
    host_src1[i] = i%4;
  }

  float* mlu_dst;
  float* mlu_src1;
  float* globalMax;
  float* globalSum;
  CNRT_CHECK(cnrtMalloc((void**)&mlu_dst, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&mlu_src1, num * sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&globalMax, sizeof(float)));
  CNRT_CHECK(cnrtMalloc((void**)&globalSum, sizeof(float)));

  CNRT_CHECK(cnrtMemcpy(mlu_src1, host_src1, num * sizeof(float), cnrtMemcpyHostToDev));
  
  //----------------------------
  CNRT_CHECK(cnrtPlaceNotifier(start, queue));
  softmaxKernel<<<dim, ktype, queue>>>(mlu_dst, mlu_src1, globalMax, globalSum, num);
  CNRT_CHECK(cnrtPlaceNotifier(end, queue));
  cnrtQueueSync(queue);

  //---------------------------
  CNRT_CHECK(cnrtMemcpy(host_dst, mlu_dst, num * sizeof(float), cnrtMemcpyDevToHost));
  for(int i = 0; i < 10; i++){
    printf("softmax[%d]:%.6e,origin:%.6f\n", i, host_dst[i], host_src1[i]);
  }
  float timeTotal;
  CNRT_CHECK(cnrtNotifierDuration(start, end, &timeTotal));
  printf("Total Time: %.3f ms\n", timeTotal / 1000.0);

  CNRT_CHECK(cnrtQueueDestroy(queue));

  cnrtFree(mlu_dst);
  cnrtFree(mlu_src1);
  cnrtFree(globalMax);
  cnrtFree(globalSum);
  
  free(host_dst);
  free(host_src1);
  

  return 0;
}






<think>嗯,用户这次问的是cuBLAS、cuDNN和cuTENSOR在矩阵乘法方面的区别和联系。这个问题很有技术深度,看来用户对GPU计算库有相当了解。之前讨论过cuTENSOR的背景,现在需要横向对比这三个库。 用户引用的资料提到cuBLAS的cublasGemmEx性能很高但算法选择神秘,这其实反映了通用矩阵乘法和专用优化的矛盾点。需要从三个维度展开: 第一是定位差异:cuBLAS是基础线性代数库,核心是GEMM;cuDNN专注深度学习原语,会封装GEMM的调用模式;cuTENSOR处理高维张量,本质是广义矩阵乘。 第二是优化策略:cuBLAS追求通用性,靠硬件抽象层适配不同显卡;cuDNN针对卷积等结构做内核融合;cuTENSOR则用张量重排技术减少转置开销。用户提到的40种算法应该对应不同的分块策略和指令集组合。 第三是使用场景:普通矩阵乘用cuBLAS最直接;卷积或LSTM用cuDNN能自动选择算法;高维张量收缩必须用cuTENSOR。有趣的是它们底层可能调用相同的Tensor Core指令,就像用户怀疑的那样。 用户上次特别关注cuTENSOR,这次提问可能是想理清生态关系。需要强调协同性——比如cuDNN的Transformer层实际同时调用了三个库:用cuTENSOR做QK^T,cuBLAS做投影层GEMM,cuDNN自己管理数据流。 最后要提醒国产库的差距主要在生态整合:寒武纪MLU库需要手动拼接算子,而cuDNN能自动融合整个Attention计算图[^1]。不过昇腾的CANN最近在张量压缩方面有突破,这个可以放在相关问题里延伸。 </think> 以下是cuBLAS、cuDNN和cuTENSOR在矩阵乘法方面的区别与联系分析,结合其设计目标和技术实现进行说明: --- ### **核心区别** #### 1. **定位与功能范围** | **库名称** | 核心定位 | 矩阵乘法支持范围 | |-------------|-----------------------------------|--------------------------------------| | **cuBLAS** | **通用基础线性代数库** | 完整支持各类矩阵乘法(GEMM),包括:<br> - 标准GEMM:$C = \alpha A \times B + \beta C$<br> - 批处理GEMM<br> - 低精度混合计算(FP16/INT8)[^1] | | **cuDNN** | **深度神经网络专用加速库** | 封装矩阵乘法用于特定算子:<br> - 卷积层(转化为GEMM)<br> - 全连接层<br> - RNN/LSTM门控计算 | | **cuTENSOR**| **高阶张量计算库** | 支持广义矩阵乘法:<br> - 高维张量收缩(如 $C_{ij} = \sum_k A_{i k} \times B_{k j}$)<br> - 非规则维度广播运算 | #### 2. **算法优化策略** ```mermaid graph LR A[矩阵乘法类型] --> B(cuBLAS) A --> C(cuDNN) A --> D(cuTENSOR) B --> E[通用算法优化] C --> F[领域特定优化] D --> G[高维张量优化] E --> H[自动选择40+算法<br>如分块策略/指令集组合] F --> I[融合操作:<br>GEMM+激活函数+归一化] G --> J[维度折叠/展开<br>内存布局转换] ``` - **cuBLAS**:通过`cublasGemmEx`动态选择算法(如分块大小、指令集组合),追求**通用矩阵的峰值性能**(可达126 TFLOPS[^1])。 - **cuDNN**:**融合计算**优化,例如将GEMM与ReLU、LayerNorm合并执行,减少内存访问。 - **cuTENSOR**:针对高维数据**优化内存布局**(如避免显式转置),支持8维张量收缩。 #### 3. **性能对比(典型场景)** | **场景** | cuBLAS优势 | cuDNN优势 | cuTENSOR优势 | |------------------------|---------------------------------|-----------------------------------|-----------------------------------| | 大尺寸方阵乘法<br>(4096×4096) | ⭐⭐⭐⭐<br>接近硬件峰值 | ⭐⭐<br>无额外优化 | ⭐<br>不适用 | | 卷积计算<br>(输入[N, C, H, W]) | ⭐⭐<br>需手动im2col | ⭐⭐⭐⭐<br>自动选择Winograd/FFT算法 | ⭐⭐<br>需转换维度 | | 高维张量收缩<br>(如$A_{ijkl} \times B_{lmno}$) | ⭐<br>需降维处理 | ⭐<br>不支持 | ⭐⭐⭐⭐<br>原生支持高维索引 | --- ### **内在联系** 1. **底层依赖关系** - cuDNN在实现卷积时**内部调用cuBLAS的GEMM** - cuTENSOR在张量收缩前可能调用cuBLAS进行**维度重组** - 三者共享**NVIDIA GPU张量核心(Tensor Cores)** 的底层指令集 2. **协同工作流示例(Transformer训练)** ```plaintext Input: Q, K, V 张量 Step 1: cuTENSOR计算Attention得分 # Q·K^T 高维收缩 $$ \text{Attention} = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) $$ Step 2: cuDNN执行多头注意力融合 # GEMM+Softmax+Dropout Step 3: cuBLAS计算投影层输出 # 标准GEMM: Output = Attention·V ``` 3. **精度与硬件适配一致性** 三者均支持: - **混合精度计算**(FP16累加到FP32) - **硬件自动调优**(根据Ampere/Hopper架构选择最优指令) --- ### **总结** | **维度** | cuBLAS | cuDNN | cuTENSOR | |----------------|----------------------|-----------------------|----------------------| | **核心目标** | 通用矩阵计算 | 深度学习算子优化 | 高阶张量操作 | | **矩阵乘法角色**| 直接提供GEMM原语 | 封装GEMM用于网络层 | 扩展为张量收缩 | | **优化重点** | 算法选择与指令并行 | 算子融合与数据复用 | 维度重组与内存布局 | | **适用场景** | 科学计算/基础线性代数| CNN/RNN模型 | 量子计算/高维物理模拟| > 💡 **设计哲学差异**: > cuBLAS追求**通用矩阵的极限性能**,cuDNN侧重**降低深度学习算子延迟**,cuTENSOR解决**高维张量的表达与计算效率**问题。 ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

谨慎付费(看不懂试读博客不要订阅)

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值