• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2023 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #ifndef NNACL_KERNEL_GATHER_H_
17 #define NNACL_KERNEL_GATHER_H_
18 
19 #include "nnacl/op_base.h"
20 #include "nnacl/tensor_c.h"
21 #include "nnacl/kernel.h"
22 
23 #define GATHER_BLOCK_INFOS_SIZE 32
24 
25 typedef struct GatherBlockBoundaryInfo {
26   int64_t begin_batch_;
27   int64_t begin_index_;
28   int64_t end_batch_;
29   int64_t end_index_;
30 } GatherBlockBoundaryInfo;
31 
32 typedef struct GatherStruct {
33   KernelBase base_;
34   int axis_;
35   int limit_;
36   int outer_size_;
37   int indices_size_;
38   int byte_inner_size_;
39   int block_infos_size_;
40   int *indices_data_;
41   GatherBlockBoundaryInfo block_infos_[GATHER_BLOCK_INFOS_SIZE];
42 } GatherStruct;
43 
44 KernelBase *CreateGather(OpParameter *param, int data_type);
45 
46 #endif  // NNACL_KERNEL_GATHER_H_
47