Prepare RF bins and objectives for weights support

cpp/src/decisiontree/batched-levelalgo/bins.cuh

@@ -3,61 +3,156 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
+
 #include <raft/util/cuda_utils.cuh>
 
 namespace ML {
 namespace DT {
 
-struct CountBin {
-  // double covers both the unweighted count path and the future weighted-count
-  // path with one bin type; 32-bit int would overflow on large weighted counts.
-  double x;
-  CountBin(CountBin const&) = default;
-  HDI CountBin(double x_) : x(x_) {}
-  HDI CountBin() : x(0.0) {}
+using BinCountT = unsigned long long int;
+static_assert(sizeof(BinCountT) == 8, "BinCountT must be 64 bits");
+
+struct ClassificationBin {
+  BinCountT count;
 
-  DI static void IncrementHistogram(CountBin* hist, int n_bins, int b, int label)
+  ClassificationBin(ClassificationBin const&) = default;
+  HDI ClassificationBin(BinCountT count_) : count(count_) {}
+  HDI ClassificationBin() : count(0) {}
+
+  DI static void IncrementHistogram(ClassificationBin* hist, int n_bins, int b, int label)
   {
     auto offset = label * n_bins + b;
-    CountBin::AtomicAdd(hist + offset, {1.0});
+    ClassificationBin::AtomicAdd(hist + offset, {1});
   }
-  DI static void AtomicAdd(CountBin* address, CountBin val) { atomicAdd(&address->x, val.x); }
-  HDI CountBin& operator+=(const CountBin& b)
+  DI static void AtomicAdd(ClassificationBin* address, ClassificationBin val)
   {
-    x += b.x;
+    atomicAdd(&address->count, val.count);
+  }
+  HDI BinCountT Count() const { return count; }
+  HDI double Weight() const { return static_cast<double>(count); }
+  HDI ClassificationBin& operator+=(const ClassificationBin& b)
+  {
+    count += b.count;
     return *this;
   }
-  HDI CountBin operator+(CountBin b) const
+  HDI ClassificationBin operator+(ClassificationBin b) const
   {
     b += *this;
     return b;
   }
 };
 
-struct AggregateBin {
+struct WeightedClassificationBin {
+  BinCountT count;
+  double weight;
+
+  WeightedClassificationBin(WeightedClassificationBin const&) = default;
+  HDI WeightedClassificationBin(BinCountT count_, double weight_) : count(count_), weight(weight_)
+  {
+  }
+  HDI WeightedClassificationBin() : count(0), weight(0.0) {}
+
+  DI static void IncrementHistogram(WeightedClassificationBin* hist, int n_bins, int b, int label)
+  {
+    WeightedClassificationBin::IncrementHistogram(hist, n_bins, b, label, 1.0);
+  }
+  DI static void IncrementHistogram(
+    WeightedClassificationBin* hist, int n_bins, int b, int label, double weight)
+  {
+    auto offset = label * n_bins + b;
+    WeightedClassificationBin::AtomicAdd(hist + offset, {1, weight});
+  }
+  DI static void AtomicAdd(WeightedClassificationBin* address, WeightedClassificationBin val)
+  {
+    atomicAdd(&address->count, val.count);
+    atomicAdd(&address->weight, val.weight);
+  }
+  HDI BinCountT Count() const { return count; }
+  HDI double Weight() const { return weight; }
+  HDI WeightedClassificationBin& operator+=(const WeightedClassificationBin& b)
+  {
+    count += b.count;
+    weight += b.weight;
+    return *this;
+  }
+  HDI WeightedClassificationBin operator+(WeightedClassificationBin b) const
+  {
+    b += *this;
+    return b;
+  }
+};
+
+struct RegressionBin {
+  double label_sum;
+  BinCountT count;
+
+  RegressionBin(RegressionBin const&) = default;
+  HDI RegressionBin() : label_sum(0.0), count(0) {}
+  HDI RegressionBin(double label_sum, BinCountT count) : label_sum(label_sum), count(count) {}
+
+  DI static void IncrementHistogram(RegressionBin* hist, int n_bins, int b, double label)
+  {
+    RegressionBin::AtomicAdd(hist + b, {label, 1});
+  }
+  DI static void AtomicAdd(RegressionBin* address, RegressionBin val)
+  {
+    atomicAdd(&address->label_sum, val.label_sum);
+    atomicAdd(&address->count, val.count);
+  }
+  HDI double LabelSum() const { return label_sum; }
+  HDI BinCountT Count() const { return count; }
+  HDI double Weight() const { return static_cast<double>(count); }
+  HDI RegressionBin& operator+=(const RegressionBin& b)
+  {
+    label_sum += b.label_sum;
+    count += b.count;
+    return *this;
+  }
+  HDI RegressionBin operator+(RegressionBin b) const
+  {
+    b += *this;
+    return b;
+  }
+};
+
+struct WeightedRegressionBin {
   double label_sum;
-  int count;
+  BinCountT count;
+  double weight;
 
-  AggregateBin(AggregateBin const&) = default;
-  HDI AggregateBin() : label_sum(0.0), count(0) {}
-  HDI AggregateBin(double label_sum, int count) : label_sum(label_sum), count(count) {}
+  WeightedRegressionBin(WeightedRegressionBin const&) = default;
+  HDI WeightedRegressionBin() : label_sum(0.0), count(0), weight(0.0) {}
+  HDI WeightedRegressionBin(double label_sum, BinCountT count, double weight)
+    : label_sum(label_sum), count(count), weight(weight)
+  {
+  }
 
-  DI static void IncrementHistogram(AggregateBin* hist, int n_bins, int b, double label)
+  DI static void IncrementHistogram(WeightedRegressionBin* hist, int n_bins, int b, double label)
+  {
+    WeightedRegressionBin::IncrementHistogram(hist, n_bins, b, label, 1.0);
+  }
+  DI static void IncrementHistogram(
+    WeightedRegressionBin* hist, int n_bins, int b, double label, double weight)
   {
-    AggregateBin::AtomicAdd(hist + b, {label, 1});
+    WeightedRegressionBin::AtomicAdd(hist + b, {label * weight, 1, weight});
   }
-  DI static void AtomicAdd(AggregateBin* address, AggregateBin val)
+  DI static void AtomicAdd(WeightedRegressionBin* address, WeightedRegressionBin val)
   {
     atomicAdd(&address->label_sum, val.label_sum);
     atomicAdd(&address->count, val.count);
+    atomicAdd(&address->weight, val.weight);
   }
-  HDI AggregateBin& operator+=(const AggregateBin& b)
+  HDI double LabelSum() const { return label_sum; }
+  HDI BinCountT Count() const { return count; }
+  HDI double Weight() const { return weight; }
+  HDI WeightedRegressionBin& operator+=(const WeightedRegressionBin& b)
   {
     label_sum += b.label_sum;
     count += b.count;
+    weight += b.weight;
     return *this;
   }
-  HDI AggregateBin operator+(AggregateBin b) const
+  HDI WeightedRegressionBin operator+(WeightedRegressionBin b) const
   {
     b += *this;
     return b;

cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu

@@ -14,7 +14,7 @@ using _DataT      = double;
 using _LabelT     = int;
 using _IdxT       = int;
 using _ObjectiveT = ClassificationObjectiveFunction<_DataT, _LabelT, _IdxT>;
-using _BinT       = CountBin;
+using _BinT       = ClassificationBin;
 using _DatasetT   = Dataset<_DataT, _LabelT, _IdxT>;
 using _NodeT      = SparseTreeNode<_DataT, _LabelT, _IdxT>;
 }  // namespace DT

cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu

@@ -14,7 +14,7 @@ using _DataT      = float;
 using _LabelT     = int;
 using _IdxT       = int;
 using _ObjectiveT = ClassificationObjectiveFunction<_DataT, _LabelT, _IdxT>;
-using _BinT       = CountBin;
+using _BinT       = ClassificationBin;
 using _DatasetT   = Dataset<_DataT, _LabelT, _IdxT>;
 using _NodeT      = SparseTreeNode<_DataT, _LabelT, _IdxT>;
 }  // namespace DT

cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu

@@ -14,7 +14,7 @@ using _DataT      = double;
 using _LabelT     = double;
 using _IdxT       = int;
 using _ObjectiveT = RegressionObjectiveFunction<_DataT, _LabelT, _IdxT>;
-using _BinT       = AggregateBin;
+using _BinT       = RegressionBin;
 using _DatasetT   = Dataset<_DataT, _LabelT, _IdxT>;
 using _NodeT      = SparseTreeNode<_DataT, _LabelT, _IdxT>;
 }  // namespace DT

cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu

@@ -14,7 +14,7 @@ using _DataT      = float;
 using _LabelT     = float;
 using _IdxT       = int;
 using _ObjectiveT = RegressionObjectiveFunction<_DataT, _LabelT, _IdxT>;
-using _BinT       = AggregateBin;
+using _BinT       = RegressionBin;
 using _DatasetT   = Dataset<_DataT, _LabelT, _IdxT>;
 using _NodeT      = SparseTreeNode<_DataT, _LabelT, _IdxT>;
 }  // namespace DT