【Hackathon 6th Fundable Projects 3 No.189】 hash static (#65400)

co63oc · web-flow · commit b88a59f24ced · 2024-06-26T00:52:27.000+08:00
* Fix

* Fix

* Fix

* Fix

* Fix

* Fix

* Fix
diff --git a/paddle/fluid/operators/hash_op.cc b/paddle/fluid/operators/hash_op.cc
diff --git a/paddle/fluid/operators/hash_op.h b/paddle/fluid/operators/hash_op.h
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -146,6 +146,7 @@
     'recv_v2',
     'c_allgather',
     'qkv_unpack_mha',
+    'hash',
 ]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -24,6 +24,7 @@ limitations under the License. */
 #include "paddle/phi/core/infermeta_utils.h"
 #include "paddle/phi/core/utils/data_type.h"
 #include "paddle/phi/kernels/funcs/flatten2_utils.h"
+#include "paddle/phi/kernels/funcs/hash_utils.h"
 #include "paddle/phi/kernels/funcs/parse_qr_mode.h"
 #include "paddle/phi/kernels/funcs/pooling.h"
 #include "paddle/phi/kernels/funcs/slice_utils.h"
@@ -2054,6 +2055,23 @@ void GumbelSoftmaxInferMeta(const MetaTensor& x,
   UnchangedInferMetaCheckAxis(x, axis, out);
 }
 
+void HashInferMeta(const MetaTensor& x,
+                   int num_hash,
+                   int64_t mod_by,
+                   MetaTensor* out) {
+  const auto& dims = x.dims();
+  PADDLE_ENFORCE_EQ(dims.size(),
+                    2UL,
+                    phi::errors::InvalidArgument(
+                        "The input of hash_op's dimensions must be 2"));
+  std::vector<int64_t> out_dims;
+  phi::funcs::HashOutputSize(dims, out_dims, num_hash);
+
+  out->set_dims(common::make_ddim(out_dims));
+  out->share_lod(x);
+  out->set_dtype(x.dtype());
+}
+
 void HistogramInferMeta(
     const MetaTensor& input, int64_t bins, int min, int max, MetaTensor* out) {
   PADDLE_ENFORCE_GE(bins,
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -350,6 +350,12 @@ void GumbelSoftmaxInferMeta(const MetaTensor& x,
                             bool hard,
                             int axis,
                             MetaTensor* out);
+
+void HashInferMeta(const MetaTensor& x,
+                   int num_hash,
+                   int64_t mod_by,
+                   MetaTensor* out);
+
 void HistogramInferMeta(
     const MetaTensor& input, int64_t bins, int min, int max, MetaTensor* out);
 
diff --git a/paddle/phi/kernels/cpu/hash_kernel.cc b/paddle/phi/kernels/cpu/hash_kernel.cc
@@ -0,0 +1,17 @@
+// 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
+//
+//     https://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/impl/hash_kernel_impl.h"
+
+PD_REGISTER_KERNEL(hash, CPU, ALL_LAYOUT, phi::HashKernel, int, int64_t) {}
diff --git a/paddle/phi/kernels/funcs/hash_utils.h b/paddle/phi/kernels/funcs/hash_utils.h
@@ -0,0 +1,35 @@
+// 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
+//
+//     https://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 <vector>
+
+#include "paddle/phi/core/enforce.h"
+namespace phi {
+namespace funcs {
+
+inline void HashOutputSize(const phi::DDim& in_dims,
+                           std::vector<int64_t>& out_dims,  // NOLINT
+                           int num_hash) {
+  out_dims.reserve(in_dims.size() + 1);
+  // copy all dims except the last one
+  for (int i = 0u; i != in_dims.size() - 1; ++i) {
+    out_dims.emplace_back(in_dims[i]);
+  }
+  out_dims.emplace_back(num_hash);
+  // keep the last dim to 1
+  out_dims.emplace_back(1);
+}
+}  // namespace funcs
+}  // namespace phi
diff --git a/paddle/phi/kernels/impl/hash_kernel_impl.h b/paddle/phi/kernels/impl/hash_kernel_impl.h
@@ -0,0 +1,58 @@
+// 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
+//
+//     https://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
+
+extern "C" {
+#include <xxhash.h>
+}
+#include <vector>
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/hash_utils.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void HashKernel(const Context& dev_ctx,
+                const DenseTensor& x,
+                int num_hash,
+                int64_t mod_by,
+                DenseTensor* out) {
+  auto* out_t = out;
+  auto* in_t = &x;
+
+  auto in_dims = in_t->dims();
+
+  std::vector<int64_t> out_dims;
+  phi::funcs::HashOutputSize(in_dims, out_dims, num_hash);
+  out_t->Resize(common::make_ddim(out_dims));
+  auto* output = dev_ctx.template Alloc<T>(out_t);
+
+  auto seq_length = in_dims[0];
+  auto last_dim = in_dims[in_dims.size() - 1];
+  auto* input = in_t->data<T>();
+  for (int idx = 0; idx < seq_length; ++idx) {
+    for (int ihash = 0; ihash != num_hash; ++ihash) {
+      output[idx * num_hash + ihash] =
+          XXH64(input, sizeof(T) * last_dim, ihash) % mod_by;
+    }
+    input += last_dim;
+  }
+
+  out_t->set_lod(in_t->lod());
+}
+
+}  // namespace phi
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -406,6 +406,17 @@
     func : hardswish
   backward : hardswish_grad
 
+- op : hash
+  args: (Tensor x, int num_hash = 1, int64_t mod_by = 100000, bool runtime_shape = true)
+  output: Tensor (out)
+  infer_meta:
+    func: HashInferMeta
+    param: [x, num_hash, mod_by]
+  kernel:
+    func: hash
+    param: [x, num_hash, mod_by]
+    data_type: x
+
 - op : less_equal
   args : (Tensor x, Tensor y)
   output : Tensor(out)
diff --git a/paddle/phi/ops/yaml/legacy/static_ops.yaml b/paddle/phi/ops/yaml/legacy/static_ops.yaml
@@ -379,6 +379,17 @@
     param : [x]
   backward : hardswish_grad
 
+- op : hash
+  args: (Tensor x, int num_hash = 1, int64_t mod_by = 100000, bool runtime_shape = true)
+  output: Tensor (out)
+  infer_meta:
+    func: HashInferMeta
+    param: [x, num_hash, mod_by]
+  kernel:
+    func: hash
+    param: [x, num_hash, mod_by]
+    data_type: x
+
 - op : less_equal
   args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
   output : Tensor(out)
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -1869,6 +1869,14 @@
   outputs :
     out : Out
 
+- op : hash
+  inputs:
+    x : X
+  attrs :
+    runtime_shape : ALL_KERNELS_MUST_COMPUTE_RUNTIME_SHAPE
+  outputs:
+    out : Out
+
 - op : heaviside (elementwise_heaviside)
   backward : heaviside_grad (elementwise_heaviside_grad)
   inputs :
diff --git a/test/legacy_test/test_hash_op.py b/test/legacy_test/test_hash_op.py

Original file line number	Diff line number	Diff line change
`@@ -146,6 +146,7 @@`
`146`	`146`	`'recv_v2',`
`147`	`147`	`'c_allgather',`
`148`	`148`	`'qkv_unpack_mha',`
	`149`	`+ 'hash',`
`149`	`150`	`]`
`150`	`151`
`151`	`152`	`NO_NEED_GEN_STATIC_ONLY_APIS = [`