【Hackathon 9th No.5】Add index check for index_sample gpu kernel

paddle/phi/kernels/gpu/index_sample_kernel.cu

@@ -15,6 +15,7 @@
 #include "paddle/phi/kernels/index_sample_kernel.h"
 
 #include <algorithm>
+#include <limits>
 #include <vector>
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
@@ -31,6 +32,35 @@ namespace {
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #define UINT32_MAX std::numeric_limits<uint32_t>::max()
 
+// Helper to check that indices are within [0, input_length) on CPU.
+template <typename IndexT, typename Context>
+inline void CheckIndexOnCPU(const Context& dev_ctx,
+                            const DenseTensor& index,
+                            int64_t input_length) {
+  DenseTensor index_tensor_cpu;
+  phi::Copy(dev_ctx, index, phi::CPUPlace(), true, &index_tensor_cpu);
+  const IndexT* index_data_cpu = index_tensor_cpu.data<IndexT>();
+  for (int64_t i = 0; i < index.numel(); i++) {
+    PADDLE_ENFORCE_GE(
+        index_data_cpu[i],
+        0,
+        errors::InvalidArgument(
+            "Variable value (index) of OP(index_sample) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            input_length,
+            index_data_cpu[i]));
+    PADDLE_ENFORCE_LT(
+        index_data_cpu[i],
+        input_length,
+        errors::InvalidArgument(
+            "Variable value (index) of OP(index_sample) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            input_length,
+            index_data_cpu[i]));
+  }
+}
 }  // namespace
 
 template <typename T, typename SampleIndexT = int, typename ElementIndexT>
@@ -81,6 +111,13 @@ void IndexSampleKernel(const Context& dev_ctx,
   size_t input_length = input_dim[1];
   size_t index_length = index_dim[1];
 
+  if (index_type == DataType::INT64) {
+    CheckIndexOnCPU<int64_t>(
+        dev_ctx, index, static_cast<int64_t>(input_length));
+  } else if (index_type == DataType::INT32) {
+    CheckIndexOnCPU<int>(dev_ctx, index, static_cast<int64_t>(input_length));
+  }
+
   auto block_width = phi::backends::gpu::RoundToPowerOfTwo(index_length);
   block_width = MIN(block_width, PREDEFINED_BLOCK_SIZE_X);
   int block_height =

test/legacy_test/test_index_sample_op.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest, convert_float_to_uint16
+from op_test import OpTest, convert_float_to_uint16, get_places
 
 import paddle
 from paddle import base
@@ -174,6 +174,41 @@ def config(self):
         self.index_type = "int32"
 
 
+class TestIndexSampleOpError(unittest.TestCase):
+    def test_errors(self):
+        places = get_places()
+        for place in places:
+            with base.dygraph.guard(place):
+                self.assertRaises(
+                    ValueError,
+                    paddle.index_sample,
+                    x=paddle.rand([10, 0], dtype="float32"),
+                    index=paddle.randint(
+                        low=0, high=5, shape=[10, 1], dtype="int32"
+                    ),
+                )
+
+            with base.dygraph.guard(place):
+                self.assertRaises(
+                    ValueError,
+                    paddle.index_sample,
+                    x=paddle.rand([10, 0], dtype="float32"),
+                    index=paddle.randint(
+                        low=0, high=5, shape=[10, 1], dtype="int64"
+                    ),
+                )
+
+            with base.dygraph.guard(place):
+                self.assertRaises(
+                    ValueError,
+                    paddle.index_sample,
+                    x=paddle.to_tensor(
+                        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype="float32"
+                    ),
+                    index=paddle.to_tensor([[0, 1], [0, 4]], dtype="int64"),
+                )
+
+
 @unittest.skipIf(core.is_compiled_with_xpu(), "complex is not supported on XPU")
 class TestIndexSampleComplex64(TestIndexSampleOp):
     def config(self):