切片数据搬运

#include "kernel_operator.h"
using namespace AscendC;
// 本样例中tensor数据类型为float
template <typename T>
class KernelDataCopySliceGM2UB {
public:
    __aicore__ inline KernelDataCopySliceGM2UB()
    {}
    __aicore__ inline void Init(__gm__ uint8_t* dstGm, __gm__ uint8_t* srcGm)
    {
        SliceInfo srcSliceInfoIn[] = {{16, 70, 7, 3, 87}, {0, 2, 1, 1, 3}};// 如输入数据示例：startIndex为16，endIndex为70，burstLen为3，stride为7, shapeValue为87。
        SliceInfo dstSliceInfoIn[] = {{0, 47, 0, 3, 48}, {0, 1, 0, 1, 2}};// UB空间相对紧张，建议设置stride为0。
        uint32_t dimValueIn = 2;
        uint32_t dstDataSize = 96;
        uint32_t srcDataSize = 261;
        dimValue = dimValueIn;

        for (uint32_t i = 0; i < dimValueIn; i++) {
            srcSliceInfo[i].startIndex = srcSliceInfoIn[i].startIndex;
            srcSliceInfo[i].endIndex = srcSliceInfoIn[i].endIndex;
            srcSliceInfo[i].stride = srcSliceInfoIn[i].stride;
            srcSliceInfo[i].burstLen = srcSliceInfoIn[i].burstLen;
            srcSliceInfo[i].shapeValue = srcSliceInfoIn[i].shapeValue;

            dstSliceInfo[i].startIndex = dstSliceInfoIn[i].startIndex;
            dstSliceInfo[i].endIndex = dstSliceInfoIn[i].endIndex;
            dstSliceInfo[i].stride = dstSliceInfoIn[i].stride;
            dstSliceInfo[i].burstLen = dstSliceInfoIn[i].burstLen;
            dstSliceInfo[i].shapeValue = dstSliceInfoIn[i].shapeValue;
        }
        srcGlobal.SetGlobalBuffer((__gm__ T *)srcGm);
        dstGlobal.SetGlobalBuffer((__gm__ T *)dstGm);

        pipe.InitBuffer(inQueueSrcVecIn, 1, dstDataSize * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {
        LocalTensor<T> srcLocal = inQueueSrcVecIn.AllocTensor<T>();
        DataCopy(srcLocal, srcGlobal,  dstSliceInfo, srcSliceInfo, dimValue);
        inQueueSrcVecIn.EnQue(srcLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<T> srcOutLocal = inQueueSrcVecIn.DeQue<T>();
        DataCopy(dstGlobal, srcOutLocal, dstSliceInfo, dstSliceInfo, dimValue);
        inQueueSrcVecIn.FreeTensor(srcOutLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, 1> inQueueSrcVecIn;
    GlobalTensor<T> srcGlobal;
    GlobalTensor<T> dstGlobal;
    SliceInfo dstSliceInfo[K_MAX_DIM];
    SliceInfo srcSliceInfo[K_MAX_DIM]; // K_MAX_DIM = 8
    uint32_t dimValue;
};

extern "C" __global__ __aicore__ void kernel_data_copy_slice_out2ub(__gm__ uint8_t* src_gm, __gm__ uint8_t* dst_gm)
{
    AscendC::KernelDataCopySliceGM2UB op;
    op.Init(dst_gm, src_gm);
    op.Process();
}