uniform_random_op.cu 6.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
L
Luo Tao 已提交
2 3 4 5 6 7 8 9 10 11 12 13

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. */
Q
qijun 已提交
14 15
#include <thrust/random.h>
#include <thrust/transform.h>
L
Leo Chen 已提交
16

Y
yaoxuefeng 已提交
17
#include "paddle/fluid/framework/generator.h"
Y
Yi Wang 已提交
18 19
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
20
#include "paddle/fluid/operators/uniform_random_op.h"
Q
qijun 已提交
21 22 23 24 25 26 27
namespace paddle {
namespace operators {

template <typename T>
struct UniformGenerator {
  T min_, max_;
  unsigned int seed_;
28 29 30 31 32 33 34 35 36 37 38
  T diag_val_;
  unsigned int diag_num_;
  unsigned int diag_step_;
  __host__ __device__ UniformGenerator(T min, T max, int seed, int diag_num,
                                       int diag_step, T diag_val)
      : min_(min),
        max_(max),
        seed_(seed),
        diag_num_(diag_num),
        diag_step_(diag_step),
        diag_val_(diag_val) {}
Q
qijun 已提交
39 40 41 42 43 44

  __host__ __device__ T operator()(const unsigned int n) const {
    thrust::minstd_rand rng;
    rng.seed(seed_);
    thrust::uniform_real_distribution<T> dist(min_, max_);
    rng.discard(n);
45 46 47 48 49 50
    T out = dist(rng);
    unsigned int remainder = n % (diag_step_ + 1);
    if (remainder == 0 && diag_num_ > n / (diag_step_ + 1)) {
      out = diag_val_;
    }
    return out;
Q
qijun 已提交
51 52 53
  }
};

Y
yaoxuefeng 已提交
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 85 86
template <typename T>
struct UniformGeneratorOffset {
  T min_, max_;
  unsigned int seed_;
  T diag_val_;
  unsigned int diag_num_;
  unsigned int diag_step_;
  int offset_;
  __host__ __device__ UniformGeneratorOffset(T min, T max, int seed,
                                             int diag_num, int diag_step,
                                             T diag_val, int offset)
      : min_(min),
        max_(max),
        seed_(seed),
        diag_num_(diag_num),
        diag_step_(diag_step),
        diag_val_(diag_val),
        offset_(offset) {}

