尾块Tiling

Tiling实现

算子的Tiling结构体定义如下：

1
2
3
4
5
6

struct AddCustomTilingData {
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;
};

Host侧Tiling实现的主要内容为计算以上四个成员变量。步骤如下：

1. 判断数据总长度totalLength是否满足32字节对齐，如不满足，则计算totalLength向上32字节对齐后的长度totalLengthAligned。

   1 2 3 4 5 6

   constexpr uint32_t BLOCK_SIZE = 32; // 为方便计算，这里根据数据类型定义变量alignNum作为对齐数 uint32_t alignNum = BLOCK_SIZE / dataTypeSize; // totalLength为数据总量 uint32_t totalLengthAligned = (totalLength % alignNum == 0)? totalLength : ((totalLength + alignNum - 1) / alignNum) * alignNum;

2. 判断totalLengthAligned是否能被使用的核数BlockDim均分，如果可以，则计算每个核上计算数据长度blockLength。

   1 2 3 4 5 6

   constexpr uint32_t BLOCK_DIM = 8; constexpr uint32_t UB_BLOCK_NUM = 20; // 此处为方便验证，使用UB_BLOCK_NUM作为Unified Buffer可用的Block数量，因此可得出可用UB空间的大小为UB_BLOCK_NUM * BLOCK_SIZE uint32_t blockLength, tileNum; if (totalLengthAligned % BLOCK_DIM == 0) { blockLength = totalLengthAligned / BLOCK_DIM; }

3. 计算tileNum。为了减少数据搬运开销，应尽量使用核内的Unified Buffer空间。基于每个核上的计算量以及可用Unified Buffer空间的大小，计算tileNum。

   1	tileNum = blockLength / alignNum / UB_BLOCK_NUM;

