数据乱序内核 - 2022.1 简体中文

AI 引擎内核编码 最佳实践指南 (UG1079)

Document ID
UG1079
Release Date
2022-05-25
Version
2022.1 简体中文

由于对于矩阵乘法形状,aie::mmul 接受基于行的矢量数据,因此在具有原始数据的 PL 或 AI 引擎中可能需要数据乱序以保证性能。本节假定对于整个矩阵,原始数据是基于行的数据。它将数据顺序打乱,以与矩阵乘法中使用的形状 4*16*8 相匹配。

以下内核代码用于为矩阵 A 打乱数据顺序,目标大小为 4*16
//element matrix size
const int M=4;
const int N=16;

//Total matrix sizes
const int rowA=64;
const int colA=64;

void shuffle_4x16(input_window<int8> * __restrict matA, output_window<int8> * __restrict matAout){
  const int sizeA=M*N;
  auto pV=aie::begin_vector<16>((int8*)matA->ptr);
  auto pOut=aie::begin_vector<sizeA>((int8*)matAout->ptr);
  aie::vector<int8,sizeA> mm;

  for(int i=0;i<rowA/M;i++){//output row number of element matrix
    for(int j=0;j<colA/N;j++){//output col number of element matrix
      for(int k=0;k<M;k++){//generate 4*16 matrix
        mm.insert(k,*pV);
        pV=pV+4;
      }
      *pOut++=mm;
      pV=pV-15;
    }
    pV=pV+12;
  }
}
以下代码示例用于为矩阵 B 打乱数据顺序,目标大小为 16*8
//element matrix size
const int M=16;
const int N=8;

//Total matrix sizes
const int rowA=64;
const int colA=64;

void shuffle_16x8(input_window<int8> * __restrict matA, output_window<int8> * __restrict matAout){
  const int sizeA=M*N;
  auto pV=aie::begin_vector<16>((int8*)matA->ptr);
  auto pOut=aie::begin_vector<16>((int8*)matAout->ptr);

  aie::vector<int8,16> sv1,sv2;
  for(int i=0;i<rowA/M;i++){
    for(int j=0;j<colA/N/2;j++){//generate two 16*8 matrices an iteration
      for(int k=0;k<M/2;k++){//generate two rows of two 16*8 matrices an iteration
        sv1=*pV;
        pV=pV+4;
        sv2=*pV;
        pV=pV+4;
        auto mm=aie::interleave_zip(sv1,sv2,8);
        *pOut=mm.first;
        pOut+=8;
        *pOut=mm.second;
        pOut-=7;
      }
      pOut+=8;
      pV-=63;
    }
    pV+=60;
  }
}