concat.cc 4.1 KB
Newer Older
C
chengduoZH 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2018 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. */

#include "paddle/fluid/operators/math/concat.h"
16
#include <vector>
C
chengduoZH 已提交
17 18 19 20 21 22

namespace paddle {
namespace operators {
namespace math {

/*
23
 * All tensors' dimension should be the same and the values of
24
 * each dimension must be the same, except the axis dimension.
C
chengduoZH 已提交
25 26 27 28 29
 */
template <typename T>
class ConcatFunctor<platform::CPUDeviceContext, T> {
 public:
  void operator()(const platform::CPUDeviceContext& context,
C
chengduoZH 已提交
30
                  const std::vector<framework::Tensor>& input, const int axis,
C
chengduoZH 已提交
31
                  framework::Tensor* output) {
32
    // TODO(zcd): Add input data validity checking
C
chengduoZH 已提交
33 34 35 36 37 38 39 40 41
    int num = input.size();

    int rows = 1;
    auto dim_0 = input[0].dims();
    for (int i = 0; i < axis; ++i) {
      rows *= dim_0[i];
    }
    int out_rows = rows, out_cols = 0;

C
chengduoZH 已提交
42
    std::vector<int64_t> input_cols(input.size());
C
chengduoZH 已提交
43 44 45
    for (int i = 0; i < num; ++i) {
      int t_cols = input[i].numel() / rows;
      out_cols += t_cols;
C
chengduoZH 已提交
46
      input_cols[i] = t_cols;
C
chengduoZH 已提交
47
    }
X
xuwei06 已提交
48
    auto cpu_place = boost::get<platform::CPUPlace>(context.GetPlace());
C
chengduoZH 已提交
49

C
chengduoZH 已提交
50
    // computation
C
chengduoZH 已提交
51
    for (int k = 0; k < out_rows; ++k) {
C
chengduoZH 已提交
52 53 54
      T* dst_ptr = output->data<T>() + k * out_cols;
      int col_idx = 0;
      for (int j = 0; j < num; ++j) {
C
chengduoZH 已提交
55
        int col_len = input_cols[j];
C
chengduoZH 已提交
56 57 58 59 60 61
        const T* src_prt = input[j].data<T>() + k * col_len;
        memory::Copy(cpu_place, dst_ptr + col_idx, cpu_place, src_prt,
                     sizeof(T) * col_len);
        col_idx += col_len;
      }
    }
C
chengduoZH 已提交
62 63 64
  }
};

65 66
/*
 * All tensors' dimension should be the same and the values of
67
 * each dimension must be the same, except the axis dimension.
68
 */
C
chengduoZH 已提交
69 70 71 72
template <typename T>
class ConcatGradFunctor<platform::CPUDeviceContext, T> {
 public:
  void operator()(const platform::CPUDeviceContext& context,
73
                  const framework::Tensor& input,
74
                  const std::vector<const framework::LoDTensor*>& ref_inputs,
75
                  const int axis, std::vector<framework::Tensor*>* outputs) {
76
    // TODO(zcd): Add input data validity checking
77
    size_t num = outputs->size();
C
chengduoZH 已提交
78

C
chengduoZH 已提交
79
    int input_rows = 1;
80
    auto dim_0 = ref_inputs[0]->dims();
C
chengduoZH 已提交
81 82 83
    for (int i = 0; i < axis; ++i) {
      input_rows *= dim_0[i];
    }
84

C
chengduoZH 已提交
85 86
    int input_cols = 0;

87
    std::vector<int64_t> output_cols(outputs->size());
88 89
    for (size_t i = 0; i < num; ++i) {
      int t_cols = ref_inputs[i]->numel() / input_rows;
C
chengduoZH 已提交
90 91 92
      input_cols += t_cols;
      output_cols[i] = t_cols;
    }
X
xuwei06 已提交
93
    auto cpu_place = boost::get<platform::CPUPlace>(context.GetPlace());
C
chengduoZH 已提交
94 95

    // computation
C
chengduoZH 已提交
96
    for (int k = 0; k < input_rows; ++k) {
C
chengduoZH 已提交
97 98
      const T* src_ptr = input.data<T>() + k * input_cols;
      int col_idx = 0;
C
chengduoZH 已提交
99
      for (size_t j = 0; j < num; ++j) {
C
chengduoZH 已提交
100
        int col_len = output_cols[j];
101
        auto* out_tensor = outputs->at(j);
102 103 104 105 106
        if (out_tensor != nullptr) {
          T* dst_ptr = out_tensor->data<T>() + k * col_len;
          memory::Copy(cpu_place, dst_ptr, cpu_place, src_ptr + col_idx,
                       sizeof(T) * col_len);
        }
C
chengduoZH 已提交
107 108 109
        col_idx += col_len;
      }
    }
C
chengduoZH 已提交
110 111 112 113 114 115 116 117
  }
};

template class ConcatFunctor<platform::CPUDeviceContext, int>;
template class ConcatFunctor<platform::CPUDeviceContext, int64_t>;
template class ConcatFunctor<platform::CPUDeviceContext, float>;
template class ConcatFunctor<platform::CPUDeviceContext, double>;

C
chengduoZH 已提交
118 119 120 121 122
template class ConcatGradFunctor<platform::CPUDeviceContext, int>;
template class ConcatGradFunctor<platform::CPUDeviceContext, int64_t>;
template class ConcatGradFunctor<platform::CPUDeviceContext, float>;
template class ConcatGradFunctor<platform::CPUDeviceContext, double>;

C
chengduoZH 已提交
123 124 125
}  // namespace math
}  // namespace operators
}  // namespace paddle