Simplify RF lower-bound and index types

cpp/include/cuml/tree/decisiontree.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -10,6 +10,7 @@
 
 #include <cuml/common/export.hpp>
 
+#include <cstdint>
 #include <string>
 #include <vector>
 

cpp/include/cuml/tree/flatnode.h

@@ -1,10 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
+#include <cstdint>
+
 // We want to define some functions as usable on device
 // But need to guard against this file being compiled by a host compiler
 #ifdef __CUDACC__
@@ -17,18 +19,20 @@
  * A node in Decision Tree.
  * @tparam T data type
  * @tparam L label type
- * @tparam IdxT type used for indexing operations
  */
-template <typename DataT, typename LabelT, typename IdxT = int>
+template <typename DataT, typename LabelT>
 struct SparseTreeNode {
  private:
-  IdxT colid            = 0;
-  DataT quesval         = DataT(0);
-  DataT best_metric_val = DataT(0);
-  IdxT left_child_id    = -1;
-  IdxT instance_count   = 0;
-  FLATNODE_HD SparseTreeNode(
-    IdxT colid, DataT quesval, DataT best_metric_val, int64_t left_child_id, IdxT instance_count)
+  std::int64_t colid          = 0;
+  DataT quesval               = DataT(0);
+  DataT best_metric_val       = DataT(0);
+  std::int64_t left_child_id  = -1;
+  std::int64_t instance_count = 0;
+  FLATNODE_HD SparseTreeNode(std::int64_t colid,
+                             DataT quesval,
+                             DataT best_metric_val,
+                             std::int64_t left_child_id,
+                             std::int64_t instance_count)
     : colid(colid),
       quesval(quesval),
       best_metric_val(best_metric_val),
@@ -38,22 +42,24 @@ struct SparseTreeNode {
   }
 
  public:
-  FLATNODE_HD IdxT ColumnId() const { return colid; }
+  FLATNODE_HD std::int64_t ColumnId() const { return colid; }
   FLATNODE_HD DataT QueryValue() const { return quesval; }
   FLATNODE_HD DataT BestMetric() const { return best_metric_val; }
-  FLATNODE_HD int64_t LeftChildId() const { return left_child_id; }
-  FLATNODE_HD int64_t RightChildId() const { return left_child_id + 1; }
-  FLATNODE_HD IdxT InstanceCount() const { return instance_count; }
+  FLATNODE_HD std::int64_t LeftChildId() const { return left_child_id; }
+  FLATNODE_HD std::int64_t RightChildId() const { return left_child_id + 1; }
+  FLATNODE_HD std::int64_t InstanceCount() const { return instance_count; }
 
-  FLATNODE_HD static SparseTreeNode CreateSplitNode(
-    IdxT colid, DataT quesval, DataT best_metric_val, int64_t left_child_id, IdxT instance_count)
+  FLATNODE_HD static SparseTreeNode CreateSplitNode(std::int64_t colid,
+                                                    DataT quesval,
+                                                    DataT best_metric_val,
+                                                    std::int64_t left_child_id,
+                                                    std::int64_t instance_count)
   {
-    return SparseTreeNode<DataT, LabelT>{
-      colid, quesval, best_metric_val, left_child_id, instance_count};
+    return SparseTreeNode{colid, quesval, best_metric_val, left_child_id, instance_count};
   }
-  FLATNODE_HD static SparseTreeNode CreateLeafNode(IdxT instance_count)
+  FLATNODE_HD static SparseTreeNode CreateLeafNode(std::int64_t instance_count)
   {
-    return SparseTreeNode<DataT, LabelT>{0, 0, 0, -1, instance_count};
+    return SparseTreeNode{0, 0, 0, -1, instance_count};
   }
   FLATNODE_HD bool IsLeaf() const { return left_child_id == -1; }
   bool operator==(const SparseTreeNode& other) const

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -12,7 +12,7 @@
 namespace ML {
 namespace DT {
 
-template <typename DataT, typename LabelT, typename IdxT>
+template <typename DataT, typename LabelT>
 struct Dataset {
   /** input dataset */
   const DataT* data;
@@ -29,18 +29,17 @@ struct Dataset {
   /** column stride in input data elements */
   std::int64_t col_stride;
   /** total sampled rows in dataset */
-  IdxT n_sampled_rows;
+  std::int64_t n_sampled_rows;
   /** total sampled cols in dataset */
-  IdxT n_sampled_cols;
+  std::int64_t n_sampled_cols;
   /** indices of sampled rows */
-  IdxT* row_ids;
+  std::int64_t* row_ids;
   /** Number of classes or regression outputs*/
-  IdxT num_outputs;
+  int num_outputs;
 
-  HDI DataT value(IdxT row, IdxT col) const
+  HDI DataT value(std::int64_t row, std::int64_t col) const
   {
-    return data[static_cast<std::int64_t>(row) * row_stride +
-                static_cast<std::int64_t>(col) * col_stride];
+    return data[row * row_stride + col * col_stride];
   }
 };
 

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

@@ -10,8 +10,10 @@
 #include "../quantiles.h"
 #include "../random_utils.cuh"
 
+#include <cuml/common/checked_arithmetic.hpp>
 #include <cuml/common/utils.hpp>
 
+#include <raft/core/error.hpp>
 #include <raft/linalg/unary_op.cuh>
 
 #include <cuda/iterator>
@@ -43,12 +45,11 @@ struct NodeWorkItem {
  * This struct has information about workload of a single threadblock of
  * computeSplit kernels of classification and regression
  */
-template <typename IdxT>
 struct WorkloadInfo {
-  IdxT nodeid;          // Node in the batch on which the threadblock needs to work
-  IdxT offset_blockid;  // Offset threadblock id among all the blocks that are
-                        // working on this node
-  IdxT num_blocks;      // Total number of blocks that are working on the node
+  std::int64_t nodeid;          // Node in the batch on which the threadblock needs to work
+  std::int64_t offset_blockid;  // Offset threadblock id among all the blocks that are
+                                // working on this node
+  std::int64_t num_blocks;      // Total number of blocks that are working on the node
 };
 
 struct SharedMemoryConfig {
@@ -63,44 +64,47 @@ DI OutT* alignPointer(InT dataset)
   return reinterpret_cast<OutT*>(raft::alignTo(reinterpret_cast<size_t>(dataset), sizeof(OutT)));
 }
 
-template <typename IdxT>
-void sample_features(IdxT* column_samples,
-                     const NodeWorkItem* work_items,
-                     size_t work_items_size,
-                     IdxT treeid,
-                     uint64_t seed,
-                     IdxT sample_offset,
-                     IdxT n,
-                     IdxT k,
-                     cudaStream_t stream)
+inline void sample_features(std::int64_t* column_samples,
+                            const NodeWorkItem* work_items,
+                            size_t work_items_size,
+                            std::int64_t treeid,
+                            uint64_t seed,
+                            std::int64_t sample_offset,
+                            std::int64_t n,
+                            std::int64_t k,
+                            cudaStream_t stream)
 {
-  auto n_column_samples = work_items_size * size_t(k);
-  auto counting         = thrust::make_counting_iterator<size_t>(0);
+  RAFT_EXPECTS(k >= 0, "k must be non-negative");
+  RAFT_EXPECTS(n >= k, "k must not exceed n");
+
+  auto sampled_cols     = ML::narrow_cast<std::size_t>(k);
+  auto n_column_samples = ML::checked_mul<std::size_t>(work_items_size, sampled_cols);
+  auto counting         = thrust::make_counting_iterator<std::size_t>(0);
 
   thrust::for_each(thrust::cuda::par.on(stream),
                    counting,
                    counting + n_column_samples,
-                   [=] __device__(size_t sample_idx) {
-                     auto node_idx     = sample_idx / size_t(k);
-                     IdxT column_index = static_cast<IdxT>(sample_idx % size_t(k));
+                   [=] __device__(std::size_t sample_idx) {
+                     auto node_idx     = sample_idx / sampled_cols;
+                     auto column_index = static_cast<std::int64_t>(sample_idx % sampled_cols);
 
-                     const uint32_t nodeid = work_items[node_idx].idx;
-                     uint32_t rng_seed     = fnv1a32_hash(seed, treeid, nodeid);
+                     auto nodeid       = work_items[node_idx].idx;
+                     uint32_t rng_seed = fnv1a32_hash(seed, treeid, nodeid);
 
-                     cuda::shuffle_iterator<IdxT> shuffled_features(
+                     cuda::shuffle_iterator<std::int64_t> shuffled_features(
                        n, cuda::std::minstd_rand(rng_seed), sample_offset);
                      column_samples[sample_idx] = shuffled_features[column_index];
                    });
 }
 
-template <typename DataT, typename LabelT, typename IdxT, int TPB>
-void launchNodeSplitKernel(const Dataset<DataT, LabelT, IdxT>& dataset,
+template <typename DataT, typename LabelT, int TPB>
+void launchNodeSplitKernel(const Dataset<DataT, LabelT>& dataset,
                            const NodeWorkItem* work_items,
-                           Split<DataT, IdxT>* splits,
-                           const WorkloadInfo<IdxT>* workload_info,
+                           Split<DataT>* splits,
+                           const WorkloadInfo* workload_info,
                            size_t n_blocks_dimx,
                            size_t n_work_items,
-                           IdxT* partition_row_ids,
+                           std::int64_t* partition_row_ids,
                            cudaStream_t builder_stream);
 
 template <typename DatasetT, typename NodeT, typename ObjectiveT, typename DataT>
@@ -112,49 +116,29 @@ void launchLeafKernel(ObjectiveT objective,
                       int batch_size,
                       size_t smem_size,
                       cudaStream_t builder_stream);
-// Returns the lowest index in `array` whose value is greater or equal to `element`.
-// Values outside the quantile range are clamped to the edge bins: values below the
-// first quantile return 0, and values above the last quantile return len - 1.
-template <typename DataT, typename IdxT>
-HDI IdxT lower_bound(DataT const* array, IdxT len, DataT element)
-{
-  IdxT start = 0;
-  IdxT end   = len - 1;
-  IdxT mid;
-  while (start < end) {
-    mid = (start + end) / 2;
-    if (array[mid] < element) {
-      start = mid + 1;
-    } else {
-      end = mid;
-    }
-  }
-  return start;
-}
-
-template <typename DataT, typename LabelT, typename IdxT, int TPB, typename ObjectiveT>
+template <typename DataT, typename LabelT, int TPB, typename ObjectiveT>
 void launchBuildHistogramsKernel(typename ObjectiveT::BinT* histograms,
-                                 IdxT n_bins,
-                                 const Dataset<DataT, LabelT, IdxT>& dataset,
-                                 const Quantiles<DataT, IdxT>& quantiles,
+                                 std::int64_t n_bins,
+                                 const Dataset<DataT, LabelT>& dataset,
+                                 const Quantiles<DataT>& quantiles,
                                  const NodeWorkItem* work_items,
-                                 IdxT colStart,
-                                 const IdxT* column_samples,
+                                 std::int64_t colStart,
+                                 const std::int64_t* column_samples,
                                  ObjectiveT& objective,
-                                 const WorkloadInfo<IdxT>* workload_info,
+                                 const WorkloadInfo* workload_info,
                                  dim3 histogram_grid,
                                  const SharedMemoryConfig& split_smem_config,
                                  cudaStream_t builder_stream);
 
-template <typename DataT, typename LabelT, typename IdxT, int TPB, typename ObjectiveT>
+template <typename DataT, typename LabelT, int TPB, typename ObjectiveT>
 void launchFindBestSplitsKernel(typename ObjectiveT::BinT* histograms,
-                                IdxT n_bins,
-                                const Dataset<DataT, LabelT, IdxT>& dataset,
-                                const Quantiles<DataT, IdxT>& quantiles,
-                                IdxT colStart,
-                                const IdxT* column_samples,
+                                std::int64_t n_bins,
+                                const Dataset<DataT, LabelT>& dataset,
+                                const Quantiles<DataT>& quantiles,
+                                std::int64_t colStart,
+                                const std::int64_t* column_samples,
                                 int* mutex,
-                                volatile Split<DataT, IdxT>* splits,
+                                volatile Split<DataT>* splits,
                                 ObjectiveT& objective,
                                 dim3 split_grid,
                                 cudaStream_t builder_stream);
@@ -168,7 +152,7 @@ inline void packHistograms(const BinT* in, double* out, std::size_t len, cudaStr
 {
   // Counts are packed as doubles so each bin can use one homogeneous arithmetic buffer. This is
   // exact for current RF problem sizes: integer values up to 2^53 are exactly representable by
-  // double, and IdxT row indexing is far below that limit.
+  // double, and RF row indexing is far below that limit.
   auto op = [in] __device__(double* out, std::size_t i) {
     auto const bin_idx = i / reduction_buffer_size_v<BinT>;
     auto const field   = i % reduction_buffer_size_v<BinT>;