decision_forest.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2023-2024, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
#include <cuml/experimental/fil/constants.hpp>
#include <cuml/experimental/fil/detail/device_initialization.hpp>
#include <cuml/experimental/fil/detail/forest.hpp>
#include <cuml/experimental/fil/detail/index_type.hpp>
#include <cuml/experimental/fil/detail/infer.hpp>
#include <cuml/experimental/fil/detail/postprocessor.hpp>
#include <cuml/experimental/fil/detail/raft_proto/buffer.hpp>
#include <cuml/experimental/fil/detail/raft_proto/cuda_stream.hpp>
#include <cuml/experimental/fil/detail/raft_proto/exceptions.hpp>
#include <cuml/experimental/fil/detail/specialization_types.hpp>
#include <cuml/experimental/fil/exceptions.hpp>
#include <cuml/experimental/fil/infer_kind.hpp>
#include <cuml/experimental/fil/postproc_ops.hpp>
#include <cuml/experimental/fil/tree_layout.hpp>
 
#include <stddef.h>
#include <stdint.h>
 
#include <algorithm>
#include <cstddef>
#include <limits>
#include <optional>
#include <variant>
 
namespace ML {
namespace experimental {
namespace fil {
 
template <tree_layout layout_v,
          typename threshold_t,
          typename index_t,
          typename metadata_storage_t,
          typename offset_t>
struct decision_forest {
  auto constexpr static const layout = layout_v;
  using forest_type = forest<layout, threshold_t, index_t, metadata_storage_t, offset_t>;
  using node_type = typename forest_type::node_type;
  using io_type = typename forest_type::io_type;
  using threshold_type = threshold_t;
  using postprocessor_type = postprocessor<io_type>;
  using categorical_storage_type = typename node_type::index_type;
 
  decision_forest()
    : nodes_{},
      root_node_indexes_{},
      node_id_mapping_{},
      vector_output_{},
      categorical_storage_{},
      num_features_{},
      num_outputs_{},
      leaf_size_{},
      has_categorical_nodes_{false},
      row_postproc_{},
      elem_postproc_{},
      average_factor_{},
      bias_{},
      postproc_constant_{}
  {
  }
 
  decision_forest(raft_proto::buffer<node_type>&& nodes,
                  raft_proto::buffer<index_type>&& root_node_indexes,
                  raft_proto::buffer<index_type>&& node_id_mapping,
                  index_type num_features,
                  index_type num_outputs                                     = index_type{2},
                  bool has_categorical_nodes                                 = false,
                  std::optional<raft_proto::buffer<io_type>>&& vector_output = std::nullopt,
                  std::optional<raft_proto::buffer<typename node_type::index_type>>&&
                    categorical_storage     = std::nullopt,
                  index_type leaf_size      = index_type{1},
                  row_op row_postproc       = row_op::disable,
                  element_op elem_postproc  = element_op::disable,
                  io_type average_factor    = io_type{1},
                  io_type bias              = io_type{0},
                  io_type postproc_constant = io_type{1})
    : nodes_{nodes},
      root_node_indexes_{root_node_indexes},
      node_id_mapping_{node_id_mapping},
      vector_output_{vector_output},
      categorical_storage_{categorical_storage},
      num_features_{num_features},
      num_outputs_{num_outputs},
      leaf_size_{leaf_size},
      has_categorical_nodes_{has_categorical_nodes},
      row_postproc_{row_postproc},
      elem_postproc_{elem_postproc},
      average_factor_{average_factor},
      bias_{bias},
      postproc_constant_{postproc_constant}
  {
    if (nodes.memory_type() != root_node_indexes.memory_type()) {
      throw raft_proto::mem_type_mismatch(
        "Nodes and indexes of forest must both be stored on either host or device");
    }
    if (nodes.device_index() != root_node_indexes.device_index()) {
      throw raft_proto::mem_type_mismatch(
        "Nodes and indexes of forest must both be stored on same device");
    }
    detail::initialize_device<forest_type>(nodes.device());
  }
 
