array.h 2.7 KB
Newer Older
X
Xin Pan 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// 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.

#pragma once

#include <cstdint>
S
sneaxiy 已提交
18 19
#include "paddle/fluid/framework/unroll_array_ops.h"
#include "paddle/fluid/platform/enforce.h"
X
Xin Pan 已提交
20 21 22

namespace paddle {
namespace framework {
S
sneaxiy 已提交
23

X
Xin Pan 已提交
24 25 26
template <typename T, size_t N>
class Array {
 public:
S
sneaxiy 已提交
27
  static constexpr size_t kSize = N;
X
Xin Pan 已提交
28

S
sneaxiy 已提交
29 30 31 32 33
  HOSTDEVICE inline Array() = default;

  template <typename... Args>
  HOSTDEVICE inline explicit Array(const T &val, Args... args) {
    UnrollVarArgsAssign<T, N>::Run(data_, val, args...);
X
Xin Pan 已提交
34 35
  }

S
sneaxiy 已提交
36 37 38
  HOSTDEVICE inline void Fill(const T &val) {
    UnrollFillConstant<N>::Run(data_, val);
  }
X
Xin Pan 已提交
39

S
sneaxiy 已提交
40
  HOSTDEVICE inline const T *Get() const { return data_; }
X
Xin Pan 已提交
41

S
sneaxiy 已提交
42
  HOSTDEVICE inline T *GetMutable() { return data_; }
X
Xin Pan 已提交
43

S
sneaxiy 已提交
44 45 46 47 48
  HOSTDEVICE inline T &operator[](size_t index) { return data_[index]; }

  HOSTDEVICE inline const T &operator[](size_t index) const {
    return data_[index];
  }
X
Xin Pan 已提交
49 50 51

  HOSTDEVICE constexpr size_t size() const { return N; }

S
sneaxiy 已提交
52 53 54 55 56 57 58 59
  HOSTDEVICE inline bool operator==(const Array<T, N> &other) const {
    return UnrollCompare<N>::Run(data_, other.data_);
  }

  HOSTDEVICE inline bool operator!=(const Array<T, N> &other) const {
    return !(*this == other);
  }

X
Xin Pan 已提交
60 61 62 63
 private:
  T data_[N];
};

S
sneaxiy 已提交
64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100
template <typename T>
class Array<T, 0> {
 public:
  static constexpr size_t kSize = 0;

  HOSTDEVICE inline Array() = default;

  HOSTDEVICE inline void Fill(const T &val) {}

  HOSTDEVICE inline constexpr T *Get() const { return nullptr; }

  // Add constexpr to GetMutable() cause warning in MAC
  HOSTDEVICE inline T *GetMutable() { return nullptr; }

  HOSTDEVICE inline T &operator[](size_t index) {
#ifndef __CUDA_ARCH__
    PADDLE_THROW("Array<T, 0> has no element");
#endif
  }

  HOSTDEVICE inline const T &operator[](size_t index) const {
#ifndef __CUDA_ARCH__
    PADDLE_THROW("Array<T, 0> has no element");
#endif
  }

  HOSTDEVICE constexpr size_t size() const { return 0; }

  HOSTDEVICE constexpr bool operator==(const Array<T, 0> &other) const {
    return true;
  }

  HOSTDEVICE constexpr bool operator!=(const Array<T, 0> &other) const {
    return false;
  }
};

X
Xin Pan 已提交
101 102
}  // namespace framework
}  // namespace paddle