use int64_t for num_class * num_data

include/LightGBM/boosting.h

@@ -94,7 +94,7 @@ public:
   * \param out_len length of returned score
   * \return training score
   */
-  virtual const score_t* GetTrainingScore(data_size_t* out_len) = 0;
+  virtual const score_t* GetTrainingScore(int64_t* out_len) = 0;
 
   /*!
   * \brief Get prediction result at data_idx data
@@ -102,7 +102,7 @@ public:
   * \param result used to store prediction result, should allocate memory before call this function
   * \param out_len lenght of returned score
   */
-  virtual void GetPredictAt(int data_idx, score_t* result, data_size_t* out_len) = 0;
+  virtual void GetPredictAt(int data_idx, score_t* result, int64_t* out_len) = 0;
 
   /*!
   * \brief Prediction for one record, not sigmoid transform

src/boosting/dart.hpp

@@ -63,13 +63,13 @@ public:
   * \param out_len length of returned score
   * \return training score
   */
-  const score_t* GetTrainingScore(data_size_t* out_len) override {
+  const score_t* GetTrainingScore(int64_t* out_len) override {
     if (!is_update_score_cur_iter_) {
       // only drop one time in one iteration
       DroppingTrees();
       is_update_score_cur_iter_ = true;
     }
-    *out_len = train_score_updater_->num_data() * num_class_;
+    *out_len = static_cast<int64_t>(train_score_updater_->num_data()) * num_class_;
     return train_score_updater_->score();
   }
 

src/boosting/gbdt.cpp

@@ -83,8 +83,9 @@ void GBDT::ResetTrainingData(const BoostingConfig* config, const Dataset* train_
     num_data_ = train_data->num_data();
     // create buffer for gradients and hessians
     if (object_function_ != nullptr) {
-      gradients_.resize(num_data_ * num_class_);
-      hessians_.resize(num_data_ * num_class_);
+      size_t total_size = static_cast<size_t>(num_data_) * num_class_;
+      gradients_.resize(total_size);
+      hessians_.resize(total_size);
     }
     // get max feature index
     max_feature_idx_ = train_data->num_total_features() - 1;
@@ -372,12 +373,12 @@ std::vector<double> GBDT::GetEvalAt(int data_idx) const {
 }
 
 /*! \brief Get training scores result */
-const score_t* GBDT::GetTrainingScore(data_size_t* out_len) {
-  *out_len = train_score_updater_->num_data() * num_class_;
+const score_t* GBDT::GetTrainingScore(int64_t* out_len) {
+  *out_len = static_cast<int64_t>(train_score_updater_->num_data()) * num_class_;
   return train_score_updater_->score();
 }
 
-void GBDT::GetPredictAt(int data_idx, score_t* out_result, data_size_t* out_len) {
+void GBDT::GetPredictAt(int data_idx, score_t* out_result, int64_t* out_len) {
   CHECK(data_idx >= 0 && data_idx <= static_cast<int>(valid_metrics_.size()));
 
   const score_t* raw_scores = nullptr;
@@ -389,7 +390,7 @@ void GBDT::GetPredictAt(int data_idx, score_t* out_result, data_size_t* out_len)
     auto used_idx = data_idx - 1;
     raw_scores = valid_score_updater_[used_idx]->score();
     num_data = valid_score_updater_[used_idx]->num_data();
-    *out_len = num_data * num_class_;
+    *out_len = static_cast<int64_t>(num_data) * num_class_;
   }
   if (num_class_ > 1) {
 #pragma omp parallel for schedule(static)
@@ -422,7 +423,7 @@ void GBDT::Boosting() {
     Log::Fatal("No object function provided");
   }
   // objective function will calculate gradients and hessians
-  int num_score = 0;
+  int64_t num_score = 0;
   object_function_->
     GetGradients(GetTrainingScore(&num_score), gradients_.data(), hessians_.data());
 }

src/boosting/gbdt.h

@@ -105,7 +105,7 @@ public:
   * \param out_len length of returned score
   * \return training score
   */
-  virtual const score_t* GetTrainingScore(data_size_t* out_len) override;
+  virtual const score_t* GetTrainingScore(int64_t* out_len) override;
 
   /*!
   * \brief Get prediction result at data_idx data
@@ -113,7 +113,7 @@ public:
   * \param result used to store prediction result, should allocate memory before call this function
   * \param out_len lenght of returned score
   */
-  void GetPredictAt(int data_idx, score_t* out_result, data_size_t* out_len) override;
+  void GetPredictAt(int data_idx, score_t* out_result, int64_t* out_len) override;
 
   /*!
   * \brief Prediction for one record without sigmoid transformation

src/boosting/score_updater.hpp

@@ -20,13 +20,14 @@ public:
   */
   ScoreUpdater(const Dataset* data, int num_class) : data_(data) {
     num_data_ = data->num_data();
-    score_ = std::vector<score_t>(num_data_ * num_class);
+    size_t total_size = static_cast<size_t>(num_data_) * num_class;
+    score_.resize(total_size);
     // default start score is zero
     std::fill(score_.begin(), score_.end(), 0.0f);
     const float* init_score = data->metadata().init_score();
     // if exists initial score, will start from it
     if (init_score != nullptr) {
-      for (data_size_t i = 0; i < num_data_ * num_class; ++i) {
+      for (size_t i = 0; i < total_size; ++i) {
         score_[i] = init_score[i];
       }
     }

src/c_api.cpp

@@ -169,7 +169,7 @@ public:
     predictor_.reset(new Predictor(boosting_.get(), is_raw_score, is_predict_leaf));
   }
 
-  void GetPredictAt(int data_idx, score_t* out_result, data_size_t* out_len) {
+  void GetPredictAt(int data_idx, score_t* out_result, int64_t* out_len) {
     boosting_->GetPredictAt(data_idx, out_result, out_len);
   }
 
@@ -681,9 +681,7 @@ DllExport int LGBM_BoosterGetPredict(BoosterHandle handle,
   float* out_result) {
   API_BEGIN();
   Booster* ref_booster = reinterpret_cast<Booster*>(handle);
-  int len = 0;
-  ref_booster->GetPredictAt(data_idx, out_result, &len);
-  *out_len = static_cast<int64_t>(len);
+  ref_booster->GetPredictAt(data_idx, out_result, out_len);
   API_END();
 }
 

src/metric/multiclass_metric.hpp

@@ -56,7 +56,8 @@ public:
       for (data_size_t i = 0; i < num_data_; ++i) {
         std::vector<double> rec(num_class_);
         for (int k = 0; k < num_class_; ++k) {
-          rec[k] = static_cast<double>(score[k * num_data_ + i]);
+          size_t idx = static_cast<size_t>(num_data_) * k + i;
+          rec[k] = static_cast<double>(score[idx]);
         }
         // add loss
         sum_loss += PointWiseLossCalculator::LossOnPoint(label_[i], rec);
@@ -66,7 +67,8 @@ public:
       for (data_size_t i = 0; i < num_data_; ++i) {
         std::vector<double> rec(num_class_);
         for (int k = 0; k < num_class_; ++k) {
-          rec[k] = static_cast<double>(score[k * num_data_ + i]);
+          size_t idx = static_cast<size_t>(num_data_) * k + i;
+          rec[k] = static_cast<double>(score[idx]);
         }
         // add loss
         sum_loss += PointWiseLossCalculator::LossOnPoint(label_[i], rec) * weights_[i];

src/objective/multiclass_objective.hpp

@@ -38,17 +38,19 @@ public:
       for (data_size_t i = 0; i < num_data_; ++i) {
         std::vector<double> rec(num_class_);
         for (int k = 0; k < num_class_; ++k){
-          rec[k] = static_cast<double>(score[k * num_data_ + i]);
+          size_t idx = static_cast<size_t>(num_data_) * k + i;
+          rec[k] = static_cast<double>(score[idx]);
         }
         Common::Softmax(&rec);
         for (int k = 0; k < num_class_; ++k) {
           score_t p = static_cast<score_t>(rec[k]);
+          size_t idx = static_cast<size_t>(num_data_) * k + i;
           if (label_int_[i] == k) {
-            gradients[k * num_data_ + i] = p - 1.0f;
+            gradients[idx] = p - 1.0f;
           } else {
-            gradients[k * num_data_ + i] = p;
+            gradients[idx] = p;
           }
-          hessians[k * num_data_ + i] = 2.0f * p * (1.0f - p);
+          hessians[idx] = 2.0f * p * (1.0f - p);
         }
       }
     } else {
@@ -56,17 +58,19 @@ public:
       for (data_size_t i = 0; i < num_data_; ++i) {
         std::vector<double> rec(num_class_);
         for (int k = 0; k < num_class_; ++k){
-          rec[k] = static_cast<double>(score[k * num_data_ + i]);
+          size_t idx = static_cast<size_t>(num_data_) * k + i;
+          rec[k] = static_cast<double>(score[idx]);
         }
         Common::Softmax(&rec);
         for (int k = 0; k < num_class_; ++k) {
           score_t p = static_cast<score_t>(rec[k]);
+          size_t idx = static_cast<size_t>(num_data_) * k + i;
           if (label_int_[i] == k) {
-            gradients[k * num_data_ + i] = (p - 1.0f) * weights_[i];
+            gradients[idx] = (p - 1.0f) * weights_[i];
           } else {
-            gradients[k * num_data_ + i] = p * weights_[i];
+            gradients[idx] = p * weights_[i];
           }
-          hessians[k * num_data_ + i] = 2.0f * p * (1.0f - p) * weights_[i];
+          hessians[idx] = 2.0f * p * (1.0f - p) * weights_[i];
         }
       }
     }