fil.h

Go to the documentation of this file.

/*
 * Copyright (c) 2019-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/ensemble/treelite_defs.hpp>
 
#include <stddef.h>
 
#include <variant>  // for std::get<>, std::variant<>
 
namespace raft {
class handle_t;
}
 
namespace ML {
namespace fil {
 
enum algo_t {
  ALGO_AUTO,
  NAIVE,
  TREE_REORG,
  BATCH_TREE_REORG
};
 
enum storage_type_t {
  AUTO,
  DENSE,
  SPARSE,
  SPARSE8,
};
static const char* storage_type_repr[] = {"AUTO", "DENSE", "SPARSE", "SPARSE8"};
 
enum precision_t {
  PRECISION_NATIVE,
  PRECISION_FLOAT32,
  PRECISION_FLOAT64
};
 
template <typename real_t>
struct forest;
 
template <typename real_t>
using forest_t = forest<real_t>*;
 
using forest32_t = forest<float>*;
using forest64_t = forest<double>*;
 
using forest_variant = std::variant<forest_t<float>, forest_t<double>>;
 
constexpr int MAX_N_ITEMS = 4;
 
struct treelite_params_t {
  // algo is the inference algorithm
  algo_t algo;
  // output_class indicates whether thresholding will be applied
  // to the model output
  bool output_class;
  // threshold may be used for thresholding if output_class == true,
  // and is ignored otherwise. threshold is ignored if leaves store
  // vectorized class labels. in that case, a class with most votes
  // is returned regardless of the absolute vote count
  float threshold;
  // storage_type indicates whether the forest should be imported as dense or sparse
  storage_type_t storage_type;
  // blocks_per_sm, if nonzero, works as a limit to improve cache hit rate for larger forests
  // suggested values (if nonzero) are from 2 to 7
  // if zero, launches ceildiv(num_rows, NITEMS) blocks
  int blocks_per_sm;
  // threads_per_tree determines how many threads work on a single tree at once inside a block
  // can only be a power of 2
  int threads_per_tree;
  // n_items is how many input samples (items) any thread processes. If 0 is given,
  // choose most (up to MAX_N_ITEMS) that fit into shared memory.
  int n_items;
  // if non-nullptr, *pforest_shape_str will be set to caller-owned string that
  // contains forest shape
  char** pforest_shape_str;
  // precision in which to load the treelite model
  precision_t precision;
};
 
void from_treelite(const raft::handle_t& handle,
                   forest_variant* pforest,
                   TreeliteModelHandle model,
                   const treelite_params_t* tl_params);
 
template <typename real_t>
void free(const raft::handle_t& h, forest_t<real_t> f);
 
template <typename real_t>
void predict(const raft::handle_t& h,
             forest_t<real_t> f,
             real_t* preds,
             const real_t* data,
             size_t num_rows,
             bool predict_proba = false);
 
}  // namespace fil
}  // namespace ML