  __host__ __device__ T operator()(const unsigned int n) const {
    thrust::minstd_rand rng;
    rng.seed(seed_);
    thrust::uniform_real_distribution<T> dist(min_, max_);
    rng.discard(n + offset_);
    T out = dist(rng);
    unsigned int remainder = n % (diag_step_ + 1);
    if (remainder == 0 && diag_num_ > n / (diag_step_ + 1)) {
      out = diag_val_;
    }
    return out;
  }
};

Q
qijun 已提交
87 88 89 90
// It seems that Eigen::Tensor::random in GPU will SEGFAULT.
// Use std::random and thrust::random(thrust is a std library in CUDA) to
// implement uniform random.
template <typename T>
Y
Yu Yang 已提交
91
class GPUUniformRandomKernel : public framework::OpKernel<T> {
Q
qijun 已提交
92 93
 public:
  void Compute(const framework::ExecutionContext& context) const override {
Y
Yancey1989 已提交
94
    framework::Tensor* tensor = nullptr;
Y
fix ci  
Yancey1989 已提交
95
    auto out_var = context.OutputVar("Out");
96 97 98 99 100 101
    std::vector<int64_t> new_shape;
    auto list_new_shape_tensor =
        context.MultiInput<framework::Tensor>("ShapeTensorList");
    if (list_new_shape_tensor.size() > 0 || context.HasInput("ShapeTensor")) {
      if (context.HasInput("ShapeTensor")) {
        auto* shape_tensor = context.Input<framework::Tensor>("ShapeTensor");
102
        new_shape = GetNewDataFromShapeTensor(shape_tensor);
103
      } else if (list_new_shape_tensor.size() > 0) {
104
        new_shape = GetNewDataFromShapeTensorList(list_new_shape_tensor);
105 106 107 108 109 110
      }
    }

    if (out_var->IsType<framework::SelectedRows>()) {
      auto* selected_rows = out_var->GetMutable<framework::SelectedRows>();
      tensor = selected_rows->mutable_value();
T
tangwei12 已提交
111
      auto shape = context.Attr<std::vector<int64_t>>("shape");
112
      if (!new_shape.empty()) shape = new_shape;
Y
Yancey1989 已提交
113
      tensor->Resize(framework::make_ddim(shape));
114 115 116 117
      selected_rows->mutable_rows()->reserve(shape[0]);
    } else if (out_var->IsType<framework::LoDTensor>()) {
      tensor = out_var->GetMutable<framework::LoDTensor>();
      if (!new_shape.empty()) tensor->Resize(framework::make_ddim(new_shape));
Y
Yancey1989 已提交
118
    } else {
119 120 121 122 123
      PADDLE_THROW(platform::errors::InvalidArgument(
          "Expected type of Output(out) in uniform_random_op must be Tensor, "
          "SelectedRows. But got "
          "unsupport type: %s.",
          framework::ToTypeName(out_var->Type())));
Y
Yancey1989 已提交
124
    }
Q
qijun 已提交
125
    T* data = tensor->mutable_data<T>(context.GetPlace());
Y
Pass CI  
Yu Yang 已提交
126
    unsigned int seed = static_cast<unsigned int>(context.Attr<int>("seed"));
Y
yaoxuefeng 已提交
127
    bool seed_flag = false;
L
Leo Chen 已提交
128 129 130
    if (seed == 0) {
      std::random_device rd;
      seed = rd();
Y
yaoxuefeng 已提交
131
      seed_flag = true;
Q
qijun 已提交
132
    }
L
Leo Chen 已提交
133

Y
Yu Yang 已提交
134 135
    T min = static_cast<T>(context.Attr<float>("min"));
    T max = static_cast<T>(context.Attr<float>("max"));
136 137 138 139 140
    unsigned int diag_num =
        static_cast<unsigned int>(context.Attr<int>("diag_num"));
    unsigned int diag_step =
        static_cast<unsigned int>(context.Attr<int>("diag_step"));
    T diag_val = static_cast<T>(context.Attr<float>("diag_val"));
Q
qijun 已提交
141
    thrust::counting_iterator<unsigned int> index_sequence_begin(0);
142
    int64_t size = tensor->numel();
Y
yaoxuefeng 已提交
143 144 145 146 147
    int device_id =
        BOOST_GET_CONST(platform::CUDAPlace, context.GetPlace()).GetDeviceId();
    auto gen_cuda = framework::GetDefaultCUDAGenerator(device_id);
    if (gen_cuda->GetIsInitPy() && seed_flag) {
      auto seed_offset = gen_cuda->IncrementOffset(1);
148
      int gen_offset = size * seed_offset.second;
Y
yaoxuefeng 已提交
149 150 151 152 153 154 155 156 157 158 159
      thrust::transform(
          index_sequence_begin, index_sequence_begin + size,
          thrust::device_ptr<T>(data),
          UniformGeneratorOffset<T>(min, max, seed_offset.first, diag_num,
                                    diag_step, diag_val, gen_offset));
    } else {
      thrust::transform(
          index_sequence_begin, index_sequence_begin + size,
          thrust::device_ptr<T>(data),
          UniformGenerator<T>(min, max, seed, diag_num, diag_step, diag_val));
    }
Q
qijun 已提交
160 161 162 163 164
  }
};

}  // namespace operators
}  // namespace paddle
Y
Yu Yang 已提交
165

166 167 168 169 170 171
REGISTER_OP_CUDA_KERNEL(uniform_random,
                        paddle::operators::GPUUniformRandomKernel<float>,
                        paddle::operators::GPUUniformRandomKernel<double>);
REGISTER_OP_CUDA_KERNEL(uniform_random_batch_size_like,
                        paddle::operators::GPUUniformRandomKernel<float>,
                        paddle::operators::GPUUniformRandomKernel<double>);