EXTERN_IMPL_BUFPOOL宏

功能说明

开发者可以通过TBufPool类手动管理Unified Buffer、L1 Buffer物理内存。

TBufPool类切分的内存块都是连续的，开发者可能有一些自定义的内存块分配需求，比如不连续内存块、内存块在不同TQue之间共享等，这时就需要开发者自定义一个TBufPool的实现。

为了简化开发者的自定义实现，提供EXTERN_IMPL_BUFPOOL宏来辅助用户自定义TBufPool。使用自定义TBufPool功能时，需要注意：

EXTERN_IMPL_BUFPOOL宏内部定义的函数Reset、Init、GetBufHandle、SetCurAddr、GetCurAddr、SetCurBufSize、GetCurBufSize接口参见后续章节描述。使用该宏后，即可使用上述接口完成自定义TBufPool功能。

自定义TBufPool相关接口为试验接口，在后续版本中可能会调整或改进，不保证后续兼容性。请开发者在使用过程中关注后续版本更新。

函数原型

1
2
// 省略宏定义具体内容
#define EXTERN_IMPL_BUFPOOL(EXT_BUFPOOL, POSITION, BUFID_SIZE) ...

参数说明

表1 EXTERN_IMPL_BUFPOOL宏原型定义参数说明

参数名称

输入/输出

含义

EXT_BUFPOOL

输入

自定义TBufPool类名。

POSITION

输入

自定义TBufPool逻辑位置，可以为VECIN、VECOUT、VECCALC、A1、B1、C1。关于TPosition的具体介绍请参考TPosition。

BUFID_SIZE

输入

自定义TBufPool分配的Buffer块数量，建议不超过16。

支持的型号

Atlas 训练系列产品

Atlas 推理系列产品AI Core

Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件

约束说明

无

返回值

无

调用示例

如下示例中，为tbufPool0划分65536 * 3大小的内存，然后自定义MyBufPool的InitBuffer函数，实现TQue和Tbuf的内存分配。

  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
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
#include "kernel_operator.h"

class MyBufPool {
public:
    __aicore__ inline MyBufPool() {
        Init();
    }

    template<class T> 
    __aicore__ inline bool InitBuffer(T& que, uint8_t num, uint32_t len) {
        len = (len + 32 - 1) / 32 * 32; // 保证内存块长度32字节对齐 
        auto ptr = this->GetBufHandle(this->GetCurBufSize());
        auto curPoolAddr = this->GetCurAddr();

        // call internal func to initnitial bufhandle
        que.InitStartBufHandle(ptr, num, len);
        for (int32_t i = 0; i < num; i++) {
            que.InitBufHandle(this, i, ptr, curPoolAddr + i * len, len);
        }

        this->SetCurAddr(curPoolAddr + num * len);
        this->SetCurBufSize(this->GetCurBufSize() + num);

        return true;
    }

    template<AscendC::TPosition bufPos>
    __aicore__ inline bool InitBuffer(AscendC::TBuf<bufPos>& buf, uint32_t len) {
        len = (len + 32 - 1) / 32 * 32; // 保证内存块长度32字节对齐         
        auto ptr = this->GetBufHandle(this->GetCurBufSize());
        auto curPoolAddr = this->GetCurAddr();

        // call internal func to initnitial bufhandle
        buf.InitStartBufHandle(ptr, 1, len);
        buf.InitBufHandle(this, 0, ptr, curPoolAddr, len);

        this->SetCurAddr(curPoolAddr + len);
        this->SetCurBufSize(this->GetCurBufSize() + 1);
        return true;
    }
    EXTERN_IMPL_BUFPOOL(MyBufPool, AscendC::TPosition::VECCALC, 16);
};

class MyTBufPoolKernel {
public:
    __aicore__ inline MyTBufPoolKernel() {}
    __aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
    {
        src0Global.SetGlobalBuffer((__gm__ half*)src0Gm);
        src1Global.SetGlobalBuffer((__gm__ half*)src1Gm);
        dstGlobal.SetGlobalBuffer((__gm__ half*)dstGm);
        pipe.InitBufPool(tbufPool0, 65536 * 3);
        tbufPool0.InitBuffer(srcQue0, 1, 65536);
        tbufPool0.InitBuffer(srcBuf1, 65536);
        tbufPool0.InitBuffer(dstQue0, 1, 65536);
    }

    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
        tbufPool0.Reset();
        pipe.Reset();
    }

private:
    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<half> src0Local = srcQue0.AllocTensor<half>();
        AscendC::LocalTensor<half> src1Local = srcBuf1.Get<half>();
        AscendC::DataCopy(src0Local, src0Global, 32768);
        AscendC::DataCopy(src1Local, src1Global, 32768);
        srcQue0.EnQue(src0Local);
    }
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<half> src0Local = srcQue0.DeQue<half>();
        AscendC::LocalTensor<half> src1Local = srcBuf1.Get<half>();
        AscendC::LocalTensor<half> dstLocal = dstQue0.AllocTensor<half>();
        AscendC::Add(dstLocal, src0Local, src1Local, 32768);
        dstQue0.EnQue<half>(dstLocal);
        srcQue0.FreeTensor(src0Local);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<half> dstLocal = dstQue0.DeQue<half>();
        AscendC::DataCopy(dstGlobal, dstLocal, 32768);
        dstQue0.FreeTensor(dstLocal);
    }

private:
    AscendC::TPipe pipe;
    MyBufPool tbufPool0;
    AscendC::TBuf<AscendC::TPosition::VECIN> srcBuf1;
    AscendC::TQue<AscendC::TPosition::VECIN, 1> srcQue0;
    AscendC::TQue<AscendC::TPosition::VECOUT, 1> dstQue0;
    AscendC::GlobalTensor<half> src0Global, src1Global, dstGlobal;
};

extern "C" __global__ __aicore__ void mytbufpool_kernel(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
{
    MyTBufPoolKernel op;
    op.Init(src0Gm, src1Gm, dstGm);
    op.Process();
}
父主题： 自定义TBufPool