GatherD

按照给定的维度 dim 和索引张量 index，从输入张量中按元素位置 抽取数据，生成新的输出张量。

该算子等价于在指定维度上执行逐元素 Gather 操作，其输出形状与 index 张量形状一致。

\[\text{output}[i_0, \dots, i_n] = \text{input}_x[i_0, \dots, i_{dim-1}, \text{index}[i_0, \dots, i_n], i_{dim+1}, \dots, i_n]\]
输入：

  • input_x - 输入张量的数据地址，大小为 input_shape 各维度之积 × sizeof(输入数据类型)。 数据类型需与所调用的 GatherD 接口类型一致。

  • dim - 指定进行 Gather 操作的维度索引，取值范围为 [0, index_shape_size)。

  • index - 索引张量的数据地址，类型为 int*， 大小为 index_shape 各维度之积 × sizeof(int)， 用于指定在 dim 维度上的取值位置。

  • input_shape - 输入张量各维度大小数组地址，大小为 input_shape_size × sizeof(int)。

  • input_shape_size - 输入张量的维度数量。

  • index_shape - 索引张量的形状数组地址， 其维度数量与 input_shape_size 相同，大小为 index_shape_size × sizeof(int)。

  • core_mask - 核掩码（仅共享存储版本使用）。

输出：

  • output - 输出张量的数据地址， 其形状与 index_shape 保持一致， 数据类型与 input_x 相同， 大小为 index_shape 各维度之积 × sizeof(输入数据类型)。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持int8, int16, int32, fp32, fp64, cplx64, cplx128

  • MT7004 支持fp16, fp32, int16, int32, cplx64

  • index 中的取值应满足 0 <= index[...] < input_shape[dim]。

  • 输出张量的元素个数等于 index_shape 各维度之积。

  • 该算子不对索引顺序进行任何排序或检查。

共享存储版本：

void fp_gather_d_s(float *input_x, int dim, int *index, float *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void dp_gather_d_s(double *input_x, int dim, int *index, double *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void hp_gather_d_s(float16 *input_x, int dim, int *index, float16 *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void i32_gather_d_s(int *input_x, int dim, int *index, int *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void i16_gather_d_s(int16_t *input_x, int dim, int *index, int16_t *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void i8_gather_d_s(int8_t *input_x, int dim, int *index, int8_t *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void c64_gather_d_s(float *input_x, int dim, int *index, float *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)
void c128_gather_d_s(double *input_x, int dim, int *index, double *output, int *input_shape, int *index_shape, int index_shape_size, int core_mask)

C调用示例：

 1// MT7004 示例
 2#include <stdio.h>
 3#include <gatherd.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input_x = (float *)0xA0000000;   // input_x 在 DDR 空间
 7    float *output  = (float *)0xB0000000;
 8    int *index   = (int *)0xA1000000;
 9
10    int input_shape[] = {4, 8, 16};
11    int index_shape[] = {4, 8, 16};
12    int index_shape_size = 3;
13    int dim = 1;
14    int core_mask = 0xff;
15
16    fp_gather_d_s(input_x, dim, index, output, input_shape, index_shape, index_shape_size, core_mask);
17    return 0;
18}

私有存储版本：

void fp_gather_d_p(float *input_x, int dim, int *index, float *output, int *input_shape, int *index_shape, int index_shape_size)
void dp_gather_d_p(double *input_x, int dim, int *index, double *output, int *input_shape, int *index_shape, int index_shape_size)
void hp_gather_d_p(float16 *input_x, int dim, int *index, float16 *output, int *input_shape, int *index_shape, int index_shape_size)
void i32_gather_d_p(int *input_x, int dim, int *index, int *output, int *input_shape, int *index_shape, int index_shape_size)
void i16_gather_d_p(int16_t *input_x, int dim, int *index, int16_t *output, int *input_shape, int *index_shape, int index_shape_size)
void i8_gather_d_p(int8_t *input_x, int dim, int *index, int8_t *output, int *input_shape, int *index_shape, int index_shape_size)
void c64_gather_d_p(float *input_x, int dim, int *index, float *output, int *input_shape, int *index_shape, int index_shape_size)
void c128_gather_d_p(double *input_x, int dim, int *index, double *output, int *input_shape, int *index_shape, int index_shape_size)

C调用示例：

 1// MT7004 示例
 2#include <stdio.h>
 3#include <gatherd.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input_x = (float *)0x10000000;   // input_x 在 L2 空间
 7    float *output  = (float *)0x10010000;
 8    int *index   = (int *)0x10020000;
 9
10    int input_shape[] = {2, 4};
11    int index_shape[] = {2, 4};
12    int index_shape_size = 2;
13    int dim = 0;
14
15    fp_gather_d_p(input_x, dim, index, output, input_shape, index_shape, index_shape_size);
16    return 0;
17}