Unify dense and sparse tests

cpp/src/fil/common.cuh

@@ -28,6 +28,12 @@
 
 #include "internal.cuh"
 
+// needed for node_traits<...>
+namespace treelite {
+template <typename, typename>
+struct ModelImpl;
+}
+
 namespace ML {
 namespace fil {
 
@@ -114,6 +120,33 @@ struct sparse_storage : storage_base {
 typedef sparse_storage<sparse_node16> sparse_storage16;
 typedef sparse_storage<sparse_node8> sparse_storage8;
 
+struct dense_forest;
+template <typename node_t>
+struct sparse_forest;
+
+template <typename node_t>
+struct node_traits {
+  using storage              = sparse_storage<node_t>;
+  using forest               = sparse_forest<node_t>;
+  static const bool IS_DENSE = false;
+  static const storage_type_t storage_type_enum =
+    std::is_same<sparse_node16, node_t>() ? SPARSE : SPARSE8;
+  template <typename threshold_t, typename leaf_t>
+  static void check(const treelite::ModelImpl<threshold_t, leaf_t>& model);
+};
+
+template <>
+struct node_traits<dense_node> {
+  using storage                                 = dense_storage;
+  using forest                                  = dense_forest;
+  static const bool IS_DENSE                    = true;
+  static const storage_type_t storage_type_enum = DENSE;
+  template <typename threshold_t, typename leaf_t>
+  static void check(const treelite::ModelImpl<threshold_t, leaf_t>& model)
+  {
+  }
+};
+
 /// all model parameters mostly required to compute shared memory footprint,
 /// also the footprint itself
 struct shmem_size_params {

cpp/src/fil/fil.cu

@@ -379,9 +379,12 @@ struct dense_forest : forest {
     }
   }
 
+  /// const int* trees is ignored and only provided for compatibility with
+  /// sparse_forest<node_t>::init()
   void init(const raft::handle_t& h,
             const categorical_sets& cat_sets,
             const std::vector<float>& vector_leaf,
+            const int* trees,
             const dense_node* nodes,
             const forest_params_t* params)
   {
@@ -560,50 +563,49 @@ void check_params(const forest_params_t* params, bool dense)
          FIL_TPB);
 }
 
-void init_dense(const raft::handle_t& h,
-                forest_t* pf,
-                const categorical_sets& cat_sets,
-                const std::vector<float>& vector_leaf,
-                const dense_node* nodes,
-                const forest_params_t* params)
-{
-  check_params(params, true);
-  dense_forest* f = new dense_forest(h);
-  f->init(h, cat_sets, vector_leaf, nodes, params);
-  *pf = f;
-}
-
+/** initializes a forest of any type
+ * When fil_node_t == dense_node, const int* trees is ignored
+ */
 template <typename fil_node_t>
-void init_sparse(const raft::handle_t& h,
-                 forest_t* pf,
-                 const categorical_sets& cat_sets,
-                 const std::vector<float>& vector_leaf,
-                 const int* trees,
-                 const fil_node_t* nodes,
-                 const forest_params_t* params)
+void init(const raft::handle_t& h,
+          forest_t* pf,
+          const categorical_sets& cat_sets,
+          const std::vector<float>& vector_leaf,
+          const int* trees,
+          const fil_node_t* nodes,
+          const forest_params_t* params)
 {
-  check_params(params, false);
-  sparse_forest<fil_node_t>* f = new sparse_forest<fil_node_t>(h);
+  check_params(params, node_traits<fil_node_t>::IS_DENSE);
+  using forest_type = typename node_traits<fil_node_t>::forest;
+  forest_type* f    = new forest_type(h);
   f->init(h, cat_sets, vector_leaf, trees, nodes, params);
   *pf = f;
 }
 
 // explicit instantiations for init_sparse()
-template void init_sparse<sparse_node16>(const raft::handle_t& h,
-                                         forest_t* pf,
-                                         const categorical_sets& cat_sets,
-                                         const std::vector<float>& vector_leaf,
-                                         const int* trees,
-                                         const sparse_node16* nodes,
-                                         const forest_params_t* params);
-
-template void init_sparse<sparse_node8>(const raft::handle_t& h,
-                                        forest_t* pf,
-                                        const categorical_sets& cat_sets,
-                                        const std::vector<float>& vector_leaf,
-                                        const int* trees,
-                                        const sparse_node8* nodes,
-                                        const forest_params_t* params);
+template void init<sparse_node16>(const raft::handle_t& h,
+                                  forest_t* pf,
+                                  const categorical_sets& cat_sets,
+                                  const std::vector<float>& vector_leaf,
+                                  const int* trees,
+                                  const sparse_node16* nodes,
+                                  const forest_params_t* params);
+
+template void init<sparse_node8>(const raft::handle_t& h,
+                                 forest_t* pf,
+                                 const categorical_sets& cat_sets,
+                                 const std::vector<float>& vector_leaf,
+                                 const int* trees,
+                                 const sparse_node8* nodes,
+                                 const forest_params_t* params);
+
+template void init<dense_node>(const raft::handle_t& h,
+                               forest_t* pf,
+                               const categorical_sets& cat_sets,
+                               const std::vector<float>& vector_leaf,
+                               const int* trees,
+                               const dense_node* nodes,
+                               const forest_params_t* params);
 
 void free(const raft::handle_t& h, forest_t f)
 {

cpp/src/fil/internal.cuh

@@ -149,7 +149,14 @@ __host__ __device__ __forceinline__ val_t base_node::output<val_t>() const
 /** dense_node is a single node of a dense forest */
 struct alignas(8) dense_node : base_node {
   dense_node() = default;
-  dense_node(val_t output, val_t split, int fid, bool def_left, bool is_leaf, bool is_categorical)
+  /// ignoring left_index, this is useful to unify import from treelite
+  dense_node(val_t output,
+             val_t split,
+             int fid,
+             bool def_left,
+             bool is_leaf,
+             bool is_categorical,
+             int left_index = -1)
     : base_node(output, split, fid, def_left, is_leaf, is_categorical)
   {
   }
@@ -492,40 +499,27 @@ struct cat_sets_device_owner {
   }
 };
 
-/** init_dense uses params and nodes to initialize the dense forest stored in pf
+/** init uses params, trees and nodes to initialize the forest
+ *  with nodes stored in pf
+ *  @tparam fil_node_t node type to use with the forest;
+ *    must be sparse_node16, sparse_node8 or dense_node
  *  @param h cuML handle used by this function
  *  @param pf pointer to where to store the newly created forest
- *  @param nodes nodes for the forest, of length
-      (2**(params->depth + 1) - 1) * params->ntrees
- *  @param params pointer to parameters used to initialize the forest
- *  @param vector_leaf optional vector leaves
- */
-void init_dense(const raft::handle_t& h,
-                forest_t* pf,
-                const categorical_sets& cat_sets,
-                const std::vector<float>& vector_leaf,
-                const dense_node* nodes,
-                const forest_params_t* params);
-
-/** init_sparse uses params, trees and nodes to initialize the sparse forest
- *  with sparse nodes stored in pf
- *  @tparam fil_node_t node type to use with the sparse forest;
- *    must be sparse_node16 or sparse_node8
- *  @param h cuML handle used by this function
- *  @param pf pointer to where to store the newly created forest
- *  @param trees indices of tree roots in the nodes arrray, of length params->ntrees
- *  @param nodes nodes for the forest, of length params->num_nodes
+ *  @param trees for sparse forests, indices of tree roots in the nodes arrray, of length
+ params->ntrees; ignored for dense forests
+ *  @param nodes nodes for the forest, of length params->num_nodes for sparse
+      or (2**(params->depth + 1) - 1) * params->ntrees for dense forests
  *  @param params pointer to parameters used to initialize the forest
  *  @param vector_leaf optional vector leaves
  */
 template <typename fil_node_t>
-void init_sparse(const raft::handle_t& h,
-                 forest_t* pf,
-                 const categorical_sets& cat_sets,
-                 const std::vector<float>& vector_leaf,
-                 const int* trees,
-                 const fil_node_t* nodes,
-                 const forest_params_t* params);
+void init(const raft::handle_t& h,
+          forest_t* pf,
+          const categorical_sets& cat_sets,
+          const std::vector<float>& vector_leaf,
+          const int* trees,
+          const fil_node_t* nodes,
+          const forest_params_t* params);
 
 struct predict_params;