【Hackathon 6th No.25】[Typing] 为 paddle.histogram 进行功能对齐与功能增强 -part (#63346)

AndPuQing · sunzhongkai588 · SigureMo · web-flow · commit 0c1eed41a990 · 2024-07-05T11:44:46.000+08:00
* Refactor histogram function to support weight and density in linalg.py

* Add histogram_bin_edges function to support bin edges computation in linalg.py

* Add class to test histogram weight and density

* Modify parameter position

* add histogram_bin_edges unit test

* add test case

* fix doc

* add test case

* update the func doc and fix histogram_bin_edges parameter

* fix doc

* remove merge mistake

* fix doc case

* Update python/paddle/tensor/linalg.py

* Update python/paddle/tensor/linalg.py

* Update python/paddle/tensor/linalg.py

* Apply suggestions from code review

* fix typing

* move to binary

* move histogramInferMeta to binary

---------

Co-authored-by: zachary sun &lt;70642955+sunzhongkai588@users.noreply.github.com&gt;
Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2248,6 +2248,47 @@ void HingeLossInferMeta(const MetaTensor& logits,
   loss->set_dtype(logits.dtype());
 }
 
+void HistogramInferMeta(const MetaTensor& input,
+                        const MetaTensor& weight,
+                        int64_t bins,
+                        int min,
+                        int max,
+                        bool density,
+                        MetaTensor* out) {
+  PADDLE_ENFORCE_GE(bins,
+                    1,
+                    phi::errors::InvalidArgument(
+                        "The bins should be greater than or equal to 1."
+                        "But received nbins is %d",
+                        bins));
+  PADDLE_ENFORCE_GE(
+      max,
+      min,
+      phi::errors::InvalidArgument("max must be larger or equal to min."
+                                   "But received max is %d, min is %d",
+                                   max,
+                                   min));
+  if (weight) {
+    auto weight_dims = weight.dims();
+    PADDLE_ENFORCE_EQ(
+        weight_dims,
+        input.dims(),
+        phi::errors::InvalidArgument(
+            "The shape of weight should be equal to the shape of input."
+            "But received weight shape is [%s], input shape is [%s]",
+            weight_dims,
+            input.dims()));
+  }
+
+  out->set_dims({bins});
+  out->share_lod(input);
+  if (density || weight) {
+    out->set_dtype(DataType::FLOAT32);
+  } else {
+    out->set_dtype(DataType::INT64);
+  }
+}
+
 void HuberLossInferMeta(const MetaTensor& input,
                         const MetaTensor& label,
                         float delta,
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -415,6 +415,14 @@ void HingeLossInferMeta(const MetaTensor& logits,
                         const MetaTensor& labels,
                         MetaTensor* loss);
 
+void HistogramInferMeta(const MetaTensor& input,
+                        const MetaTensor& weight,
+                        int64_t bins,
+                        int min,
+                        int max,
+                        bool density,
+                        MetaTensor* out);
+
 void HuberLossInferMeta(const MetaTensor& input_meta,
                         const MetaTensor& label_meta,
                         float delta,
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -2072,27 +2072,6 @@ void HashInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
-void HistogramInferMeta(
-    const MetaTensor& input, int64_t bins, int min, int max, MetaTensor* out) {
-  PADDLE_ENFORCE_GE(bins,
-                    1,
-                    phi::errors::InvalidArgument(
-                        "The bins should be greater than or equal to 1."
-                        "But received nbins is %d",
-                        bins));
-  PADDLE_ENFORCE_GE(
-      max,
-      min,
-      phi::errors::InvalidArgument("max must be larger or equal to min."
-                                   "But received max is %d, min is %d",
-                                   max,
-                                   min));
-
-  out->set_dims({bins});
-  out->share_lod(input);
-  out->set_dtype(DataType::INT64);
-}
-
 void IdentityLossInferMeta(const MetaTensor& x,
                            int reduction,
                            MetaTensor* out) {
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -356,9 +356,6 @@ void HashInferMeta(const MetaTensor& x,
                    int64_t mod_by,
                    MetaTensor* out);
 
-void HistogramInferMeta(
-    const MetaTensor& input, int64_t bins, int min, int max, MetaTensor* out);
-
 void IdentityLossInferMeta(const MetaTensor& x, int reduction, MetaTensor* out);
 
 void IncrementInferMeta(const MetaTensor& x, float value, MetaTensor* out);
diff --git a/paddle/phi/kernels/cpu/histogram_kernel.cc b/paddle/phi/kernels/cpu/histogram_kernel.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// 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.
@@ -17,27 +17,27 @@
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/utils/optional.h"
 
 namespace phi {
-
 template <typename T, typename Context>
 void HistogramKernel(const Context& dev_ctx,
                      const DenseTensor& input,
+                     const paddle::optional<DenseTensor>& weight,
                      int64_t bins,
                      int min,
                      int max,
+                     bool density,
                      DenseTensor* output) {
   auto& nbins = bins;
   auto& minval = min;
   auto& maxval = max;
 
   const T* input_data = input.data<T>();
+  auto weight_data = weight.get_ptr() ? weight.get_ptr()->data<T>() : nullptr;
   auto input_numel = input.numel();
 
-  int64_t* out_data = dev_ctx.template Alloc<int64_t>(output);
-  phi::funcs::SetConstant<Context, int64_t>()(
-      dev_ctx, output, static_cast<int64_t>(0));
-
   if (input_data == nullptr) return;
 
   T output_min = static_cast<T>(minval);
@@ -82,11 +82,38 @@ void HistogramKernel(const Context& dev_ctx,
           maxval,
           minval));
 
-  for (int64_t i = 0; i < input_numel; i++) {
-    if (input_data[i] >= output_min && input_data[i] <= output_max) {
-      const int64_t bin = (int64_t)((input_data[i] - output_min) * nbins /
-                                    (output_max - output_min));
-      out_data[std::min(bin, nbins - 1)] += 1;
+  if (density || weight_data) {
+    float* out_data = dev_ctx.template Alloc<float>(output);
+    phi::funcs::SetConstant<Context, float>()(
+        dev_ctx, output, static_cast<float>(0));
+    for (int64_t i = 0; i < input_numel; i++) {
+      if (input_data[i] >= output_min && input_data[i] <= output_max) {
+        const int64_t bin = (int64_t)((input_data[i] - output_min) * nbins /
+                                      (output_max - output_min));
+        out_data[std::min(bin, nbins - 1)] +=
+            weight_data ? static_cast<float>(weight_data[i]) : 1;
+      }
+    }
+    if (density) {
+      DenseTensor sum = phi::Sum<float, Context>(
+          dev_ctx, *output, phi::IntArray({0}), phi::DataType::FLOAT32, false);
+      float* sum_data = sum.data<float>();
+      float gap = static_cast<float>(nbins) /
+                  static_cast<float>((output_max - output_min)) / *sum_data;
+      for (int64_t i = 0; i < nbins; i++) {
+        out_data[i] *= gap;
+      }
+    }
+  } else {
+    int64_t* out_data = dev_ctx.template Alloc<int64_t>(output);
+    phi::funcs::SetConstant<Context, int64_t>()(
+        dev_ctx, output, static_cast<int64_t>(0));
+    for (int64_t i = 0; i < input_numel; i++) {
+      if (input_data[i] >= output_min && input_data[i] <= output_max) {
+        const int64_t bin = (int64_t)((input_data[i] - output_min) * nbins /
+                                      (output_max - output_min));
+        out_data[std::min(bin, nbins - 1)] += 1;
+      }
     }
   }
 }
@@ -100,6 +127,4 @@ PD_REGISTER_KERNEL(histogram,
                    float,
                    double,
                    int,
-                   int64_t) {
-  kernel->OutputAt(0).SetDataType(paddle::DataType::INT64);
-}
+                   int64_t) {}
diff --git a/paddle/phi/kernels/gpu/histogram_kernel.cu b/paddle/phi/kernels/gpu/histogram_kernel.cu
@@ -15,11 +15,13 @@
 #include "paddle/phi/kernels/histogram_kernel.h"
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
-#include "paddle/phi/backends/gpu/gpu_launch_config.h"
 #include "paddle/phi/backends/gpu/gpu_primitives.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/funcs/functors.h"
 #include "paddle/phi/kernels/funcs/math_cuda_utils.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -41,14 +43,15 @@ __device__ static IndexType GetBin(T input_value,
   return output_index;
 }
 
-template <typename T, typename IndexType>
+template <typename T, typename IndexType, typename Out_T>
 __global__ void KernelHistogram(const T* input,
+                                const T* weight,
                                 const int total_elements,
                                 const int64_t nbins,
                                 const T* min_value,
                                 const T* max_value,
-                                int64_t* output) {
-  extern __shared__ int64_t buf_hist[];
+                                Out_T* output) {
+  extern __shared__ float buf_hist[];
   for (int i = threadIdx.x; i < nbins; i += blockDim.x) {
     buf_hist[i] = 0;
   }
@@ -60,7 +63,8 @@ __global__ void KernelHistogram(const T* input,
     if (input_value >= *min_value && input_value <= *max_value) {
       const IndexType output_index =
           GetBin<T, IndexType>(input_value, *min_value, *max_value, nbins);
-      phi::CudaAtomicAdd(&buf_hist[output_index], 1);
+      phi::CudaAtomicAdd(&buf_hist[output_index],
+                         weight ? static_cast<float>(weight[input_index]) : 1);
     }
   }
   __syncthreads();
@@ -124,23 +128,29 @@ __global__ void KernelMinMax(const T min_value,
   }
 }
 
+__global__ void KernelMul(float* data, float* scale, int64_t numel) {
+  size_t index = threadIdx.x + blockIdx.x * blockDim.x;
+  if (index < numel) {
+    data[index] /= *scale;
+  }
+}
+
 template <typename T, typename Context>
 void HistogramKernel(const Context& dev_ctx,
                      const DenseTensor& input,
+                     const paddle::optional<DenseTensor>& weight,
                      int64_t bins,
                      int min,
                      int max,
+                     bool density,
                      DenseTensor* output) {
   auto& nbins = bins;
   auto& minval = min;
   auto& maxval = max;
 
   const T* input_data = input.data<T>();
   const int input_numel = input.numel();
-
-  int64_t* out_data = dev_ctx.template Alloc<int64_t>(output);
-  phi::funcs::SetConstant<Context, int64_t>()(
-      dev_ctx, output, static_cast<int64_t>(0));
+  auto weight_data = weight.get_ptr() ? weight.get_ptr()->data<T>() : nullptr;
 
   if (input_data == nullptr) return;
 
@@ -200,13 +210,50 @@ void HistogramKernel(const Context& dev_ctx,
           minval));
 
   auto stream = dev_ctx.stream();
-  KernelHistogram<T, IndexType><<<GET_BLOCKS(input_numel),
-                                  PADDLE_CUDA_NUM_THREADS,
-                                  nbins * sizeof(int64_t),
-                                  stream>>>(
-      input_data, input_numel, nbins, min_block_ptr, max_block_ptr, out_data);
-}
 
+  if (!density && !weight_data) {
+    int64_t* out_data = dev_ctx.template Alloc<int64_t>(output);
+    phi::funcs::SetConstant<Context, int64_t>()(dev_ctx, output, 0);
+    KernelHistogram<T, IndexType, int64_t><<<GET_BLOCKS(input_numel),
+                                             PADDLE_CUDA_NUM_THREADS,
+                                             nbins * sizeof(int64_t),
+                                             stream>>>(input_data,
+                                                       weight_data,
+                                                       input_numel,
+                                                       nbins,
+                                                       min_block_ptr,
+                                                       max_block_ptr,
+                                                       out_data);
+    return;
+
+  } else {
+    float* out_data = dev_ctx.template Alloc<float>(output);
+    phi::funcs::SetConstant<Context, float>()(
+        dev_ctx, output, static_cast<float>(0));
+    KernelHistogram<T, IndexType, float><<<GET_BLOCKS(input_numel),
+                                           PADDLE_CUDA_NUM_THREADS,
+                                           nbins * sizeof(int64_t),
+                                           stream>>>(input_data,
+                                                     weight_data,
+                                                     input_numel,
+                                                     nbins,
+                                                     min_block_ptr,
+                                                     max_block_ptr,
+                                                     out_data);
+    if (density) {
+      DenseTensor sum = phi::Sum<float, Context>(
+          dev_ctx, *output, phi::IntArray({0}), phi::DataType::FLOAT32, false);
+      float gap = static_cast<float>(nbins) /
+                  static_cast<float>(output_max - output_min);
+      std::vector<const DenseTensor*> ins = {output};
+      std::vector<DenseTensor*> outs = {output};
+      auto functor = phi::funcs::ScaleFunctor<float>(gap);
+      phi::funcs::ElementwiseKernel<float>(dev_ctx, ins, &outs, functor);
+      KernelMul<<<GET_BLOCKS(static_cast<int>(bins)),
+                  PADDLE_CUDA_NUM_THREADS>>>(out_data, sum.data<float>(), bins);
+    }
+  }
+}
 }  // namespace phi
 
 PD_REGISTER_KERNEL(histogram,
@@ -216,6 +263,4 @@ PD_REGISTER_KERNEL(histogram,
                    float,
                    double,
                    int,
-                   int64_t) {
-  kernel->OutputAt(0).SetDataType(paddle::DataType::INT64);
-}
+                   int64_t) {}
diff --git a/paddle/phi/kernels/histogram_kernel.h b/paddle/phi/kernels/histogram_kernel.h
@@ -20,9 +20,11 @@ namespace phi {
 template <typename T, typename Context>
 void HistogramKernel(const Context& dev_ctx,
                      const DenseTensor& input,
+                     const paddle::optional<DenseTensor>& weight,
                      int64_t bins,
                      int min,
                      int max,
+                     bool density,
                      DenseTensor* output);
 
 }  // namespace phi
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -1922,7 +1922,7 @@
 
 - op : histogram
   inputs :
-    input : X
+    {input: X, weight : Weight}
   outputs :
     out : Out
 
diff --git a/paddle/phi/ops/yaml/op_version.yaml b/paddle/phi/ops/yaml/op_version.yaml
@@ -326,6 +326,16 @@
           comment : In order to specify interpolation mode
           default : std::string("bilinear")
 
+- op : histogram
+  version :
+    - checkpoint : Upgrade histogram, add a new Input [Weight] and a new attribute [density]
+      action :
+        - add_input : Weight
+          comment : The weight of each value in the input tensor.
+        - add_attr : density
+          comment : If true, the histogram is normalized to form a probability density.
+          default : "false"
+
 - op : instance_norm
   version :
     - checkpoint : Change dispensable of attribute from False to True in instance_norm.
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2273,10 +2273,11 @@
   backward: hinge_loss_grad
 
 - op : histogram
-  args : (Tensor input, int64_t bins = 100, int min = 0, int max = 0)
+  args : (Tensor input, Tensor weight, int64_t bins = 100, int min = 0, int max = 0, bool density = false)
   output : Tensor(out)
   infer_meta :
     func : HistogramInferMeta
+  optional : weight
   kernel :
     func : histogram
 
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -227,6 +227,7 @@
     dot,
     eigvalsh,
     histogram,
+    histogram_bin_edges,
     histogramdd,
     matmul,
     mv,
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py
diff --git a/test/legacy_test/test_histogram_bin_edges_op.py b/test/legacy_test/test_histogram_bin_edges_op.py
diff --git a/test/legacy_test/test_histogram_op.py b/test/legacy_test/test_histogram_op.py
