lod_tensor.h 3.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

   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 <memory>
18
#ifdef PADDLE_WITH_CUDA
19 20
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
21
#include <thrust/system/cuda/experimental/pinned_allocator.h>
22 23
#endif

24
#include <glog/logging.h>
25 26 27 28 29 30 31
#include "paddle/framework/ddim.h"
#include "paddle/framework/tensor.h"
#include "paddle/platform/enforce.h"

namespace paddle {
namespace framework {

32
#ifndef PADDLE_WITH_CUDA
Q
qijun 已提交
33 34 35 36
template <typename T>
using Vector = std::vector<T>;
#else
template <typename T>
37 38
using Vector = thrust::host_vector<
    T, thrust::system::cuda::experimental::pinned_allocator<T>>;
Q
qijun 已提交
39 40
#endif

41 42 43 44 45 46 47 48 49 50 51 52
/*
 * 3-level LoD stores
 *
 * 0 10 20
 * 0 5 10 15 20
 * 0 2 5 7 10 12 15 20
 *
 * - in a level, each element indicates offset in the underlying Tensor
 * - the first element should be 0 and that indicates that this sequence start
 * from 0
 * - each sequence's begin and end(no-inclusive) is level[id, id+1]
 */
53
using LoD = std::vector<Vector<size_t>>;
Q
qijun 已提交
54

55
LoD SliceLevels(const LoD& in, size_t level_begin, size_t level_end);
Q
qijun 已提交
56

57
LoD SliceInLevel(const LoD& in, size_t level, size_t elem_begin,
Q
qijun 已提交
58 59
                 size_t elem_end);

60
bool operator==(const LoD& a, const LoD& b);
Q
qijun 已提交
61

62
/*
63
 * LoDTensor (Level of details Tensor)
64 65
 * see https://en.wikipedia.org/wiki/Level_of_details for reference.
 */
66
class LoDTensor : public Tensor {
67
 public:
68
  LoDTensor() {}
69

70
  explicit LoDTensor(const LoD& lod) : lod_(lod) {}
Q
qijun 已提交
71

72
  void set_lod(const LoD& lod) { lod_ = lod; }
Q
qijun 已提交
73

74
  LoD lod() const { return lod_; }
Q
qijun 已提交
75

76
  /*
W
wanghaoshuang 已提交
77
   * Get the start offset and end offset of an  element from LoD.
78
   */
W
wanghaoshuang 已提交
79
  std::pair<size_t, size_t> lod_element(size_t level, size_t elem) const {
80
    PADDLE_ENFORCE_LT(level, NumLevels());
W
wanghaoshuang 已提交
81 82
    PADDLE_ENFORCE_LT(elem, NumElements(level));
    return std::make_pair((lod_)[level][elem], (lod_)[level][elem + 1]);
83 84 85
  }

  /*
86
   * Number of LoDTensor's levels, each level has units of data, for example,
87 88
   * in the sentence's view, article, paragraph, sentence are 3 levels.
   */
89
  size_t NumLevels() const { return lod_.size(); }
90 91 92 93
  /*
   * Number of elements in a level.
   */
  size_t NumElements(size_t level = 0) const {
94
    PADDLE_ENFORCE_LT(level, NumLevels());
95
    // the last offset is the end of last element
Q
qijun 已提交
96
    return (lod_)[level].size() - 1;
97 98
  }

99 100 101 102 103 104 105 106 107 108 109 110
  /*
   * Number of lower-level elements.
   * For example, a 2-level lod-tensor
   *
   * 0-th level   |   |
   * 1-th level   ||  |||
   *
   * NumElements(0, 0) get 2
   * NumElements(0, 1) get 3
   */
  size_t NumElements(size_t level, size_t idx) const;

111
  /*
112
   * Shrink levels[level_begin:level_end]
113
   */
114
  void ShrinkLevels(size_t level_begin, size_t level_end);
115 116

  /*
117
   * Shrink elements of a level, [elem_begin: elem_end]
118
   * @note: low performance in slice lod_.
119
   */
120
  void ShrinkInLevel(size_t level, size_t elem_begin, size_t elem_end);
121

Q
qijun 已提交
122
 private:
123
  LoD lod_;
124 125 126
};
}  // namespace framework
}  // namespace paddle