Support row-major RandomForest training inputs

cpp/bench/sg/dataset.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -223,6 +223,7 @@ namespace {
 std::ostream& operator<<(std::ostream& os, const DatasetParams& d)
 {
   os << "/" << d.nrows << "x" << d.ncols;
+  os << "/" << (d.rowMajor ? "row" : "col");
   return os;
 }
 }  // namespace

cpp/bench/sg/rf_classifier.cu

@@ -45,19 +45,21 @@ class RFClassifier : public BlobsFixture<D> {
   void runBenchmark(::benchmark::State& state) override
   {
     using MLCommon::Bench::CudaEventTimer;
-    if (this->params.rowMajor) {
-      state.SkipWithError("RFClassifier only supports col-major inputs");
-    }
     this->loopOnState(state, [this]() {
       auto* mPtr = &model.model;
+      mPtr->trees.clear();
       fit(*this->handle,
           mPtr,
           this->data.X.data(),
           this->params.nrows,
           this->params.ncols,
           this->data.y.data(),
           this->params.nclasses,
-          rfParams);
+          rfParams,
+          rapids_logger::level_enum::info,
+          nullptr,
+          nullptr,
+          this->params.rowMajor);
       this->handle->sync_stream(this->stream);
     });
   }
@@ -119,6 +121,9 @@ std::vector<Params> getInputs()
     p.rf.tree_params.max_features    = 1.f / std::sqrt(float(cfg.ncols));
     for (auto max_depth : std::vector<int>({7, 9})) {
       p.rf.tree_params.max_depth = max_depth;
+      p.data.rowMajor            = false;
+      out.push_back(p);
+      p.data.rowMajor = true;
       out.push_back(p);
     }
   }

cpp/bench/sg/rf_regressor.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -45,18 +45,20 @@ class RFRegressor : public RegressionFixture<D> {
   void runBenchmark(::benchmark::State& state) override
   {
     using MLCommon::Bench::CudaEventTimer;
-    if (this->params.rowMajor) {
-      state.SkipWithError("RFRegressor only supports col-major inputs");
-    }
     this->loopOnState(state, [this]() {
       auto* mPtr = &model.model;
+      mPtr->trees.clear();
       fit(*this->handle,
           mPtr,
-          this->data.X,
+          this->data.X.data(),
           this->params.nrows,
           this->params.ncols,
-          this->data.y,
-          rfParams);
+          this->data.y.data(),
+          rfParams,
+          rapids_logger::level_enum::info,
+          nullptr,
+          nullptr,
+          this->params.rowMajor);
       handle->sync_stream(this->stream);
     });
   }
@@ -107,6 +109,9 @@ std::vector<RegParams> getInputs()
     p.rf.tree_params.max_features = 1.f;
     for (auto max_depth : std::vector<int>({7, 11, 15})) {
       p.rf.tree_params.max_depth = max_depth;
+      p.data.rowMajor            = false;
+      out.push_back(p);
+      p.data.rowMajor = true;
       out.push_back(p);
     }
   }

cpp/include/cuml/ensemble/randomforest.hpp

@@ -152,7 +152,8 @@ void fit(const raft::handle_t& user_handle,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
          bool* bootstrap_masks               = nullptr,
-         const double* sample_weight         = nullptr);
+         const double* sample_weight         = nullptr,
+         bool input_row_major                = false);
 void fit(const raft::handle_t& user_handle,
          RandomForestClassifierD* forest,
          double* input,
@@ -163,8 +164,27 @@ void fit(const raft::handle_t& user_handle,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
          bool* bootstrap_masks               = nullptr,
-         const double* sample_weight         = nullptr);
+         const double* sample_weight         = nullptr,
+         bool input_row_major                = false);
 
+/**
+ * @brief Train a random forest classifier and export it as a Treelite model.
+ *
+ * @param[in] user_handle RAFT handle for stream and allocator resources.
+ * @param[out] model Treelite model handle populated by training.
+ * @param[in] input Training data, column-major by default or row-major when
+ *   input_row_major is true.
+ * @param[in] n_rows Number of rows in input.
+ * @param[in] n_cols Number of columns in input.
+ * @param[in] labels Training labels.
+ * @param[in] n_unique_labels Number of unique classes in labels.
+ * @param[in] rf_params Random forest training parameters.
+ * @param[in] bootstrap_masks Optional bootstrap masks.
+ * @param[out] feature_importances Output feature importances.
+ * @param[in] verbosity Logging verbosity.
+ * @param[in] sample_weight Optional per-row sample weights.
+ * @param[in] input_row_major Whether input is row-major instead of column-major.
+ */
 template <typename T, typename L>
 void fit_treelite(const raft::handle_t& user_handle,
                   TreeliteModelHandle* model,
@@ -177,7 +197,8 @@ void fit_treelite(const raft::handle_t& user_handle,
                   bool* bootstrap_masks,
                   T* feature_importances,
                   rapids_logger::level_enum verbosity,
-                  const double* sample_weight = nullptr);
+                  const double* sample_weight = nullptr,
+                  bool input_row_major        = false);
 
 void predict(const raft::handle_t& user_handle,
              const RandomForestClassifierF* forest,
@@ -236,7 +257,8 @@ void fit(const raft::handle_t& user_handle,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
          bool* bootstrap_masks               = nullptr,
-         const double* sample_weight         = nullptr);
+         const double* sample_weight         = nullptr,
+         bool input_row_major                = false);
 void fit(const raft::handle_t& user_handle,
          RandomForestRegressorD* forest,
          double* input,
@@ -246,8 +268,26 @@ void fit(const raft::handle_t& user_handle,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
          bool* bootstrap_masks               = nullptr,
-         const double* sample_weight         = nullptr);
+         const double* sample_weight         = nullptr,
+         bool input_row_major                = false);
 
