下载
中文
注册

通过K_MAX_SHAPE_DIM宏设置ShapeInfo维度为0,优化栈空间

【优先级】中

【描述】GlobalTensor和LocalTensor中通过ShapeInfo类型的成员变量来保存shape信息,可以通过SetShapeInfo、GetShapeInfo来进行设置或者获取,通常用于算子实现内部的shape信息保存和传递。默认情况下其维度值为8。在不使用上述ShapeInfo功能的情况下,不需要这些信息,可以通过K_MAX_SHAPE_DIM宏将其设置为0。经实测减小K_MAX_SHAPE_DIM值,可缩减栈空间,减少scalar指令和cache miss几率,提升算子性能。

...
#ifndef K_MAX_SHAPE_DIM
#define K_MAX_SHAPE_DIM 8
#endif
...
struct ShapeInfo {
public:
    ...
    uint32_t shape[K_MAX_SHAPE_DIM];
    uint32_t originalShape[K_MAX_SHAPE_DIM];
};

template <typename T> class GlobalTensor {
....
private:
    ShapeInfo shapeInfo_;
}
template <typename T> class LocalTensor {
....
private:
    ShapeInfo shapeInfo_;
}
...

【反例】

算子无需使用ShapeInfo,但未对ShapeInfo大小进行限制(使用默认值8),导致浪费K_MAX_SHAPE_DIM * sizeof(uint32_t) * 2 * 4字节的栈空间。2表示有shape和originalShape2个数组,4表示该样例中使用了GlobalTensor和LocalTensor共4个Tensor。
...
#include "kernel_operator.h" ...
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    ...
    GlocalTensor<T> dataIn;
    GlocalTensor<T> dataOut;
    LocalTensor<T> vecIn;
    LocalTensor<T> vecOut;
    ...
}
...

【正例】

算子无需使用ShapeInfo,设置 #define K_MAX_SHAPE_DIM 0,有效缩减了K_MAX_SHAPE_DIM * sizeof(uint32_t) * 2 * 4大小的栈空间。
#define K_MAX_SHAPE_DIM 0
...
#include "kernel_operator.h" //需注意定义K_MAX_SHAPE_DIM宏的位置须在包含Ascend C相关头文件之前
...
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    ...
    GlocalTensor<T> dataIn;
    GlocalTensor<T> dataOut;
    LocalTensor<T> vecIn;
    LocalTensor<T> vecOut;
    ...
}
...