4. 根据计算出的tileNum，计算tileLength和lastTileLength。

   如果每个核的计算量能够被当前可用Unified Buffer空间均分，则按照无尾块场景处理。

   1 2 3 4 5

   if ((blockLength / alignNum) % UB_BLOCK_NUM == 0) { // 单核的计算量能被当前可用UB空间均分，仅有主块，无尾块 tileLength = UB_BLOCK_NUM * alignNum; lastTileLength = 0; }

   反之，按照尾块场景处理，尾块长度为单核计算数据长度 - tileNum * tileLength。

   1 2 3 4 5 6 7 8 9

   if (tileNum == 0) { // 单核需要计算的长度小于UB可用空间，按照仅有尾块处理 tileLength = 0; lastTileLength = ((blockLength + alignNum - 1) / alignNum) * alignNum; } else { // 同时有主块和尾块 tileLength = UB_BLOCK_NUM * alignNum; lastTileLength = blockLength - tileNum * tileLength; }

Host侧Tiling实现的代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

constexpr uint32_t BLOCK_SIZE = 32;
constexpr uint32_t BLOCK_DIM = 8;
constexpr uint32_t UB_BLOCK_NUM = 20;  // 此处为方便验证，使用UB_BLOCK_NUM作为UB可用的Block数量，因此可得出可用UB空间的大小为UB_BLOCK_NUM * BLOCK_SIZE
...

uint32_t alignNum = BLOCK_SIZE / dataTypeSize;  // 为方便计算，这里根据数据类型定义变量alignNum作为对齐数，dataTypeSize为运算数据的数据类型对应的字节数
// totalLength为数据总量
uint32_t totalLengthAligned = (totalLength % alignNum == 0)?
        totalLength : ((totalLength + alignNum - 1) / alignNum) * alignNum;
uint32_t blockLength, tileNum;
if (totalLengthAligned % BLOCK_DIM == 0) {
    blockLength = totalLengthAligned / BLOCK_DIM;
    tileNum = blockLength / alignNum / UB_BLOCK_NUM;

    if (tileNum == 0) {
        // 单核需要计算的长度小于UB可用空间，按照仅有尾块处理
        tileLength = 0;
        lastTileLength = ((blockLength + alignNum - 1) / alignNum) * alignNum;
    } else if ((blockLength / alignNum) % UB_BLOCK_NUM == 0) {
        // 单核的计算量能被当前可用UB空间均分，仅有主块，无尾块
        tileLength = UB_BLOCK_NUM * alignNum;
        lastTileLength = 0;
    } else {
        // 同时有主块和尾块
        tileLength = UB_BLOCK_NUM * alignNum;
        lastTileLength = blockLength - tileNum * tileLength;
    }
    ...
}

(1，1904)形状的输入数据计算后，tiling结构体内各个变量的值如下：

1
2
3
4
5
6

struct AddCustomTilingData {
    uint32_t blockLength = 238;      // 每个核计算238个half，8个核共计算1904个half
    uint32_t tileNum = 0;            // 可用的UB空间足够，为仅有尾块的场景
    uint32_t tileLength = 0;         // 没有主块，主块长度为0
    uint32_t lastTileLength = 240;   // 238个half未32B对齐，对齐到240个half搬运
};

算子类实现

与多核Tiling相比，在Init函数中通过Pipe内存管理对象为输入输出Queue分配内存时，取tileLength与lastTileLength中的最大值作为分配内存的长度。例如，当单核需要计算的长度小于UB可用空间时，按照仅有尾块处理，此时tileLength为0，而lastTileLength为数据块长度。因此，需要取两者中的较大值来分配内存。

1
2

uint32_t initBufferLength = AscendC::Std::max(this->tileLength, this->lastTileLength);
pipe.InitBuffer(inQueueX, 1, initBufferLength * sizeof(dataType));

由于尾块长度为lastTileLength，与主块数据块的长度不同，因此在CopyIn函数、Compute函数、CopyOut函数中传入本次循环待处理的数据块长度参数tileLength，即待处理的主块或尾块的数据长度。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

__aicore__ inline void Process()
{
    // 计算主块数据，对应数据块长度为tileLength
    for (uint32_t i = 0; i < this->tileNum; i++) {
        CopyIn(i, this->tileLength);
        Compute(i, this->tileLength);
        CopyOut(i, this->tileLength);
    }
    // 计算尾块数据，对应数据块长度为lastTileLength
    if (this->lastTileLength > 0) {
        CopyIn(this->tileNum, this->lastTileLength);
        Compute(this->tileNum, this->lastTileLength);
        CopyOut(this->tileNum, this->lastTileLength);
    }
}

1
2
3
4
5
6
7
8
9

__aicore__ inline void CopyIn(int32_t progress, uint32_t tileLength)
{
    AscendC::LocalTensor<dataType> xLocal = inQueueX.AllocTensor<dataType>();
    AscendC::LocalTensor<dataType> yLocal = inQueueY.AllocTensor<dataType>();
    AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], tileLength);
    AscendC::DataCopy(yLocal, yGm[progress * this->tileLength], tileLength);
    inQueueX.EnQue(xLocal);
    inQueueY.EnQue(yLocal);
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

__aicore__ inline void Compute(int32_t progress, uint32_t tileLength)
{
    AscendC::LocalTensor<dataType> xLocal = inQueueX.DeQue<dataType>();
    AscendC::LocalTensor<dataType> yLocal = inQueueY.DeQue<dataType>();
    AscendC::LocalTensor<dataType> zLocal = outQueueZ.AllocTensor<dataType>();
    AscendC::Add(zLocal, xLocal, yLocal, tileLength);
    outQueueZ.EnQue<dataType>(zLocal);
    inQueueX.FreeTensor(xLocal);
    inQueueY.FreeTensor(yLocal);
}

1
2
3
4
5
6

__aicore__ inline void CopyOut(int32_t progress, uint32_t tileLength)
{
    AscendC::LocalTensor<dataType> zLocal = outQueueZ.DeQue<dataType>();
    AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, tileLength);
    outQueueZ.FreeTensor(zLocal);
}