Add proper check for generic_histogram_binning_calibration_by_feature_cpu (#5042)

Summary:
Pull Request resolved: #5042

X-link: https://github.com/facebookresearch/FBGEMM/pull/2052

Add check to catch any index overflow

Reviewed By: ionuthristodorescu

Differential Revision: D85206442

fbshipit-source-id: b6e38949344e7c450263933dbe1add2bad057cef

Author: spcyppt

fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp

@@ -42,12 +42,19 @@ void _to_dense_representation(
     const int64_t num_lengths,
     const SegmentValueType* const segment_value_data,
     const SegmentLengthType* const segment_lengths_data,
-    SegmentValueType* const dense_segment_value_data) {
+    SegmentValueType* const dense_segment_value_data,
+    const int64_t num_segment_value) {
   int k = 0;
   for (const auto i : c10::irange(num_lengths)) {
     if (segment_lengths_data[i] == 1) {
       // Add 1 to distinguish between 0 inserted by densification vs. original
       // value.
+      TORCH_CHECK(
+          k < num_segment_value,
+          "k should be less than num_segment_value ",
+          num_segment_value,
+          " but found k = ",
+          k);
       dense_segment_value_data[i] = segment_value_data[k] + 1;
     } else {
       dense_segment_value_data[i] = 0;
@@ -2504,7 +2511,8 @@ std::tuple<Tensor, Tensor> histogram_binning_calibration_by_feature_cpu(
                   segment_lengths.numel(),
                   segment_value.data_ptr<segment_value_t>(),
                   segment_lengths.data_ptr<segment_length_t>(),
-                  dense_segment_value.data_ptr<segment_value_t>());
+                  dense_segment_value.data_ptr<segment_value_t>(),
+                  segment_value.numel());
             });
       });
 
@@ -2613,6 +2621,16 @@ std::tuple<Tensor, Tensor> generic_histogram_binning_calibration_by_feature_cpu(
   // dense_segment_value is used as a temporary storage.
   Tensor dense_segment_value =
       at::empty({logit.numel()}, segment_value.options());
+
+  // _to_dense_representation will access dense_segment_value[i] where i <=
+  // num_length, so num_length should be within the range of dense_segment_value
+  TORCH_CHECK(
+      segment_lengths.numel() <= dense_segment_value.numel(),
+      "segment_lengths numel (num_length) should be less than dense_segment_value numel of ",
+      dense_segment_value.numel(),
+      " but found num_length = ",
+      segment_lengths.numel());
+
   AT_DISPATCH_INDEX_TYPES(
       segment_value.scalar_type(), "to_dense_representation_cpu_wrapper", [&] {
         using segment_value_t = index_t;
@@ -2623,7 +2641,8 @@ std::tuple<Tensor, Tensor> generic_histogram_binning_calibration_by_feature_cpu(
                   segment_lengths.numel(),
                   segment_value.data_ptr<segment_value_t>(),
                   segment_lengths.data_ptr<segment_length_t>(),
-                  dense_segment_value.data_ptr<segment_value_t>());
+                  dense_segment_value.data_ptr<segment_value_t>(),
+                  segment_value.numel());
             });
       });