chore: import upstream snapshot with attribution

paddle/phi/kernels/impl/dirichlet_kernel_impl.h

@@ -0,0 +1,326 @@
+// 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.
+
+#pragma once
+
+#include <cmath>
+#include <random>
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#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/kernels/cpu/elementwise.h"
+#include "paddle/phi/kernels/dirichlet_kernel.h"
+#include "paddle/phi/kernels/elementwise_divide_kernel.h"
+#include "paddle/phi/kernels/funcs/broadcast_function.h"
+#include "paddle/phi/kernels/funcs/elementwise_functor.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/funcs/reduce_functor.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+
+// ROCM hcc doesn't work well with using std:: in kernel functions
+#if defined(PADDLE_WITH_CUDA)
+#define COMPAT_LOG log
+#define COMPAT_POW pow
+#define COMPAT_SQRT sqrt
+#else
+#define COMPAT_LOG std::log
+#define COMPAT_POW std::pow
+#define COMPAT_SQRT std::sqrt
+#endif
+
+#ifdef PADDLE_WITH_CUDA
+#include <curand_kernel.h>
+#endif
+#ifdef PADDLE_WITH_HIP
+#include <hiprand_kernel.h>
+#endif
+
+#if defined(PADDLE_WITH_CUDA)
+using COMPAT_RANDSTATEPHILOX4_32_10_T = curandStatePhilox4_32_10_t;
+#define COMPAT_RAND_INIT curand_init
+#define COMPAT_RAND_UNIFORM curand_uniform
+#define COMPAT_RAND_NORMAL curand_normal
+#elif defined(PADDLE_WITH_HIP)
+using COMPAT_RANDSTATEPHILOX4_32_10_T = hiprandStatePhilox4_32_10_t;
+#define COMPAT_RAND_INIT hiprand_init
+#define COMPAT_RAND_UNIFORM hiprand_uniform
+#define COMPAT_RAND_NORMAL hiprand_normal
+#endif
+
+namespace phi {
+
+template <typename ScalarT, typename SamplerT>
+struct BaseSampler {
+  SamplerT sampler_;
+  HOSTDEVICE BaseSampler(const SamplerT& sampler) : sampler_(sampler) {}
+  HOSTDEVICE ScalarT sample() {
+    // Sometimes convert float to float16/bfloat16
+    return static_cast<ScalarT>(sampler_());
+  }
+};
+
+template <typename Context, typename T>
+struct GammaSampler {
+  void operator()(const Context& dev_ctx,
+                  const DenseTensor& alpha,
+                  DenseTensor* out);
+};
+
+template <typename Context, typename T>
+struct DirichletSampler {
+  void operator()(const Context& dev_ctx,
+                  const DenseTensor& alpha,
+                  DenseTensor* out);
+};
+
+// `sample_gamma` is d from Numpy's distributions.c, and add support for
+//  paddle data type and code style.
+//  Source MIT licensed:
+/* Copyright 2005 Robert Kern (robert.kern@gmail.com)
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the
+ * "Software"), to deal in the Software without restriction, including
+ * without limitation the rights to use, copy, modify, merge, publish,
+ * distribute, sublicense, and/or sell copies of the Software, and to
+ * permit persons to whom the Software is furnished to do so, subject to
+ * the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included
+ * in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+ * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+template <typename ScalarT,
+          typename AccscalarT,
+          typename UniformSamplerT,
+          typename NormalSamplerT>
+HOSTDEVICE ScalarT
+sample_gamma(ScalarT alpha,
+             BaseSampler<AccscalarT, UniformSamplerT> standard_uniform,
+             BaseSampler<AccscalarT, NormalSamplerT> standard_normal) {
+  using MPTypeScalar = typename MPTypeTrait<ScalarT>::Type;
+  using MPTypeAccscalar = typename MPTypeTrait<AccscalarT>::Type;
+
+  MPTypeAccscalar mp_scale = static_cast<MPTypeAccscalar>(1.0f);
+  MPTypeScalar mp_alpha = static_cast<MPTypeScalar>(alpha);
+
+  // Boost alpha for higher acceptance probability.
+  if (mp_alpha < 1.0f) {
+    if (mp_alpha == 0.f) return static_cast<ScalarT>(0.f);
+    MPTypeAccscalar mp_sample =
+        static_cast<MPTypeAccscalar>(standard_uniform.sample());
+    mp_scale *= COMPAT_POW(1 - mp_sample, 1.0f / mp_alpha);
+    mp_alpha += 1.0f;
+  }
+
+  // This implements the acceptance-rejection method of Marsaglia and Tsang
+  // (2000)
+  // doi:10.1145/358407.358414
+  const MPTypeAccscalar d = mp_alpha - 1.0f / 3.0f;
+  const MPTypeAccscalar c = 1.0f / COMPAT_SQRT(9.0f * d);
+  for (;;) {
+    MPTypeAccscalar x, y;
+    do {
+      x = static_cast<MPTypeAccscalar>(standard_normal.sample());
+      y = 1.0f + c * x;
+    } while (y <= 0);
+    const MPTypeAccscalar v = y * y * y;
+    const MPTypeAccscalar u =
+        1 - static_cast<MPTypeAccscalar>(standard_uniform.sample());
+    const MPTypeAccscalar xx = x * x;
+    if (u < 1.0f - 0.0331f * xx * xx)
+      return static_cast<ScalarT>(mp_scale * d * v);
+    if (COMPAT_LOG(u) < 0.5f * xx + d * (1.0f - v + COMPAT_LOG(v)))
+      return static_cast<ScalarT>(mp_scale * d * v);
+  }
+}
+
+template <typename T, typename UniformSamplerT, typename NormalSamplerT>
+struct GammaCPUFunctor {
+  GammaCPUFunctor(const T* alpha,
+                  T* gamma,
+                  BaseSampler<T, UniformSamplerT> uniform,
+                  BaseSampler<T, NormalSamplerT> normal)
+      : alpha_(alpha), gamma_(gamma), uniform_(uniform), normal_(normal) {}
+
+  HOST void operator()(int64_t index) {
+    auto sample = sample_gamma<T, T, UniformSamplerT, NormalSamplerT>(
+        alpha_[index], uniform_, normal_);
+    gamma_[index] = std::max(std::numeric_limits<T>::min(), sample);
+  }
+
+  const T* alpha_;
+  T* gamma_;
+  BaseSampler<T, UniformSamplerT> uniform_;
+  BaseSampler<T, NormalSamplerT> normal_;
+};
+
+template <typename T>
+struct GammaSampler<CPUContext, T> {
+  void operator()(const CPUContext& dev_ctx,
+                  const DenseTensor& alpha,
+                  DenseTensor* out) {
+    auto generator = dev_ctx.GetGenerator()->GetCPUEngine();
+
+    auto uniform = [&generator]() -> T {
+      std::uniform_real_distribution<T> u(0.0, 1.0);
+      return u(*generator);
+    };
+    BaseSampler<T, decltype(uniform)> standard_uniform(uniform);
+
+    auto normal = [&generator]() {
+      std::normal_distribution<T> n(0.0, 1.0);
+      return n(*generator);
+    };
+    BaseSampler<T, decltype(normal)> standard_normal(normal);
+
+    GammaCPUFunctor<T, decltype(uniform), decltype(normal)> gamma_functor(
+        alpha.data<T>(), out->data<T>(), standard_uniform, standard_normal);
+    funcs::ForRange<CPUContext> for_range(dev_ctx, out->numel());
+    for_range(gamma_functor);
+  }
+};
+
+template <typename T>
+struct DirichletSampler<CPUContext, T> {
+  void operator()(const CPUContext& dev_ctx,
+                  const DenseTensor& alpha,
+                  DenseTensor* out) {
+    // sample from K gamma distributions, where K=alpha.numel()
+    DenseTensor gamma_samples;
+    gamma_samples.Resize(alpha.dims());
+    dev_ctx.template Alloc<T>(&gamma_samples);
+
+    GammaSampler<CPUContext, T> gamma_sampler;
+    gamma_sampler(dev_ctx, alpha, &gamma_samples);
+
+    // normalize them into a simplex, along the last axis
+    DenseTensor gamma_sum;
+    auto new_shape = gamma_samples.dims();
+    new_shape[new_shape.size() - 1] = 1;
+    gamma_sum.Resize(new_shape);
+    dev_ctx.template Alloc<T>(&gamma_sum);
+
+    funcs::ReduceKernelImpl<CPUContext, T, T, funcs::SumFunctor>(
+        dev_ctx,
+        gamma_samples,
+        &gamma_sum,
+        {new_shape.size() - 1},
+        true,
+        false);
+
+    funcs::ElementwiseCompute<funcs::DivideFunctor<T>, T>(
+        dev_ctx, gamma_samples, gamma_sum, funcs::DivideFunctor<T>(), out);
+  }
+};
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+template <typename T>
+struct GammaCUDAFunctor {
+  GammaCUDAFunctor(const T* alpha, T* gamma, uint64_t seed, uint64_t offset)
+      : alpha_(alpha), gamma_(gamma), seed_(seed), offset_(offset) {}
+
+  DEVICE void operator()(int64_t index) {
+    // curand initialization
+    COMPAT_RANDSTATEPHILOX4_32_10_T state;
+    COMPAT_RAND_INIT(
+        /*seed=*/seed_, /*subsequence=*/index, /*offset=*/offset_, &state);
+
+    // sample
+    auto uniform_lambda = [&state]() { return COMPAT_RAND_UNIFORM(&state); };
+    BaseSampler<T, decltype(uniform_lambda)> standard_uniform(uniform_lambda);
+    auto normal_lambda = [&state]() { return COMPAT_RAND_NORMAL(&state); };
+    BaseSampler<T, decltype(normal_lambda)> standard_normal(normal_lambda);
+
+    auto sample =
+        sample_gamma<T, T, decltype(uniform_lambda), decltype(normal_lambda)>(
+            alpha_[index], standard_uniform, standard_normal);
+    gamma_[index] = std::max(std::numeric_limits<T>::min(), sample);
+  }
+
+  const T* alpha_;
+  T* gamma_;
+  const uint64_t seed_;
+  const uint64_t offset_;
+};
+
+template <typename T>
+struct GammaSampler<GPUContext, T> {
+  void operator()(const GPUContext& dev_ctx,
+                  const DenseTensor& alpha,
+                  DenseTensor* out) {
+    auto p_gen = dev_ctx.GetGenerator();
+    auto seed_and_offset = p_gen->IncrementOffset(10);  // hard-coded offset
+    auto seed = seed_and_offset.first;
+    auto offset = seed_and_offset.second;
+
+    GammaCUDAFunctor<T> gamma_functor(
+        alpha.data<T>(), out->data<T>(), seed, offset);
+    funcs::ForRange<GPUContext> for_range(dev_ctx, out->numel());
+    for_range(gamma_functor);
+  }
+};
+
+template <typename T>
+struct DirichletSampler<GPUContext, T> {
+  void operator()(const GPUContext& dev_ctx,
+                  const DenseTensor& alpha,
+                  DenseTensor* out) {
+    // sample from K gamma distributions, where K=alpha.numel()
+    DenseTensor gamma_samples;
+    gamma_samples.Resize(alpha.dims());
+    dev_ctx.template Alloc<T>(&gamma_samples);
+
+    GammaSampler<GPUContext, T> gamma_sampler;
+    gamma_sampler(dev_ctx, alpha, &gamma_samples);
+
+    // normalize them into a simplex, along the last axis
+    DenseTensor gamma_sum;
+    auto new_shape = gamma_samples.dims();
+    new_shape[new_shape.size() - 1] = 1;
+    gamma_sum.Resize(new_shape);
+    dev_ctx.template Alloc<T>(&gamma_sum);
+
+    SumRawKernel<T, GPUContext>(dev_ctx,
+                                gamma_samples,
+                                {new_shape.size() - 1},
+                                true,
+                                false,
+                                gamma_sum.dtype(),
+                                &gamma_sum);
+    DivideKernel<T, GPUContext>(dev_ctx, gamma_samples, gamma_sum, out);
+  }
+};
+#endif
+
+template <typename T, typename Context>
+void DirichletKernel(const Context& dev_ctx,
+                     const DenseTensor& alpha,
+                     DenseTensor* out) {
+  dev_ctx.template Alloc<T>(out);
+  DirichletSampler<Context, T> sampler;
+  sampler(dev_ctx, alpha, out);
+}
+
+}  // namespace phi