Validate GatherBlockQuantized CUDA indices bounds

onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc

@@ -5,11 +5,43 @@
 #include "contrib_ops/cuda/quantization/gather_block_quantized.h"
 #include "contrib_ops/cuda/quantization/gather_block_quantized.cuh"
 
+#include <vector>
+
 namespace onnxruntime {
 namespace contrib {
 namespace cuda {
 using namespace onnxruntime::cuda;
 
+namespace {
+
+template <typename Tind>
+Status ValidateIndicesRangeForCuda(const Tensor* indices, int64_t gather_axis_dim, cudaStream_t stream) {
+  const size_t indices_size = static_cast<size_t>(indices->Shape().Size());
+  std::vector<Tind> indices_host(indices_size);
+
+  if (indices->Location().device.Type() == OrtDevice::CPU) {
+    const Tind* indices_ptr = indices->Data<Tind>();
+    for (size_t i = 0; i < indices_size; ++i) {
+      indices_host[i] = indices_ptr[i];
+    }
+  } else {
+    CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(indices_host.data(), indices->Data<Tind>(), indices_size * sizeof(Tind),
+                                         cudaMemcpyDefault, stream));
+    CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream));
+  }
+
+  for (size_t i = 0; i < indices_size; ++i) {
+    const int64_t indices_val = static_cast<int64_t>(indices_host[i]);
+    ORT_RETURN_IF_NOT(indices_val >= -gather_axis_dim && indices_val < gather_axis_dim,
+                      "indices element out of data bounds, idx=", indices_val,
+                      " must be within the inclusive range [", -gather_axis_dim, ",", gather_axis_dim - 1, "]");
+  }
+
+  return Status::OK();
+}
+
+}  // namespace
+
 #define REGISTER_GATHERBLOCKQUANTIZED(T1, T2, Tind)                     \
   ONNX_OPERATOR_THREE_TYPED_KERNEL_EX(                                  \
       GatherBlockQuantized,                                             \
@@ -140,6 +172,8 @@
   param.gather_axis = gather_axis_;
   param.N = N;
 
+  ORT_RETURN_IF_ERROR(ValidateIndicesRangeForCuda<Tind>(indices, param.gather_axis_dim, param.stream));
+
   const auto dequantized_type = scales->GetElementType();
   if (dequantized_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
     const auto* scales_ptr = static_cast<const float*>(scales->DataRaw());

onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc

@@ -499,6 +499,33 @@ void Test_InvalidIndices_WithZeroPoints() {
                                 gather_axis, quantize_axis, block_size, bits, output, output_shape, false, true);
 }
 
+template <typename T1, typename T2, typename Tind>
+void Test_NegativeInvalidIndices_WithZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {-3};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 3, 1};
+  std::vector<int> zero_points = {-1, 1, 0, 0, 1, -1};
+  std::vector<float> output = {8.f, 10.f, 12.f, 14.f,
+                               3.f, 4.f, 5.f, 6.f,
+                               -6.f, -4.f, -2.f, 0.f};
+  std::vector<int64_t> output_shape = {1, 3, 4};
+
+  constexpr int64_t gather_axis = 0;
+  constexpr int64_t quantize_axis = 2;
+  constexpr int64_t block_size = 16;
+  constexpr int64_t bits = 4;
+  RunUnpackedData<T1, T2, Tind>(data, data_shape, indices, indices_shape, scales, scales_shape, zero_points,
+                                gather_axis, quantize_axis, block_size, bits, output, output_shape, false, true);
+}
+
 #ifndef USE_CUDA
 TEST(GatherBlockQuantizedOpTest, InvalidIndices) {
   Test_InvalidIndices_WithZeroPoints<UInt4x2, float, int32_t>();
@@ -507,6 +534,20 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) {
 }
 #endif
 
+#ifdef USE_CUDA
+TEST(GatherBlockQuantizedOpTest, InvalidIndices_Cuda) {
+  Test_InvalidIndices_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_InvalidIndices_WithZeroPoints<UInt4x2, float, int64_t>();
+  Test_InvalidIndices_WithZeroPoints<uint8_t, float, int32_t>();
+}
+
+TEST(GatherBlockQuantizedOpTest, NegativeInvalidIndices_Cuda) {
+  Test_NegativeInvalidIndices_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_NegativeInvalidIndices_WithZeroPoints<UInt4x2, float, int64_t>();
+  Test_NegativeInvalidIndices_WithZeroPoints<uint8_t, float, int32_t>();
+}
+#endif
+
 template <typename T1, typename T2, typename Tind>
 void Test_GatherAxis0_WithZeroPoints(int bits = 4) {
   std::vector<int> data = {-8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8,