cpu.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2023, 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 <cstddef>
#include <cuml/experimental/fil/detail/cpu_introspection.hpp>
#include <cuml/experimental/fil/detail/evaluate_tree.hpp>
#include <cuml/experimental/fil/detail/index_type.hpp>
#include <cuml/experimental/fil/detail/postprocessor.hpp>
#include <cuml/experimental/fil/detail/raft_proto/ceildiv.hpp>
#include <cuml/experimental/fil/infer_kind.hpp>
#include <iostream>
#include <new>
#include <numeric>
#include <vector>
 
namespace ML {
namespace experimental {
namespace fil {
namespace detail {
 
template <bool has_categorical_nodes,
          bool predict_leaf,
          typename forest_t,
          typename vector_output_t    = std::nullptr_t,
          typename categorical_data_t = std::nullptr_t>
void infer_kernel_cpu(forest_t const& forest,
                      postprocessor<typename forest_t::io_type> const& postproc,
                      typename forest_t::io_type* output,
                      typename forest_t::io_type const* input,
                      index_type row_count,
                      index_type col_count,
                      index_type num_outputs,
                      index_type chunk_size               = hardware_constructive_interference_size,
                      index_type grove_size               = hardware_constructive_interference_size,
                      vector_output_t vector_output_p     = nullptr,
                      categorical_data_t categorical_data = nullptr,
                      infer_kind infer_type               = infer_kind::default_kind)
{
  auto constexpr has_vector_leaves       = !std::is_same_v<vector_output_t, std::nullptr_t>;
  auto constexpr has_nonlocal_categories = !std::is_same_v<categorical_data_t, std::nullptr_t>;
 
  using node_t = typename forest_t::node_type;
 
  using output_t = typename forest_t::template raw_output_type<vector_output_t>;
 
  auto const num_tree  = forest.tree_count();
  auto const num_grove = raft_proto::ceildiv(num_tree, grove_size);
  auto const num_chunk = raft_proto::ceildiv(row_count, chunk_size);
 
  auto output_workspace = std::vector<output_t>(row_count * num_outputs * num_grove, output_t{});
  auto const task_count = num_grove * num_chunk;
 
  // Infer on each grove and chunk
#pragma omp parallel for
  for (auto task_index = index_type{}; task_index < task_count; ++task_index) {
    auto const grove_index = task_index / num_chunk;
    auto const chunk_index = task_index % num_chunk;
    auto const start_row   = chunk_index * chunk_size;
    auto const end_row     = std::min(start_row + chunk_size, row_count);
    auto const start_tree  = grove_index * grove_size;
    auto const end_tree    = std::min(start_tree + grove_size, num_tree);
 
    for (auto row_index = start_row; row_index < end_row; ++row_index) {
      for (auto tree_index = start_tree; tree_index < end_tree; ++tree_index) {
        auto tree_output =
          std::conditional_t<predict_leaf,
                             index_type,
                             std::conditional_t<has_vector_leaves,
                                                typename node_t::index_type,
                                                typename node_t::threshold_type>>{};
        tree_output = evaluate_tree<has_vector_leaves,
                                    has_categorical_nodes,
                                    has_nonlocal_categories,
                                    predict_leaf>(
          forest, tree_index, input + row_index * col_count, categorical_data);
        if constexpr (predict_leaf) {
          output_workspace[row_index * num_outputs * num_grove + tree_index * num_grove +
                           grove_index] = static_cast<typename forest_t::io_type>(tree_output);
        } else {
          auto const default_num_outputs = forest.num_outputs();
          if constexpr (has_vector_leaves) {
            auto output_offset =
              (row_index * num_outputs * num_grove +
               tree_index * default_num_outputs * num_grove * (infer_type == infer_kind::per_tree) +
               grove_index);
            for (auto output_index = index_type{}; output_index < default_num_outputs;
                 ++output_index) {
              output_workspace[output_offset + output_index * num_grove] +=
                vector_output_p[tree_output * default_num_outputs + output_index];
            }
          } else {
            auto output_offset =
              (row_index * num_outputs * num_grove +
               (tree_index % default_num_outputs) * num_grove *
                 (infer_type == infer_kind::default_kind) +
               tree_index * num_grove * (infer_type == infer_kind::per_tree) + grove_index);
            output_workspace[output_offset] += tree_output;
          }
        }
      }  // Trees
    }    // Rows
  }      // Tasks
 
  // Sum over grove and postprocess
#pragma omp parallel for
  for (auto row_index = index_type{}; row_index < row_count; ++row_index) {
    for (auto output_index = index_type{}; output_index < num_outputs; ++output_index) {
      auto grove_offset = (row_index * num_outputs * num_grove + output_index * num_grove);
 
      output_workspace[grove_offset] =
        std::accumulate(std::begin(output_workspace) + grove_offset,
                        std::begin(output_workspace) + grove_offset + num_grove,
                        output_t{});
    }
    postproc(output_workspace.data() + row_index * num_outputs * num_grove,
             num_outputs,
             output + row_index * num_outputs,
             num_grove);
  }
}
 
}  // namespace detail
}  // namespace fil
}  // namespace experimental
}  // namespace ML