  auto num_features() const { return num_features_; }
  auto num_trees() const { return root_node_indexes_.size(); }
  auto has_vector_leaves() const { return vector_output_.has_value(); }
 
  auto num_outputs(infer_kind inference_kind = infer_kind::default_kind) const
  {
    auto result = num_outputs_;
    if (inference_kind == infer_kind::per_tree) {
      result = num_trees();
      if (has_vector_leaves()) { result *= num_outputs_; }
    } else if (inference_kind == infer_kind::leaf_id) {
      result = num_trees();
    }
    return result;
  }
 
  auto row_postprocessing() const { return row_postproc_; }
  // Setter for row_postprocessing
  void set_row_postprocessing(row_op val) { row_postproc_ = val; }
  auto elem_postprocessing() const { return elem_postproc_; }
 
  auto memory_type() { return nodes_.memory_type(); }
  auto device_index() { return nodes_.device_index(); }
 
  void predict(raft_proto::buffer<typename forest_type::io_type>& output,
               raft_proto::buffer<typename forest_type::io_type> const& input,
               raft_proto::cuda_stream stream                          = raft_proto::cuda_stream{},
               infer_kind predict_type                                 = infer_kind::default_kind,
               std::optional<index_type> specified_rows_per_block_iter = std::nullopt)
  {
    if (output.memory_type() != memory_type() || input.memory_type() != memory_type()) {
      throw raft_proto::wrong_device_type{
        "Tried to use host I/O data with model on device or vice versa"};
    }
    if (output.device_index() != device_index() || input.device_index() != device_index()) {
      throw raft_proto::wrong_device{"I/O data on different device than model"};
    }
    auto* vector_output_data =
      (vector_output_.has_value() ? vector_output_->data() : static_cast<io_type*>(nullptr));
    auto* categorical_storage_data =
      (categorical_storage_.has_value() ? categorical_storage_->data()
                                        : static_cast<categorical_storage_type*>(nullptr));
    switch (nodes_.device().index()) {
      case 0:
        fil::detail::infer(obj(),
                           get_postprocessor(predict_type),
                           output.data(),
                           input.data(),
                           index_type(input.size() / num_features_),
                           num_features_,
                           num_outputs(predict_type),
                           has_categorical_nodes_,
                           vector_output_data,
                           categorical_storage_data,
                           predict_type,
                           specified_rows_per_block_iter,
                           std::get<0>(nodes_.device()),
                           stream);
        break;
      case 1:
        fil::detail::infer(obj(),
                           get_postprocessor(predict_type),
                           output.data(),
                           input.data(),
                           index_type(input.size() / num_features_),
                           num_features_,
                           num_outputs(predict_type),
                           has_categorical_nodes_,
                           vector_output_data,
                           categorical_storage_data,
                           predict_type,
                           specified_rows_per_block_iter,
                           std::get<1>(nodes_.device()),
                           stream);
        break;
    }
  }
 
 private:
  raft_proto::buffer<node_type> nodes_;
  raft_proto::buffer<index_type> root_node_indexes_;
  raft_proto::buffer<index_type> node_id_mapping_;
  std::optional<raft_proto::buffer<io_type>> vector_output_;
  std::optional<raft_proto::buffer<categorical_storage_type>> categorical_storage_;
 
  // Metadata
  index_type num_features_;
  index_type num_outputs_;
  index_type leaf_size_;
  bool has_categorical_nodes_ = false;
  // Postprocessing constants
  row_op row_postproc_;
  element_op elem_postproc_;
  io_type average_factor_;
  io_type bias_;
  io_type postproc_constant_;
 
  auto obj() const
  {
    return forest_type{nodes_.data(),
                       root_node_indexes_.data(),
                       node_id_mapping_.data(),
                       static_cast<index_type>(root_node_indexes_.size()),
                       num_outputs_};
  }
 
