chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/lookup_table_kernel.cu

@@ -0,0 +1,134 @@
+// Copyright (c) 2024 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/gpu/lookup_table_kernel.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/mixed_vector.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+
+namespace phi {
+
+template <typename T,
+          int BlockDimX,
+          int BlockDimY,
+          int GridDimX,
+          bool PaddingFlag>
+__global__ void LookupTable(T *output,
+                            const T *table,
+                            const int64_t *ids,
+                            const int64_t N,
+                            const int64_t K,
+                            const int64_t D,
+                            const int64_t padding_idx) {
+  int idx = threadIdx.x;
+  int64_t idy = static_cast<int64_t>(blockIdx.x) +
+                static_cast<int64_t>(threadIdx.y) * GridDimX;
+
+  while (idy < K) {
+    int64_t id = ids[idy];
+    PADDLE_ENFORCE(
+        id >= 0,
+        "Variable value (input) of OP(lookup_table) "
+        "expected >= 0 and < %ld, but got %ld. Please check input value.",
+        N,
+        id);
+    PADDLE_ENFORCE(
+        id < N,
+        "Variable value (input) of OP(lookup_table) "
+        "expected >= 0 and < %ld, but got %ld. Please check input value.",
+        N,
+        id);
+    T *out = output + idy * D;
+    const T *tab = table + id * D;
+    for (int64_t i = idx; i < D; i += BlockDimX) {
+      if (PaddingFlag) {
+        if (id == padding_idx)
+          out[i] = static_cast<T>(0);
+        else
+          out[i] = tab[i];
+      } else {
+        out[i] = tab[i];
+      }
+    }
+    idy += BlockDimY * GridDimX;
+  }
+}
+
+template <typename T, typename Context>
+void LookupTableCUDAKernel(const Context &dev_ctx,
+                           const DenseTensor &w,
+                           const DenseTensor &ids_in,
+                           bool is_sparse,
+                           bool is_distributed UNUSED,
+                           int64_t padding_idx,
+                           bool remote_prefetch UNUSED,
+                           const std::string &entry_config UNUSED,
+                           bool is_test,
+                           const std::string &entry UNUSED,
+                           const std::string &table_class UNUSED,
+                           const std::vector<std::string> &table_names UNUSED,
+                           int trainer_id UNUSED,
+                           bool grad_inplace UNUSED,
+                           const std::vector<std::string> &epmap UNUSED,
+                           const std::vector<int64_t> &height_sections UNUSED,
+                           DenseTensor *out) {
+  auto *table_t = &w;
+  auto *ids_t = &ids_in;
+  auto *output_t = out;
+
+  size_t N = table_t->dims()[0];
+  size_t D = table_t->dims()[1];
+  size_t K = ids_t->numel();
+
+  auto *ids = ids_t->data<int64_t>();
+  auto *table = table_t->data<T>();
+  auto *output = dev_ctx.template Alloc<T>(output_t);
+
+#ifdef PADDLE_WITH_HIP
+  dim3 threads(64, 4);
+#else
+  dim3 threads(128, 8);
+#endif  // PADDLE_WITH_HIP
+  dim3 grids(8, 1);
+#ifdef PADDLE_WITH_HIP
+  if (padding_idx == -1)
+    LookupTable<T, 64, 4, 8, false><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+  else
+    LookupTable<T, 64, 4, 8, true><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+#else
+  if (padding_idx == -1)
+    LookupTable<T, 128, 8, 8, false><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+  else
+    LookupTable<T, 128, 8, 8, true><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+#endif  // PADDLE_WITH_HIP
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(lookup_table,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::LookupTableCUDAKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   int8_t,
+                   int16_t) {}