+/**
+ * @brief Train a random forest regressor and export it as a Treelite model.
+ *
+ * @param[in] user_handle RAFT handle for stream and allocator resources.
+ * @param[out] model Treelite model handle populated by training.
+ * @param[in] input Training data, column-major by default or row-major when
+ *   input_row_major is true.
+ * @param[in] n_rows Number of rows in input.
+ * @param[in] n_cols Number of columns in input.
+ * @param[in] labels Training labels.
+ * @param[in] rf_params Random forest training parameters.
+ * @param[in] bootstrap_masks Optional bootstrap masks.
+ * @param[out] feature_importances Output feature importances.
+ * @param[in] verbosity Logging verbosity.
+ * @param[in] sample_weight Optional per-row sample weights.
+ * @param[in] input_row_major Whether input is row-major instead of column-major.
+ */
 template <typename T, typename L>
 void fit_treelite(const raft::handle_t& user_handle,
                   TreeliteModelHandle* model,
@@ -259,7 +299,8 @@ void fit_treelite(const raft::handle_t& user_handle,
                   bool* bootstrap_masks,
                   T* feature_importances,
                   rapids_logger::level_enum verbosity,
-                  const double* sample_weight = nullptr);
+                  const double* sample_weight = nullptr,
+                  bool input_row_major        = false);
 
 void predict(const raft::handle_t& user_handle,
              const RandomForestRegressorF* forest,

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

@@ -212,7 +212,8 @@ struct Builder {
           IdxT n_cols,
           rmm::device_uvector<IdxT>* row_ids,
           IdxT n_classes,
-          const QuantilesT& q)
+          const QuantilesT& q,
+          bool row_major = false)
     : handle(handle),
       builder_stream(s),
       treeid(treeid),
