TPU编程竞赛

《第二届TPU编程竞赛》赛题分析-ReduceSum

sophgo logo
算能SOPHGO
智算赋能数字世界

功能描述

reduce_sum算子的作用是将某些维度的所有元素求和。

输入格式:[N, C, H, W]的tensor, 以及需要进行reduce的维度dims

输出格式:对应dims的维度为1,如dims={1,2},输出为[N, 1, 1, W]

Host端

typedef struct {
    int N, C, H, W;
    int axis_list[4];
    int axis_num;
    unsigned long long output_addr;
    unsigned long long input_addr;
} __attribute__((packed)) param_t;

其中[N, C, H, W]为输入tensor的四个维度的大小,axis_list为需要进行reduce_sum操作的维度索引,axis_num为需要进行reduce_sum操作的维度个数。

Device端

以dims={3}为例,可以通过使用avgpool算子,kernel_size设为(1, W),stride设为(1,1),padding为(0,0),这样可以计算W维度所有元素的均值,之后将avgpool的结果乘以W,就可以得到W维度所有元素的和。

如果dims={2},可以进行类似的操作,将kernel_size设为(H, 1),其他相同。最后结果乘以H,可以得到H维度所有元素的和。

dims={0}和dims={1}的情况,各位参赛选手结合上述方法,以及文档中的基础算子进行完成。