切片数据搬运

函数功能

切片数据搬运，主要适用于非连续vector数据搬运。

函数原型

参数说明

表3 切片数据搬运接口参数说明

参数名称

输入/输出

含义

dstLocal, dstGlobal

输出

目的操作数，类型为LocalTensor或GlobalTensor。支持的数据类型为：half/int16_t/uint16_t/float/int32_t/uint32_t/int8_t/uint8_t。

srcLocal, srcGlobal

输入

源操作数，类型为LocalTensor或GlobalTensor。支持的数据类型为：half/int16_t/uint16_t/float/int32_t/uint32_t/int8_t/uint8_t。

srcSliceInfo/dstSliceInfo

输入

目的操作数/源操作数切片信息，SliceInfo类型，SliceInfo具体参数请参考表4。

dimValue

输入

操作数维度信息，默认值为1。
表4 SliceInfo结构参数定义

参数名称

含义

startIndex

切片的起始元素位置。

endIndex

切片的终止元素位置。

stride

切片的间隔元素个数。

burstLen

横向切片，每一片数据的长度，仅在维度为1时生效，超出1维的情况下，必须配置为1，不支持配置成其他值。单位为datablock（32B）。比如，srcSliceInfo的List为 {{16, 70, 7, 3, 87}, {0, 2, 1, 1, 3}}，{16, 70, 7, 3, 87}表示第一维的切片信息，burstLen设置为3，表示一个切片数据段大小为3个datablock； {0, 2, 1, 1, 3}为第二维的切片信息，burstLen仅能设置为1。

shapeValue

当前维度的原始长度。单位为元素个数。

通过具体的示例对上述参数进行解析，示意图如下：

图1 参数解析示意图

支持的型号

Atlas推理系列产品AI Core

Atlas A2训练系列产品/Atlas 800I A2推理产品

注意事项

调用示例

切片数据搬运，非连续转为连续
 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
#include "kernel_operator.h"
// 本样例中tensor数据类型为float
template <typename T>
class KernelDataCopySliceGM2UB {
public:
    __aicore__ inline KernelDataCopySliceGM2UB()
    {}
    __aicore__ inline void Init(__gm__ uint8_t* dstGm, __gm__ uint8_t* srcGm)
    {
        AscendC::SliceInfo srcSliceInfoIn[] = {{16, 70, 7, 3, 87}, {0, 2, 1, 1, 3}};// 如输入数据示例：startIndex为16，endIndex为70，burstLen为3，stride为7, shapeValue为87。
        AscendC::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()
    {
        AscendC::LocalTensor<T> srcLocal = inQueueSrcVecIn.AllocTensor<T>();
        AscendC::DataCopy(srcLocal, srcGlobal,  dstSliceInfo, srcSliceInfo, dimValue);
        inQueueSrcVecIn.EnQue(srcLocal);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<T> srcOutLocal = inQueueSrcVecIn.DeQue<T>();
        AscendC::DataCopy(dstGlobal, srcOutLocal, dstSliceInfo, dstSliceInfo, dimValue);
        inQueueSrcVecIn.FreeTensor(srcOutLocal);
    }
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrcVecIn;
    AscendC::GlobalTensor<T> srcGlobal;
    AscendC::GlobalTensor<T> dstGlobal;
    AscendC::SliceInfo dstSliceInfo[K_MAX_DIM];
    AscendC::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)
{
    KernelDataCopySliceGM2UB<TYPE> op;
    op.Init(dst_gm, src_gm);
    op.Process();
}

结果示例请参考图1。

父主题： DataCopy