@@ -223,6 +224,8 @@ struct Builder {
               sample_weight,
               n_rows,
               n_cols,
+              row_major ? n_cols : IdxT{1},
+              row_major ? IdxT{1} : n_rows,
               int(row_ids->size()),
               max(1, IdxT(params.max_features * n_cols)),
               row_ids->data(),
@@ -397,10 +400,10 @@ struct Builder {
     RAFT_CUDA_TRY(cudaMemsetAsync(n_nodes, 0, sizeof(IdxT), builder_stream));
 
     const IdxT original_n_sampled_cols = dataset.n_sampled_cols;
-    ASSERT(original_n_sampled_cols > 0 && original_n_sampled_cols <= dataset.N,
+    ASSERT(original_n_sampled_cols > 0 && original_n_sampled_cols <= dataset.n_cols,
            "n_sampled_cols must be in [1, n_cols]");
     const std::size_t max_sampling_rounds =
-      std::size_t((dataset.N + original_n_sampled_cols - 1) / original_n_sampled_cols);
+      std::size_t((dataset.n_cols + original_n_sampled_cols - 1) / original_n_sampled_cols);
     struct HostSplit {
       DataT quesval;
       IdxT colid;
@@ -429,8 +432,9 @@ struct Builder {
     // Match sklearn's behavior of searching beyond max_features when the
     // sampled features do not yield a valid split.
     for (std::size_t round = 0; !active_items.empty() && round < max_sampling_rounds; ++round) {
-      IdxT sample_offset     = IdxT(round) * original_n_sampled_cols;
-      dataset.n_sampled_cols = std::min(original_n_sampled_cols, dataset.N - sample_offset);
+      IdxT sample_offset = IdxT(round) * original_n_sampled_cols;
+      dataset.n_sampled_cols =
+        std::min(original_n_sampled_cols, static_cast<IdxT>(dataset.n_cols) - sample_offset);
       computeBestSplits(active_items, seed, sample_offset);
 
       std::vector<NodeWorkItem> retry_items;
@@ -515,7 +519,7 @@ struct Builder {
                           treeid,
                           sampling_seed,
                           sample_offset,
-                          dataset.N,
+                          static_cast<IdxT>(dataset.n_cols),
                           dataset.n_sampled_cols,
                           builder_stream);
     RAFT_CUDA_TRY(cudaPeekAtLastError());

cpp/src/decisiontree/batched-levelalgo/dataset.h

@@ -5,21 +5,29 @@
 
 #pragma once
 
+#include <raft/util/cuda_utils.cuh>
+
+#include <cstdint>
+
 namespace ML {
 namespace DT {
 
 template <typename DataT, typename LabelT, typename IdxT>
 struct Dataset {
-  /** input dataset (assumed to be col-major) */
+  /** input dataset */
   const DataT* data;
   /** input labels */
   const LabelT* labels;
   /** optional input sample weights */
   const double* sample_weight;
   /** total rows in dataset */
-  IdxT M;
+  std::int64_t n_rows;
   /** total cols in dataset */
-  IdxT N;
+  std::int64_t n_cols;
+  /** row stride in input data elements */
+  std::int64_t row_stride;
+  /** column stride in input data elements */
+  std::int64_t col_stride;
   /** total sampled rows in dataset */
   IdxT n_sampled_rows;
   /** total sampled cols in dataset */
@@ -28,6 +36,12 @@ struct Dataset {
   IdxT* row_ids;
   /** Number of classes or regression outputs*/
   IdxT num_outputs;
+
+  HDI DataT value(IdxT row, IdxT col) const
+  {
+    return data[static_cast<std::int64_t>(row) * row_stride +
+                static_cast<std::int64_t>(col) * col_stride];
+  }
 };
 
 }  // namespace DT

cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh

@@ -144,8 +144,7 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf,
     }
 
     const auto row       = dataset.row_ids[work_item.instances.begin + range_pos];
-    const auto col_idx   = std::size_t(split.colid) * dataset.M + row;
-    const auto goes_left = dataset.data[col_idx] <= split.quesval;
+    const auto goes_left = dataset.value(row, split.colid) <= split.quesval;
     return NodeSplitPartitionState<IdxT>{goes_left ? IdxT(1) : IdxT(0), true, goes_left};
   };
 
@@ -328,11 +327,10 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms,
   // compute pdf histogram for all bins for all classes
 
   // Must be 64 bit - can easily grow larger than a 32 bit int
-  std::size_t col_offset = std::size_t(col) * dataset.M;
   for (auto i = range_start + tid; i < end; i += stride) {
     // each thread works over a data point and strides to the next
     auto row   = dataset.row_ids[i];
-    auto data  = dataset.data[row + col_offset];
+    auto data  = dataset.value(row, col);
     auto label = dataset.labels[row];
 
     // `start` is lowest index such that data <= quantiles_for_split[start]

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

@@ -46,6 +46,8 @@ static __global__ void sampleOwnedColumnsKernel(T* out,
                                                 int sample_count,
                                                 int rank,
                                                 int n_rows,
+                                                int n_cols,
+                                                bool row_major,
                                                 std::uint64_t seed)
 {
   int col        = blockIdx.x;
@@ -70,7 +72,8 @@ static __global__ void sampleOwnedColumnsKernel(T* out,
   if (sample_rank == rank) {
     int local_row = static_cast<int>(global_row - local_begin);
     out[static_cast<std::size_t>(col) * sample_count + sample_idx] =
-      data[static_cast<int64_t>(col) * n_rows + local_row];
+      row_major ? data[static_cast<std::size_t>(local_row) * n_cols + col]
+                : data[static_cast<std::size_t>(col) * n_rows + local_row];
   }
 }
 
@@ -125,7 +128,9 @@ struct QuantileResult {
  *
  * @tparam T Floating-point input type.
  * @param handle RAFT handle used for stream and resource access.
- * @param data Column-major input matrix with shape `[n_cols, n_rows]`.
+ * @param data Input matrix. When `row_major` is false, data is column-major with
+ * shape `[n_cols, n_rows]`; when `row_major` is true, data is row-major with
+ * shape `[n_rows, n_cols]`.
  * @param max_n_bins Maximum number of quantile candidates to retain per feature.
  * @param n_rows Number of local rows in `data` for this rank.
  * @param n_cols Number of columns in `data`.
@@ -135,6 +140,7 @@ struct QuantileResult {
  * `max_n_bins`, rank error decreases like O(1 / sqrt(oversampling_factor)), so
  * returns from increasing this are strongly diminishing.
  * @param seed User seed for deterministic sampling.
+ * @param row_major Whether `data` is row-major instead of column-major.
  * @return Quantile metadata and owning device buffers for quantile values and bin counts.
  */
 template <typename T>
@@ -144,7 +150,8 @@ CUML_EXPORT QuantileResult<T> computeQuantiles(const raft::handle_t& handle,
                                                int n_rows,
                                                int n_cols,
                                                int oversampling_factor = 4,
-                                               uint64_t seed           = uint64_t{0})
+                                               uint64_t seed           = uint64_t{0},
+                                               bool row_major          = false)
 {
   raft::common::nvtx::push_range("computeQuantiles");
   RAFT_EXPECTS(data != nullptr, "data pointer must not be null");
@@ -213,6 +220,8 @@ CUML_EXPORT QuantileResult<T> computeQuantiles(const raft::handle_t& handle,
                                                                           sample_count,
                                                                           rank,
                                                                           n_rows,
+                                                                          n_cols,
+                                                                          row_major,
                                                                           seed);
   RAFT_CUDA_TRY(cudaGetLastError());
   if (distributed) {