forest_model.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/decision_forest.hpp>
#include <cuml/experimental/fil/detail/index_type.hpp>
#include <cuml/experimental/fil/detail/raft_proto/buffer.hpp>
#include <cuml/experimental/fil/detail/raft_proto/gpu_support.hpp>
#include <cuml/experimental/fil/detail/raft_proto/handle.hpp>
#include <cuml/experimental/fil/infer_kind.hpp>
 
#include <cstddef>
#include <type_traits>
#include <variant>
 
namespace ML {
namespace experimental {
namespace fil {
 
struct forest_model {
  forest_model(decision_forest_variant&& forest = decision_forest_variant{})
    : decision_forest_{forest}
  {
  }
 
  auto num_features()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.num_features(); },
                      decision_forest_);
  }
 
  auto num_outputs()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.num_outputs(); },
                      decision_forest_);
  }
 
  auto num_trees()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.num_trees(); },
                      decision_forest_);
  }
 
  auto has_vector_leaves()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.has_vector_leaves(); },
                      decision_forest_);
  }
 
  auto row_postprocessing()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.row_postprocessing(); },
                      decision_forest_);
  }
 
  void set_row_postprocessing(row_op val)
  {
    return std::visit(
      [&val](auto&& concrete_forest) { concrete_forest.set_row_postprocessing(val); },
      decision_forest_);
  }
 
  auto elem_postprocessing()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.elem_postprocessing(); },
                      decision_forest_);
  }
 
  auto memory_type()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.memory_type(); },
                      decision_forest_);
  }
 
  auto device_index()
  {
    return std::visit([](auto&& concrete_forest) { return concrete_forest.device_index(); },
                      decision_forest_);
  }
 
  auto is_double_precision()
  {
    return std::visit(
      [](auto&& concrete_forest) {
        return std::is_same_v<typename std::remove_reference_t<decltype(concrete_forest)>::io_type,
                              double>;
      },
      decision_forest_);
  }
 
  template <typename io_t>
  void predict(raft_proto::buffer<io_t>& output,
               raft_proto::buffer<io_t> const& input,
               raft_proto::cuda_stream stream                 = raft_proto::cuda_stream{},
               infer_kind predict_type                        = infer_kind::default_kind,
               std::optional<index_type> specified_chunk_size = std::nullopt)
  {
    std::visit(
      [this, predict_type, &output, &input, &stream, &specified_chunk_size](
        auto&& concrete_forest) {
        if constexpr (std::is_same_v<
                        typename std::remove_reference_t<decltype(concrete_forest)>::io_type,
                        io_t>) {
          concrete_forest.predict(output, input, stream, predict_type, specified_chunk_size);
        } else {
          throw type_error("Input type does not match model_type");
        }
      },
      decision_forest_);
  }
 
  template <typename io_t>
  void predict(raft_proto::handle_t const& handle,
               raft_proto::buffer<io_t>& output,
               raft_proto::buffer<io_t> const& input,
               infer_kind predict_type                        = infer_kind::default_kind,
               std::optional<index_type> specified_chunk_size = std::nullopt)
  {
    std::visit(
      [this, predict_type, &handle, &output, &input, &specified_chunk_size](
        auto&& concrete_forest) {
        using model_io_t = typename std::remove_reference_t<decltype(concrete_forest)>::io_type;
        if constexpr (std::is_same_v<model_io_t, io_t>) {
          if (output.memory_type() == memory_type() && input.memory_type() == memory_type()) {
            concrete_forest.predict(
              output, input, handle.get_next_usable_stream(), predict_type, specified_chunk_size);
          } else {
            auto constexpr static const MIN_CHUNKS_PER_PARTITION = std::size_t{64};
            auto constexpr static const MAX_CHUNK_SIZE           = std::size_t{64};
 
            auto row_count = input.size() / num_features();
            auto partition_size =
              std::max(raft_proto::ceildiv(row_count, handle.get_usable_stream_count()),
                       specified_chunk_size.value_or(MAX_CHUNK_SIZE) * MIN_CHUNKS_PER_PARTITION);
            auto partition_count = raft_proto::ceildiv(row_count, partition_size);
            for (auto i = std::size_t{}; i < partition_count; ++i) {
              auto stream = handle.get_next_usable_stream();
              auto rows_in_this_partition =
                std::min(partition_size, row_count - i * partition_size);
              auto partition_in = raft_proto::buffer<io_t>{};
              if (input.memory_type() != memory_type()) {
                partition_in =
                  raft_proto::buffer<io_t>{rows_in_this_partition * num_features(), memory_type()};
                raft_proto::copy<raft_proto::DEBUG_ENABLED>(partition_in,
                                                            input,
                                                            0,
                                                            i * partition_size * num_features(),
                                                            partition_in.size(),
                                                            stream);
              } else {
                partition_in =
                  raft_proto::buffer<io_t>{input.data() + i * partition_size * num_features(),
                                           rows_in_this_partition * num_features(),
                                           memory_type()};
              }
              auto partition_out = raft_proto::buffer<io_t>{};
              if (output.memory_type() != memory_type()) {
                partition_out =
                  raft_proto::buffer<io_t>{rows_in_this_partition * num_outputs(), memory_type()};
              } else {
                partition_out =
                  raft_proto::buffer<io_t>{output.data() + i * partition_size * num_outputs(),
                                           rows_in_this_partition * num_outputs(),
                                           memory_type()};
              }
              concrete_forest.predict(
                partition_out, partition_in, stream, predict_type, specified_chunk_size);
              if (output.memory_type() != memory_type()) {
                raft_proto::copy<raft_proto::DEBUG_ENABLED>(output,
                                                            partition_out,
                                                            i * partition_size * num_outputs(),
                                                            0,
                                                            partition_out.size(),
                                                            stream);
              }
            }
          }
        } else {
          throw type_error("Input type does not match model_type");
        }
      },
      decision_forest_);
  }
 
  template <typename io_t>
  void predict(raft_proto::handle_t const& handle,
               io_t* output,
               io_t* input,
               std::size_t num_rows,
               raft_proto::device_type out_mem_type,
               raft_proto::device_type in_mem_type,
               infer_kind predict_type                        = infer_kind::default_kind,
               std::optional<index_type> specified_chunk_size = std::nullopt)
  {
    // TODO(wphicks): Make sure buffer lands on same device as model
    auto out_buffer = raft_proto::buffer{output, num_rows * num_outputs(), out_mem_type};
    auto in_buffer  = raft_proto::buffer{input, num_rows * num_features(), in_mem_type};
    predict(handle, out_buffer, in_buffer, predict_type, specified_chunk_size);
  }
 
 private:
  decision_forest_variant decision_forest_;
};
 
}  // namespace fil
}  // namespace experimental
}  // namespace ML