【Hackathon 6th Fundable Projects 3 No.290】fluid operator pyramid_hash

paddle/phi/infermeta/multiary.cc

@@ -3879,6 +3879,137 @@ void PsroiPoolInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
+void PyramidHashInferMeta(const MetaTensor& x,
+                          const MetaTensor& w,
+                          const MetaTensor& white_list,
+                          const MetaTensor& black_list,
+                          int num_emb,
+                          int space_len,
+                          int pyramid_layer,
+                          int rand_len,
+                          float drop_out_percent,
+                          int is_training,
+                          bool use_filter,
+                          int white_list_len,
+                          int black_list_len,
+                          int seed,
+                          float lr,
+                          const std::string& distribute_update_vars,
+                          MetaTensor* out,
+                          MetaTensor* drop_pos,
+                          MetaTensor* x_temp_out,
+                          MetaConfig config) {
+  const auto& x_dims = x.dims();
+  PADDLE_ENFORCE_EQ(x_dims.size(),
+                    2,
+                    phi::errors::InvalidArgument(
+                        "The rank of Input(X) of PyramidHashOP is invalid. "
+                        "It should be 2, but got %d",
+                        x_dims.size()));
+
+  const auto& w_dims = w.dims();
+  PADDLE_ENFORCE_EQ(w_dims.size(),
+                    2,
+                    phi::errors::InvalidArgument(
+                        "The rank of Input(W) of PyramidHashOP is invalid. "
+                        "It should be 2, but got %d",
+                        w_dims.size()));
+
+  PADDLE_ENFORCE_EQ(
+      w_dims[0],
+      space_len + rand_len,
+      phi::errors::InvalidArgument(
+          "The first dimension of Input(W) of PyramidHashOP is invalid. "
+          "It should be space_len + rand_len, but now %d != %d + %d",
+          w_dims[0],
+          space_len,
+          rand_len));
+  PADDLE_ENFORCE_EQ(
+      w_dims[1],
+      1,
+      phi::errors::InvalidArgument(
+          "The second dimension of Input(W) of PyramidHashOP is invalid."
+          " It should be 1, but got %d",
+          w_dims[1]));
+
+  PADDLE_ENFORCE_EQ(
+      num_emb % rand_len,
+      0,
+      phi::errors::InvalidArgument(
+          "The PyramidHashOP's Attr(num_emb) should mod Attr(rand_len), "
+          "but num_emb is %d, rand_len is %d",
+          num_emb,
+          rand_len));
+
+  if (white_list_len > 0) {
+    PADDLE_ENFORCE_EQ(
+        white_list.initialized(),
+        true,
+        phi::errors::NotFound("Input(WhiteList) of PyramidHashOP is not "
+                              "found but white_list_len > 0."));
+    const auto& wl_dims = white_list.dims();
+    PADDLE_ENFORCE_EQ(
+        wl_dims.size(),
+        2,
+        phi::errors::InvalidArgument(
+            "The rank of Input(WhiteList) of PyramidHashOP is invalid."
+            " It should be 2, but got %d",
+            wl_dims.size()));
+    PADDLE_ENFORCE_EQ(wl_dims[0],
+                      white_list_len,
+                      phi::errors::InvalidArgument(
+                          "The first dimension of Input(WhiteList) of "
+                          "PyramidHashOP is invalid."
+                          " It should be equal to Attr(white_list_len) "
+                          ", but first dimension is %d, white_list_len is %d",
+                          wl_dims[0],
+                          white_list_len));
+    PADDLE_ENFORCE_EQ(wl_dims[1],
+                      1,
+                      phi::errors::InvalidArgument(
+                          "The second dimension of Input(WhiteList) of "
+                          "PyramidHashOP is invalid."
+                          " It should be 1, but got %d",
+                          wl_dims[1]));
+  }
+
+  if (black_list_len > 0) {
+    const auto& bl_dims = black_list.dims();
+    PADDLE_ENFORCE_EQ(
+        bl_dims.size(),
+        2,
+        phi::errors::InvalidArgument(
+            "The rank of Input(BlackList) of PyramidHashOP is invalid."
+            " It should be 2, but got %d",
+            bl_dims.size()));
+    PADDLE_ENFORCE_EQ(bl_dims[0],
+                      black_list_len,
+                      phi::errors::InvalidArgument(
+                          "The first dimension of Input(BlackList) of "
+                          "PyramidHashOP is invalid."
+                          " It should be equal to Attr(black_list_len)"
+                          ", but first dimension is %d, black_list_len is %d",
+                          bl_dims[0],
+                          black_list_len));
+    PADDLE_ENFORCE_EQ(bl_dims[1],
+                      1,
+                      phi::errors::InvalidArgument(
+                          "The second dimension of Input(BlackList) of "
+                          "PyramidHashOP is invalid."
+                          " It should be 1, but got %d",
+                          bl_dims[1]));
+  }
+
+  if (config.is_runtime) {
+    // something to do in runtime.
+  } else {
+    // compile time
+    out->set_dims(common::make_ddim({-1, num_emb}));
+    x_temp_out->set_dims(x_dims);
+    out->share_lod(x);
+  }
+}
+
 void QuantizeLinearInferMeta(const MetaTensor& x,
                              const MetaTensor& scale,
                              const MetaTensor& zero_point,

paddle/phi/infermeta/multiary.h

@@ -730,6 +730,27 @@ void PsroiPoolInferMeta(const MetaTensor& x,
                         float spatial_scale,
                         MetaTensor* out);
 
+void PyramidHashInferMeta(const MetaTensor& x,
+                          const MetaTensor& w,
+                          const MetaTensor& white_list,
+                          const MetaTensor& black_list,
+                          int num_emb,
+                          int space_len,
+                          int pyramid_layer,
+                          int rand_len,
+                          float drop_out_percent,
+                          int is_training,
+                          bool use_filter,
+                          int white_list_len,
+                          int black_list_len,
+                          int seed,
+                          float lr,
+                          const std::string& distribute_update_vars,
+                          MetaTensor* out,
+                          MetaTensor* drop_pos,
+                          MetaTensor* x_temp_out,
+                          MetaConfig config = MetaConfig());
+
 void QuantizeLinearInferMeta(const MetaTensor& x,
                              const MetaTensor& scale,
                              const MetaTensor& zero_point,

paddle/phi/kernels/cpu/pyramid_hash_grad_kernel.cc

@@ -0,0 +1,127 @@
+// 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 <xxhash.h>
+
+#include <algorithm>
+#include <cmath>
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/search_compute.h"
+
+namespace phi {
+
+template <typename T>
+void hash_embedding_bp(const T* hash_id,
+                       int len,
+                       const T* top_pos,
+                       T* weights,
+                       T mlr,
+                       int _num_emb,
+                       int _rand_len,
+                       int _space_len) {
+  for (int j = 0; j != _num_emb; j += _rand_len) {
+    unsigned int pos = XXH32(hash_id, len * sizeof(T), j) % _space_len;
+    phi::funcs::axpy(top_pos + j, weights + pos, _rand_len, mlr);
+  }
+}
+
+template <typename T, typename Context>
+void CPUPyramidHashOPGradKernel(const Context& dev_ctx,
+                                const DenseTensor& x,
+                                const DenseTensor& w,
+                                const DenseTensor& drop_pos,
+                                const DenseTensor& x_temp_out,
+                                const DenseTensor& out_grad,
+                                int num_emb,
+                                int space_len,
+                                int pyramid_layer,
+                                int rand_len,
+                                float drop_out_percent UNUSED,
+                                int is_training,
+                                bool use_filter,
+                                int white_list_len UNUSED,
+                                int black_list_len UNUSED,
+                                int seed UNUSED,
+                                float lr,
+                                const std::string& distribute_update_vars
+                                    UNUSED,
+                                DenseTensor* x_grad) {
+  auto* bottom = &x;
+  auto* _blobs = &w;
+  auto* drop_pos_p = &drop_pos;
+  auto* top = &out_grad;
+
+  int _num_emb = num_emb;
+  float _lr = lr;
+  int _rand_len = rand_len;
+  int _space_len = space_len;
+  int _pyramid_layer = pyramid_layer;
+
+  auto* buff = &x_temp_out;
+  auto* bottom_data = buff->data<T>();
+
+  int _slot_len = static_cast<int>(bottom->dims()[0]);
+  if (static_cast<size_t>(_slot_len) == bottom->lod()[0].size() - 1 &&
+      std::count(bottom_data, bottom_data + _slot_len, -1) == _slot_len) {
+    return;
+  }
+
+  auto& offset = bottom->lod()[0];
+  auto& drop_pos_offset = drop_pos_p->lod()[0];
+
+  const auto* top_diff = top->data<T>();
+  // in-place update weight, so need const_cast
+  T* weights = const_cast<T*>(_blobs->data<T>());
+  T mlr = -1.0 * _lr;
+
+  const int* iter = drop_pos_p->data<int>();
+  int top_counter = 0;
+  for (size_t i = 0; i < offset.size() - 1; ++i) {
+    int w = static_cast<int>(offset[i + 1] - offset[i]);
+    int w_drop = static_cast<int>(drop_pos_offset[i + 1] - drop_pos_offset[i]);
+    if (w_drop == 0) {
+      top_counter++;
+    }
+    if (w > 1) {
+      for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
+        for (int l = 0; l < w - ilayer; ++l) {
+          if (*(iter++) == 0) {
+            // do nothing
+          } else {
+            const T* top_pos = top_diff + top_counter++ * _num_emb;
+            hash_embedding_bp<T>((const T*)(bottom_data + offset[i] + l),
+                                 ilayer + 1,
+                                 top_pos,
+                                 weights,
+                                 mlr,
+                                 _num_emb,
+                                 _rand_len,
+                                 _space_len);
+          }
+        }
+      }
+    } else {
+      // do nothing
+    }
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(pyramid_hash_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::CPUPyramidHashOPGradKernel,
+                   float) {}