forked from PaddlePaddle/Paddle
-
Notifications
You must be signed in to change notification settings - Fork 0
/
concat_and_split.h
84 lines (76 loc) · 2.67 KB
/
concat_and_split.h
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <iostream>
#include <vector>
#include "paddle/extension.h"
int64_t GetRows(std::vector<int64_t> shape, int64_t axis) {
int64_t rows = 1;
for (int64_t i = 0; i < axis; ++i) {
rows *= shape[i];
}
return rows;
}
std::vector<int64_t> GetCols(const std::vector<paddle::Tensor>& ins,
int64_t rows,
int64_t* cols) {
std::vector<int64_t> cols_vec(ins.size());
for (size_t i = 0; i < ins.size(); ++i) {
int64_t t_cols = ins[i].size() / rows;
*cols += t_cols;
cols_vec[i] = t_cols;
}
return cols_vec;
}
template <typename data_t>
void ConcatCpuKernel(const std::vector<paddle::Tensor>& ins,
paddle::Tensor* out,
int64_t axis) {
size_t num = ins.size();
int64_t out_rows = GetRows(ins[0].shape(), axis);
int64_t out_cols = 0;
auto ins_cols = GetCols(ins, out_rows, &out_cols);
auto* out_data = out->mutable_data<data_t>();
int64_t col_idx = 0;
for (size_t i = 0; i < num; ++i) {
int64_t col_len = ins_cols[i];
auto* in_data = ins[i].data<data_t>();
for (int j = 0; j < out_rows; ++j) {
std::memcpy(out_data + j * out_cols + col_idx,
in_data + j * col_len,
sizeof(data_t) * col_len);
}
col_idx += col_len;
}
}
template <typename data_t>
void SplitCpuKernel(const paddle::Tensor& in,
const std::vector<paddle::Tensor>& ref_ins,
std::vector<paddle::Tensor>* outs,
int64_t axis) {
size_t num = outs->size();
int64_t in_rows = GetRows(ref_ins[0].shape(), axis);
int64_t in_cols = 0;
auto out_cols = GetCols(ref_ins, in_rows, &in_cols);
for (size_t i = 0; i < in_rows; ++i) {
auto* in_data = in.data<data_t>() + i * in_cols;
int64_t col_idx = 0;
for (size_t j = 0; j < num; ++j) {
int64_t col_len = out_cols[j];
auto* out_data = outs->at(j).mutable_data<data_t>() + i * col_len;
std::memcpy(out_data, in_data + col_idx, sizeof(data_t) * col_len);
col_idx += col_len;
}
}
}