fix out of boundary vector per class in SVM

onnxruntime/core/providers/cpu/ml/svmclassifier.cc

@@ -50,29 +50,29 @@
     vector_count_ += narrow<ptrdiff_t>(vectors_per_class_[i]);
   }
 
+  ORT_ENFORCE(classlabels_strings_.size() > 0 || classlabels_ints_.size() > 0, "One of classlabels_strings, classlabels_ints is required.");
+
   using_strings_ = false;
   if (classlabels_strings_.size() > 0) {
     using_strings_ = true;
     class_count_ = classlabels_strings_.size();
-  } else if (classlabels_ints_.size() > 0) {
-    class_count_ = classlabels_ints_.size();
   } else {
-    class_count_ = 1;
+    class_count_ = classlabels_ints_.size();
   }
 
+  ORT_ENFORCE(proba_.size() == probb_.size(), "proba and probb must have the same size.");
+  ORT_ENFORCE(coefficients_.size() > 0, "coefficients are empty.");
+
   if (vector_count_ > 0) {
     feature_count_ = support_vectors_.size() / vector_count_;  // length of each support vector
     mode_ = SVM_TYPE::SVM_SVC;
+    ORT_ENFORCE(vectors_per_class_.size() == class_count_, "Mismatch between class_labels and vector_per_class dimensions.");
   } else {
     feature_count_ = coefficients_.size() / class_count_;  // liblinear mode
     mode_ = SVM_TYPE::SVM_LINEAR;
     set_kernel_type(KERNEL::LINEAR);
   }
 
-  ORT_ENFORCE(classlabels_strings_.size() > 0 || classlabels_ints_.size() > 0);
-  ORT_ENFORCE(proba_.size() == probb_.size());
-  ORT_ENFORCE(coefficients_.size() > 0);
-
   // Validate attribute array sizes against the declared dimensions to prevent
   // out-of-bounds reads from crafted models.
   if (mode_ == SVM_TYPE::SVM_SVC) {
@@ -233,10 +233,10 @@
   }
 
   int write_additional_scores = -1;
