AI Engine Kernel and Graph Programming--知识分享4

使用AI引擎API的示例设计

矩阵乘法

        AI引擎API为基于向量的矩阵乘法提供了一个aie：：mmul类模板。多个中间矩阵相乘的结果被累加以给出最终结果。有关支持的矩阵乘法形状(M*K*N)和数据类型的更多详细信息，请参阅《AI Engine API用户指南》(UG1529)中的矩阵乘法。AIE：：MMUL操作MUL和MAC接受用于基于向量的矩阵乘法的行主格式数据。然后，对于aie：：mmul的Mac操作，按M*K或K*N排列数据。这种数据洗牌可以在PL或AI引擎中完成。本节给出了A(*)×B(*)矩阵乘法的一个例子。数据类型为int8 x int8。选择矩阵乘法形状4*16*8用于AIE：：MMUL运算。假定输入数据为行优先格式。将数据作为4×16矩阵和16×8矩阵输入到矩阵乘法核。在矩阵乘法核之前，对输入数据进行混洗。例如，在混洗之前，将矩阵A(*)与a0、a1、…一起存储在存储器中，A63，A64，…，按顺序排列为A4096。AIE：：MMUL运算将使用形状4*16。矩阵A被分割成大小为4*16的较小矩阵。对于较小的矩阵A00，对于AIE：：MMUL，应顺序获取A0到A15、A64到A79、A128到A143和A192到A207。因此，数据混洗的目的是将A0到A15、A64到A79、A128到A143和A192到A207放入矩阵乘法内核的连续存储中。这种数据洗牌如下图所示。

类似地，输出数据被混洗。下图显示了设计图。

向量化矩阵乘法

        下面是一个矩阵乘法的例子，矩阵大小为（64 * 64）x（64 * 64），数据类型为int8 x int8。矩阵乘法的形状是4*16*8。输入数据被重新整形以用于矩阵乘法。

const int SHIFT=10;
//For element mmul
const int M=4;
const int K=16;
const int N=8;
//Total matrix sizes
const int rowA=64;
const int colA=64;
const int colB=64;
//mmul numbers
const int num_rowA=rowA/M;
const int num_colA=colA/K;
const int num_colB=colB/N;
void matrix_mul(input_buffer<int8> & __restrict matA, input_buffer<int8> &
__restrict matB, output_buffer<int8> & __restrict matC){
using MMUL = aie::mmul<M, K, N, int8, int8>;
const int8* __restrict pA=(int8*)matA.data();
const int8* __restrict pB=(int8*)matB.data();
int8* __restrict pC=(int8*)matC.data();
//For profiling only
unsigned cycle_num[2];
aie::tile tile=aie::tile::current();
cycle_num[0]=tile.cycles();//cycle counter of the AI Engine tile
for (unsigned i = 0; i < num_rowA; i++) { //for output row number of
    element matrix
for (unsigned j = 0; j < num_colB; j++) { //for output col number ofelement matrix
    const int8 * __restrict pA1 = pA + ( i * num_colA + 0) * MMUL::size_A;
    const int8 * __restrict pB1 = pB + ( 0 * num_colB + j) * MMUL::size_B;
    aie::vector<int8, MMUL::size_A> A0 = aie::load_v<MMUL::size_A>(pA1);
    pA1 += MMUL::size_A;
    aie::vector<int8, MMUL::size_B> B0 = aie::load_v<MMUL::size_B>(pB1);
    pB1 += MMUL::size_B * num_colB;
    MMUL C00; C00.mul(A0, B0);
    for (unsigned k = 0; k < num_colA-1; k++) {
        A0 = aie::load_v<MMUL::size_A>(pA1); pA1 += MMUL::size_A;
        B0 = aie::load_v<MMUL::size_B>(pB1); pB1 += MMUL::size_B * num_colB;
        C00.mac(A0, B0);
    }
    aie::store_v(pC, C00.template to_vector<int8>(SHIFT)); pC +=
    MMUL::size_C;
    }
}
//For profiling only
cycle_num[1]=tile.cycles();//cycle counter of the AI Engine tile
printf("start=%d,end=%d,total=%d
\n",cycle_num[0],cycle_num[1],cycle_num[1]-cycle_num[0]);
}

        分析的结果表明，循环大约需要3254个周期。总共，这是对int 8 * int 8数据类型的64 *64*64=262144次乘法，即每个周期262144/3254 ~=80次int 8 * int 8 MAC运算。注意：根据特定编译器设置和所用工具的版本，准确的周期数可能会略有波动。然而，无论这些变化如何，本节中描述的分析技术仍然是相关的和适用的。

数据混洗内核

        由于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_buffer<int8> & __restrict matA, output_buffer<int8>& __restrict matAout){
const int sizeA=M*N;
auto pV=aie::begin_vector<16>((int8*)matA.data());
auto pOut=aie::begin_vector<sizeA>((int8*)matAout.data());
aie::vector<int8,sizeA> mm;
for(int i=0;i<rowA/M;i++){
    for(int j=0;j<colA/N;j++){
        for(int k=0;k<M;k++){
            mm.insert(k,*pV);
            pV=pV+4;
        }
        *pOut++=mm;
        pV=pV-15;
    }
        pV=pV+12;
    }
}

//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_buffer<int8> & __restrict matA, output_buffer<int8>
& __restrict matAout){
const int sizeA=M*N;
auto pV=aie::begin_vector<16>((int8*)matA.data());
auto pOut=aie::begin_vector<16>((int8*)matAout.data());
aie::vector<int8,16> sv1,sv2;
for(int i=0;i<rowA/M;i++){
    for(int j=0;j<colA/N/2;j++){
        for(int k=0;k<M/2;k++){
            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;
    }
}

以下是用于对矩阵C的数据进行混洗的代码示例，输入形状为4*8：

//element matrix size
const int M=4;
const int N=8;
//Total matrix sizes
const int rowA=64;
const int colA=64;
void shuffle_4x8(input_buffer<int8> & __restrict matA, output_buffer<int8>
& __restrict matAout){
const int sizeA=M*N;
auto pV=aie::begin_vector<sizeA>((int8*)matA.data());
auto pOut=aie::begin_vector<sizeA>((int8*)matAout.data());
aie::vector<int8,sizeA> mm1,mm2,mm3,mm4;
for(int i=0;i<rowA/M;i++){
    for(int j=0;j<colA/N/4;j++){
        mm1=*pV++;
        mm2=*pV++;
        mm3=*pV++;
        mm4=*pV++;
        auto mm12=aie::interleave_zip(mm1,mm2,8);
        auto mm34=aie::interleave_zip(mm3,mm4,8);
        auto mm1234_low=aie::interleave_zip(mm12.first,mm34.first,16);
                auto mm1234_high=aie::interleave_zip(mm12.second,mm34.second,16);
        *pOut=mm1234_low.first;
        pOut=pOut+2;
        *pOut=mm1234_low.second;
        pOut=pOut+2;
        *pOut=mm1234_high.first;
        pOut=pOut+2;
        *pOut=mm1234_high.second;
        pOut=pOut-5;
    }
    pOut=pOut+6;
    }
}

阅读更多精彩文章，请关注订阅号：威视锐科技