chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/truncated_gaussian_random_kernel.cu

@@ -0,0 +1,145 @@
+// Copyright (c) 2022 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/phi/kernels/truncated_gaussian_random_kernel.h"
+
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+#include <thrust/random.h>
+#include <thrust/transform.h>
+
+#include <limits>
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename MT>
+struct GPUTruncatedNormal {
+  MT mean, std, a, b;
+  MT a_normal_cdf;
+  MT b_normal_cdf;
+  unsigned int seed;
+  MT numeric_min;
+
+  __host__ __device__
+  GPUTruncatedNormal(MT mean, MT std, MT numeric_min, int seed, MT a, MT b)
+      : mean(mean), std(std), seed(seed), numeric_min(numeric_min), a(a), b(b) {
+    a_normal_cdf = (1.0 + erff((a - mean) / std / sqrtf(2.0))) / 2.0;
+    b_normal_cdf = (1.0 + erff((b - mean) / std / sqrtf(2.0))) / 2.0;
+  }
+
+  __host__ __device__ T operator()(const unsigned int n) const {
+    thrust::minstd_rand rng;
+    rng.seed(seed);
+    thrust::uniform_real_distribution<MT> dist(numeric_min, 1);
+    rng.discard(n);
+    MT value = dist(rng);
+    auto p = a_normal_cdf + (b_normal_cdf - a_normal_cdf) * value;
+    MT ret = std::sqrt(2.0) * erfinvf(2 * p - 1) * std + mean;
+    return static_cast<T>(std::clamp(ret, a, b));
+  }
+};
+
+template <typename T, typename MT>
+struct TruncatedNormalOffset {
+  MT mean, std, a, b;
+  MT a_normal_cdf;
+  MT b_normal_cdf;
+  unsigned int seed;
+  MT numeric_min;
+  int offset_;
+
+  __host__ __device__ TruncatedNormalOffset(
+      MT mean, MT std, MT numeric_min, int seed, int offset, MT a, MT b)
+      : mean(mean),
+        std(std),
+        seed(seed),
+        numeric_min(numeric_min),
+        offset_(offset),
+        a(a),
+        b(b) {
+    a_normal_cdf = (1.0 + erff((a - mean) / std / sqrtf(2.0))) / 2.0;
+    b_normal_cdf = (1.0 + erff((b - mean) / std / sqrtf(2.0))) / 2.0;
+  }
+
+  __host__ __device__ T operator()(const unsigned int n) const {
+    thrust::minstd_rand rng;
+    rng.seed(seed);
+    thrust::uniform_real_distribution<MT> dist(numeric_min, 1);
+    rng.discard(n + offset_);
+    MT value = dist(rng);
+    auto p = a_normal_cdf + (b_normal_cdf - a_normal_cdf) * value;
+    MT ret = std::sqrt(2.0) * erfinvf(2 * p - 1) * std + mean;
+    return static_cast<T>(std::clamp(ret, a, b));
+  }
+};
+
+template <typename T, typename Context>
+void TruncatedGaussianRandomKernel(const Context& dev_ctx,
+                                   const std::vector<int>& shape,
+                                   float mean,
+                                   float std,
+                                   int seed,
+                                   float a,
+                                   float b,
+                                   DataType dtype,
+                                   DenseTensor* out) {
+  T* data = dev_ctx.template Alloc<T>(out);
+
+  using MT = typename MPTypeTrait<T>::Type;
+
+  thrust::counting_iterator<int64_t> index_sequence_begin(0);
+  int64_t size = out->numel();
+
+  auto gen_cuda = dev_ctx.GetGenerator();
+  if (seed == 0) {
+    // use global Generator seed
+    auto seed_offset = gen_cuda->IncrementOffset(1);
+    uint64_t seed = seed_offset.first;
+    uint64_t offset = seed_offset.second;
+    thrust::transform(
+        index_sequence_begin,
+        index_sequence_begin + size,
+        thrust::device_ptr<T>(data),
+        TruncatedNormalOffset<T, MT>(mean,
+                                     std,
+                                     std::numeric_limits<MT>::min(),
+                                     seed,
+                                     size * offset,
+                                     a,
+                                     b));
+  } else {
+    // use OP seed
+    thrust::transform(
+        index_sequence_begin,
+        index_sequence_begin + size,
+        thrust::device_ptr<T>(data),
+        GPUTruncatedNormal<T, MT>(
+            mean, std, std::numeric_limits<MT>::min(), seed, a, b));
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(truncated_gaussian_random,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::TruncatedGaussianRandomKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}