chore: import upstream snapshot with attribution

paddle/phi/kernels/custom/c_embedding_kernel.cc

@@ -0,0 +1,82 @@
+// Copyright (c) 2023 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/c_embedding_kernel.h"
+#include "glog/logging.h"
+#include "paddle/phi/api/backward/backward_api_base.h"
+#include "paddle/phi/api/include/api.h"
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+template <typename T, typename Context>
+void CEmbeddingKernel(const Context& dev_ctx,
+                      const DenseTensor& w,
+                      const DenseTensor& ids,
+                      int64_t start_index,
+                      int64_t vocab_size,
+                      DenseTensor* out) {
+  const auto& index_type = ids.dtype();
+  if (index_type == phi::DataType::INT32 ||
+      index_type == phi::DataType::INT64) {
+    auto out_dims = out->dims();
+    auto K = ids.numel();
+    auto N = w.dims()[0];
+    auto D = w.dims()[1];
+
+    auto x_tmp = std::make_shared<DenseTensor>();
+    x_tmp->ShareDataWith(ids).Resize({K});
+    auto w_tmp = std::make_shared<DenseTensor>();
+    w_tmp->ShareDataWith(w).Resize({N, D});
+    paddle::Tensor x_tensor(x_tmp), w_tensor(w_tmp);
+
+    auto start_index_tensor = paddle::experimental::full_like(
+        x_tensor, start_index, x_tensor.dtype(), x_tensor.place());
+    auto end_index_tensor = paddle::experimental::full_like(
+        x_tensor, start_index + N, x_tensor.dtype(), x_tensor.place());
+    auto ids_mask_tensor = paddle::experimental::logical_and(
+        x_tensor.greater_equal(start_index_tensor),
+        x_tensor.less_than(end_index_tensor));
+    auto ids_tensor = (x_tensor - start_index_tensor)
+                          .multiply(paddle::experimental::cast(
+                              ids_mask_tensor, x_tensor.dtype()));
+    auto out_tensor =
+        paddle::experimental::reshape(
+            paddle::experimental::cast(ids_mask_tensor, w_tensor.dtype()),
+            {K, 1})
+            .multiply(paddle::experimental::reshape(
+                paddle::experimental::embedding(
+                    ids_tensor, w_tensor, -1, false),
+                {K, D}));
+    out->ShareDataWith(*reinterpret_cast<DenseTensor*>(out_tensor.impl().get()))
+        .Resize(out_dims);
+  } else {
+    PADDLE_THROW(common::errors::Unavailable(
+        "Custom Device c_embedding ids only support int32 or int64."));
+  }
+}
+#endif
+}  // namespace phi
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+PD_REGISTER_KERNEL(c_embedding,
+                   Custom,
+                   ALL_LAYOUT,
+                   phi::CEmbeddingKernel,
+                   float,
+                   phi::float16,
+                   phi::bfloat16) {}
+#endif