Fix Histogram kernel to check range error (#64631)

qili93 · web-flow · commit 42a277445fd1 · 2024-05-28T17:00:36.000+08:00
* Fix Histogram kernel to check range error

* fix windows openblas ci fail
diff --git a/paddle/phi/kernels/cpu/histogram_kernel.cc b/paddle/phi/kernels/cpu/histogram_kernel.cc
@@ -51,6 +51,21 @@ void HistogramKernel(const Context& dev_ctx,
     output_max = output_max + 1;
   }
 
+  // check if out of range
+  double range =
+      static_cast<double>(output_max) - static_cast<double>(output_min);
+  PADDLE_ENFORCE_LT(
+      range,
+      static_cast<double>(std::numeric_limits<T>::max()),
+      phi::errors::InvalidArgument(
+          "The range of max - min is out of range for target type, "
+          "current kernel type is %s, the range should less than %f "
+          "but now min is %f, max is %f.",
+          typeid(T).name(),
+          std::numeric_limits<T>::max(),
+          output_min,
+          output_max));
+
   PADDLE_ENFORCE_EQ((std::isinf(static_cast<float>(output_min)) ||
                      std::isnan(static_cast<float>(output_max)) ||
                      std::isinf(static_cast<float>(output_min)) ||
diff --git a/paddle/phi/kernels/gpu/histogram_kernel.cu b/paddle/phi/kernels/gpu/histogram_kernel.cu
@@ -35,8 +35,8 @@ __device__ static IndexType GetBin(T input_value,
                                    T min_value,
                                    T max_value,
                                    int64_t nbins) {
-  IndexType bin = static_cast<int>((input_value - min_value) * nbins /
-                                   (max_value - min_value));
+  IndexType bin = static_cast<IndexType>((input_value - min_value) * nbins /
+                                         (max_value - min_value));
   IndexType output_index = bin < nbins - 1 ? bin : nbins - 1;
   return output_index;
 }
@@ -151,7 +151,7 @@ void HistogramKernel(const Context& dev_ctx,
   min_max.Resize({2 * block_num});
   auto* min_block_ptr = dev_ctx.template Alloc<T>(&min_max);
   auto* max_block_ptr = min_block_ptr + block_num;
-  if (output_min == output_max) {
+  if (min == max) {
     KernelMinMax<T><<<GET_BLOCKS(input_numel),
                       PADDLE_CUDA_NUM_THREADS,
                       0,
@@ -162,6 +162,27 @@ void HistogramKernel(const Context& dev_ctx,
         output_min, output_max, min_block_ptr, max_block_ptr);
   }
 
+  // copy min max value from GPU to CPU
+  std::vector<T> min_max_vec;
+  phi::TensorToVector(min_max, dev_ctx, &min_max_vec);
+  output_min = min_max_vec[0];
+  output_max = min_max_vec[1];
+
+  // check if out of range
+  double range =
+      static_cast<double>(output_max) - static_cast<double>(output_min);
+  PADDLE_ENFORCE_LT(
+      range,
+      static_cast<double>(std::numeric_limits<T>::max()),
+      phi::errors::InvalidArgument(
+          "The range of max - min is out of range for target type, "
+          "current kernel type is %s, the range should less than %f "
+          "but now min is %f, max is %f.",
+          typeid(T).name(),
+          std::numeric_limits<T>::max(),
+          output_min,
+          output_max));
+
   PADDLE_ENFORCE_EQ((std::isinf(static_cast<float>(output_min)) ||
                      std::isnan(static_cast<float>(output_max)) ||
                      std::isinf(static_cast<float>(output_min)) ||
diff --git a/test/legacy_test/test_histogram_op.py b/test/legacy_test/test_histogram_op.py
@@ -117,6 +117,29 @@ def net_func():
         with self.assertRaises(TypeError):
             self.run_network(net_func)
 
+    @test_with_pir_api
+    def test_input_range_error(self):
+        """Test range of input is out of bound"""
+
+        def net_func():
+            input_value = paddle.to_tensor(
+                [
+                    -7095538316670326452,
+                    -6102192280439741006,
+                    2040176985344715288,
+                    -6276983991026997920,
+                    -6570715756420355710,
+                    -5998045007776667296,
+                    -6763099356862306438,
+                    3166073479842736625,
+                ],
+                dtype=paddle.int64,
+            )
+            paddle.histogram(input=input_value, bins=1, min=0, max=0)
+
+        with self.assertRaises(ValueError):
+            self.run_network(net_func)
+
     @test_with_pir_api
     def test_type_errors(self):
         with paddle.static.program_guard(paddle.static.Program()):
