forest.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
#include <cuml/fil/detail/index_type.hpp>
#include <cuml/fil/detail/node.hpp>
#include <cuml/fil/detail/raft_proto/gpu_support.hpp>
 
#include <stddef.h>
 
#include <type_traits>
 
namespace ML {
namespace fil {
 
/* A collection of trees which together form a forest model
 */
template <tree_layout layout_v,
          typename threshold_t,
          typename index_t,
          typename metadata_storage_t,
          typename offset_t>
struct forest {
  using node_type = node<layout_v, threshold_t, index_t, metadata_storage_t, offset_t>;
  using io_type   = threshold_t;
  template <typename vector_output_t>
  using raw_output_type = std::conditional_t<!std::is_same_v<vector_output_t, std::nullptr_t>,
                                             std::remove_pointer_t<vector_output_t>,
                                             typename node_type::threshold_type>;
 
  HOST DEVICE forest(node_type* forest_nodes,
                     index_type* forest_root_indexes,
                     index_type* node_id_mapping,
                     io_type* bias,
                     index_type num_trees,
                     index_type num_outputs)
    : nodes_{forest_nodes},
      root_node_indexes_{forest_root_indexes},
      node_id_mapping_{node_id_mapping},
      bias_{bias},
      num_trees_{num_trees},
      num_outputs_{num_outputs}
  {
  }
 
  /* Return pointer to the root node of the indicated tree */
  HOST DEVICE auto* get_tree_root(index_type tree_index) const
  {
    return nodes_ + root_node_indexes_[tree_index];
  }
 
  /* Return pointer to the mapping from internal node IDs to final node ID outputs.
   * Only used when infer_type == infer_kind::leaf_id */
  HOST DEVICE const auto* get_node_id_mapping() const { return node_id_mapping_; }
 
  /* Return pointer to the bias term */
  HOST DEVICE const auto* bias() const { return bias_; }
 
  /* Return the number of trees in this forest */
  HOST DEVICE auto tree_count() const { return num_trees_; }
 
  /* Return the number of outputs per row for default evaluation of this
   * forest */
  HOST DEVICE auto num_outputs() const { return num_outputs_; }
 
 private:
  node_type* nodes_;
  index_type* root_node_indexes_;
  index_type* node_id_mapping_;
  io_type* bias_;
  index_type num_trees_;
  index_type num_outputs_;
};
 
}  // namespace fil
}  // namespace ML