  auto get_postprocessor(infer_kind inference_kind = infer_kind::default_kind) const
  {
    auto result = postprocessor_type{};
    if (inference_kind == infer_kind::default_kind) {
      result = postprocessor_type{
        row_postproc_, elem_postproc_, average_factor_, bias_, postproc_constant_};
    }
    return result;
  }
 
  auto leaf_size() const { return leaf_size_; }
};
 
namespace detail {
template <tree_layout layout, bool double_precision, bool large_trees>
using preset_decision_forest = decision_forest<
  layout,
  typename specialization_types<layout, double_precision, large_trees>::threshold_type,
  typename specialization_types<layout, double_precision, large_trees>::index_type,
  typename specialization_types<layout, double_precision, large_trees>::metadata_type,
  typename specialization_types<layout, double_precision, large_trees>::offset_type>;
 
}  // namespace detail
 
using decision_forest_variant =
  std::variant<detail::preset_decision_forest<
                 std::variant_alternative_t<0, detail::specialization_variant>::layout,
                 std::variant_alternative_t<0, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<0, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<1, detail::specialization_variant>::layout,
                 std::variant_alternative_t<1, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<1, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<2, detail::specialization_variant>::layout,
                 std::variant_alternative_t<2, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<2, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<3, detail::specialization_variant>::layout,
                 std::variant_alternative_t<3, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<3, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<4, detail::specialization_variant>::layout,
                 std::variant_alternative_t<4, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<4, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<5, detail::specialization_variant>::layout,
                 std::variant_alternative_t<5, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<5, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<6, detail::specialization_variant>::layout,
                 std::variant_alternative_t<6, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<6, detail::specialization_variant>::has_large_trees>,
               detail::preset_decision_forest<
                 std::variant_alternative_t<7, detail::specialization_variant>::layout,
                 std::variant_alternative_t<7, detail::specialization_variant>::is_double_precision,
                 std::variant_alternative_t<7, detail::specialization_variant>::has_large_trees>>;
 
inline auto get_forest_variant_index(bool use_double_thresholds,
                                     index_type max_node_offset,
                                     index_type num_features,
                                     index_type num_categorical_nodes = index_type{},
                                     index_type max_num_categories    = index_type{},
                                     index_type num_vector_leaves     = index_type{},
                                     tree_layout layout               = preferred_tree_layout)
{
  using small_index_t =
    typename detail::specialization_types<preferred_tree_layout, false, false>::index_type;
  auto max_local_categories = index_type(sizeof(small_index_t) * 8);
  // If the index required for pointing to categorical storage bins or vector
  // leaf output exceeds what we can store in a uint32_t, uint64_t will be used
  //
  // TODO(wphicks): We are overestimating categorical storage required here
  auto double_indexes_required =
    (max_num_categories > max_local_categories &&
     ((raft_proto::ceildiv(max_num_categories, max_local_categories) + 1 * num_categorical_nodes) >
      std::numeric_limits<small_index_t>::max())) ||
    num_vector_leaves > std::numeric_limits<small_index_t>::max();
 
  auto double_precision = use_double_thresholds || double_indexes_required;
 
  using small_metadata_t =
    typename detail::specialization_types<preferred_tree_layout, false, false>::metadata_type;
  using small_offset_t =
    typename detail::specialization_types<preferred_tree_layout, false, false>::offset_type;
 
  auto large_trees =
    (num_features > (std::numeric_limits<small_metadata_t>::max() >> reserved_node_metadata_bits) ||
     max_node_offset > std::numeric_limits<small_offset_t>::max());
 
  auto layout_value = static_cast<std::underlying_type_t<tree_layout>>(layout);
 
  return ((index_type{layout_value} << index_type{2}) +
          (index_type{double_precision} << index_type{1}) + index_type{large_trees});
}
}  // namespace fil
}  // namespace experimental
}  // namespace ML