-  int64_t num_scores_per_batch = class_count_;
+  size_t num_scores_per_batch = class_count_;
 
   if (mode_ == SVM_TYPE::SVM_SVC && !have_proba) {
     num_scores_per_batch = num_classifiers;
     if (class_count_ <= 2) {
       write_additional_scores = post_transform_ == POST_EVAL_TRANSFORM::NONE ? 2 : 0;
     }
@@ -303,39 +303,39 @@
       // e.g. AB combines with BA.
       // If A has 3 support vectors and B has 2, there's a 3x2 block for AB and a 2x3 block for BA to combine
 
-      auto cur_kernels = kernels_span.subspan(n * SafeInt<size_t>(vector_count_), onnxruntime::narrow<size_t>(vector_count_));
-      auto cur_scores = classifier_scores.subspan(n * SafeInt<size_t>(num_slots_per_iteration), onnxruntime::narrow<size_t>(num_classifiers));
-      auto cur_votes = votes_span.subspan(n * SafeInt<size_t>(class_count_), onnxruntime::narrow<size_t>(class_count_));
+      auto cur_kernels = kernels_span.subspan(n * SafeInt<size_t>(vector_count_), vector_count_);
+      auto cur_scores = classifier_scores.subspan(n * SafeInt<size_t>(num_slots_per_iteration), num_classifiers);
+      auto cur_votes = votes_span.subspan(n * SafeInt<size_t>(class_count_), class_count_);
       auto scores_iter = cur_scores.begin();
 
       size_t classifier_idx = 0;
-      for (int64_t i = 0; i < class_count_ - 1; i++) {
-        int64_t start_index_i = starting_vector_[onnxruntime::narrow<size_t>(i)];  // start of support vectors for class i
-        int64_t class_i_support_count = vectors_per_class_[onnxruntime::narrow<size_t>(i)];
-        int64_t i_coeff_row_offset = vector_count_ * i;
+      for (size_t i = 0; i < class_count_ - 1; i++) {
+        size_t start_index_i = starting_vector_[i];  // start of support vectors for class i
+        size_t class_i_support_count = vectors_per_class_[i];
+        size_t i_coeff_row_offset = vector_count_ * i;
 
-        for (int64_t j = i + 1; j < class_count_; j++) {
-          int64_t start_index_j = starting_vector_[onnxruntime::narrow<size_t>(j)];  // start of support vectors for class j
-          int64_t class_j_support_count = vectors_per_class_[onnxruntime::narrow<size_t>(j)];
-          int64_t j_coeff_row_offset = vector_count_ * (j - 1);
+        for (size_t j = i + 1; j < class_count_; j++) {
+          size_t start_index_j = starting_vector_[j];  // start of support vectors for class j
+          size_t class_j_support_count = vectors_per_class_[j];
+          size_t j_coeff_row_offset = vector_count_ * (j - 1);
 
           double sum = 0;
 
           const float* val1 = &(coefficients_[j_coeff_row_offset + SafeInt<size_t>(start_index_i)]);
-          const float* val2 = &(cur_kernels[onnxruntime::narrow<size_t>(start_index_i)]);
-          for (int64_t m = 0; m < class_i_support_count; ++m, ++val1, ++val2)
+          const float* val2 = &(cur_kernels[start_index_i]);
+          for (size_t m = 0; m < class_i_support_count; ++m, ++val1, ++val2)
             sum += *val1 * *val2;
 
           val1 = &(coefficients_[i_coeff_row_offset + SafeInt<size_t>(start_index_j)]);
           val2 = &(cur_kernels[onnxruntime::narrow<size_t>(start_index_j)]);
 
-          for (int64_t m = 0; m < class_j_support_count; ++m, ++val1, ++val2)
+          for (size_t m = 0; m < class_j_support_count; ++m, ++val1, ++val2)
             sum += *val1 * *val2;
 
           sum += rho_[classifier_idx++];
 
           *scores_iter++ = static_cast<float>(sum);
-          ++(cur_votes[onnxruntime::narrow<size_t>(sum > 0 ? i : j)]);
+          ++(cur_votes[sum > 0 ? i : j]);
         }
       }
     }
@@ -346,23 +346,23 @@
                          &classifier_scores_data, num_classifiers, &votes_data, &Y,
                          num_scores_per_batch, write_additional_scores](ptrdiff_t idx) {
     int n = SafeInt<int32_t>(idx);  // convert to a usable sized type
-    auto cur_scores = final_scores.subspan(n * SafeInt<size_t>(final_scores_per_batch), onnxruntime::narrow<size_t>(final_scores_per_batch));
+    auto cur_scores = final_scores.subspan(n * SafeInt<size_t>(final_scores_per_batch), final_scores_per_batch);
 
     if (mode_ == SVM_TYPE::SVM_SVC && have_proba) {
-      auto probsp2 = gsl::make_span<float>(probsp2_data.data() + (n * class_count_squared), onnxruntime::narrow<size_t>(class_count_squared));
+      auto probsp2 = gsl::make_span<float>(probsp2_data.data() + (n * class_count_squared), class_count_squared);
 
       float* classifier_scores = classifier_scores_data.data() + (n * num_classifiers);
 
       size_t index = 0;
-      for (int64_t i = 0; i < class_count_ - 1; ++i) {
-        int64_t p1 = i * class_count_ + i + 1;
-        int64_t p2 = (i + 1) * class_count_ + i;
-        for (int64_t j = i + 1; j < class_count_; ++j, ++index) {
+      for (size_t i = 0; i < class_count_ - 1; ++i) {
+        size_t p1 = i * class_count_ + i + 1;
+        size_t p2 = (i + 1) * class_count_ + i;
+        for (size_t j = i + 1; j < class_count_; ++j, ++index) {
           float val1 = sigmoid_probability(classifier_scores[index], proba_[index], probb_[index]);
           float val2 = std::max(val1, 1.0e-7f);
           val2 = std::min(val2, 1 - 1.0e-7f);
-          probsp2[onnxruntime::narrow<size_t>(p1)] = val2;
-          probsp2[onnxruntime::narrow<size_t>(p2)] = 1 - val2;
+          probsp2[p1] = val2;
+          probsp2[p2] = 1 - val2;
           ++p1;
           p2 += class_count_;
         }

onnxruntime/core/providers/cpu/ml/svmclassifier.h

@@ -120,9 +120,9 @@ class SVMClassifier final : public OpKernel, private SVMCommon {
   Status ComputeImpl(OpKernelContext& ctx, gsl::span<const float> x_data, const TensorShape& x_shape) const;
 
   bool weights_are_all_positive_;
-  ptrdiff_t feature_count_;
-  ptrdiff_t class_count_;
-  ptrdiff_t vector_count_;
+  size_t feature_count_;
+  size_t class_count_;
+  size_t vector_count_;
   bool using_strings_;
   std::vector<int64_t> vectors_per_class_;
   std::vector<int64_t> starting_vector_;

onnxruntime/test/providers/cpu/ml/svmclassifier_test.cc

@@ -348,5 +348,47 @@ TEST(MLOpTest, SVMClassifierUndersizedProba) {
   test.Run(OpTester::ExpectResult::kExpectFailure, "prob_a attribute size");
 }
 
+TEST(MLOpTest, SVMClassifierLinearUndersizedVectorPerClass) {
+  OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain);
+
+  std::vector<float> coefficients = {0.766398549079895f, 0.0871576070785522f, 0.110420741140842f,
+                                     -0.963976919651031f};
+  std::vector<float> support_vectors = {4.80000019073486f, 3.40000009536743f, 1.89999997615814f,
+                                        5.f, 3.f, 1.60000002384186f,
+                                        4.5f, 2.29999995231628f, 1.29999995231628f,
+                                        5.09999990463257f, 2.5f, 3.f};
+  std::vector<float> rho = {2.23510527610779f};
+  std::vector<float> kernel_params = {0.122462183237076f, 0.f, 3.f};  // gamma, coef0, degree
+  std::vector<int64_t> classes = {0, 1};
+  std::vector<int64_t> vectors_per_class = {3};
+
+  std::vector<float> X = {5.1f, 3.5f, 1.4f,
+                          4.9f, 3.f, 1.4f,
+                          4.7f, 3.2f, 1.3f,
+                          4.6f, 3.1f, 1.5f,
+                          5.f, 3.6f, 1.4f};
+  std::vector<float> scores_predictions = {-1.5556798f, 1.5556798f,
+                                           -1.2610321f, 1.2610321f,
+                                           -1.5795376f, 1.5795376f,
+                                           -1.3083477f, 1.3083477f,
+                                           -1.6572928f, 1.6572928f};
+
+  std::vector<int64_t> class_predictions = {0, 0, 0, 0, 0};
+
+  test.AddAttribute("kernel_type", std::string("LINEAR"));
+  test.AddAttribute("coefficients", coefficients);
+  test.AddAttribute("support_vectors", support_vectors);
+  test.AddAttribute("vectors_per_class", vectors_per_class);
+  test.AddAttribute("rho", rho);
+  test.AddAttribute("kernel_params", kernel_params);
+  test.AddAttribute("classlabels_ints", classes);
+
+  test.AddInput<float>("X", {5, 3}, X);
+  test.AddOutput<int64_t>("Y", {5}, class_predictions);
+  test.AddOutput<float>("Z", {5, 2}, scores_predictions);
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "Mismatch between class_labels and vector_per_class dimensions.");
+}
+
 }  // namespace test
 }  // namespace onnxruntime