Unify dense and sparse tests

cpp/src/fil/common.cuh

@@ -28,12 +28,6 @@
 
 #include "internal.cuh"
 
-// needed for node_traits<...>
-namespace treelite {
-template <typename, typename>
-struct ModelImpl;
-}
-
 namespace ML {
 namespace fil {
 
@@ -120,33 +114,6 @@ 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/internal.cuh

@@ -33,6 +33,12 @@ namespace raft {
 class handle_t;
 }
 
+// needed for node_traits<...>
+namespace treelite {
+template <typename, typename>
+struct ModelImpl;
+}
+
 namespace ML {
 namespace fil {
 
@@ -215,6 +221,37 @@ struct alignas(8) sparse_node8 : base_node {
   __host__ __device__ int left(int curr) const { return left_index(); }
 };
 
+struct dense_forest;
+template <typename node_t>
+struct sparse_forest;
+
+struct dense_storage;
+template <typename node_t>
+struct sparse_storage;
+
+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)
+  {
+  }
+};
+
 /** leaf_algo_t describes what the leaves in a FIL forest store (predict)
     and how FIL aggregates them into class margins/regression result/best class
 **/

cpp/test/sg/fil_test.cu

@@ -646,31 +646,8 @@ class BaseFilTest : public testing::TestWithParam<FilTestParams> {
   FilTestParams ps;
 };
 
-class PredictDenseFilTest : public BaseFilTest {
- protected:
-  void init_forest(fil::forest_t* pforest) override
-  {
-    // init FIL model
-    fil::forest_params_t fil_ps;
-    fil_ps.depth            = ps.depth;
-    fil_ps.num_trees        = ps.num_trees;
-    fil_ps.num_cols         = ps.num_cols;
-    fil_ps.algo             = ps.algo;
-    fil_ps.output           = ps.output;
-    fil_ps.threshold        = ps.threshold;
-    fil_ps.global_bias      = ps.global_bias;
-    fil_ps.leaf_algo        = ps.leaf_algo;
-    fil_ps.num_classes      = ps.num_classes;
-    fil_ps.blocks_per_sm    = ps.blocks_per_sm;
-    fil_ps.threads_per_tree = ps.threads_per_tree;
-    fil_ps.n_items          = ps.n_items;
-
-    fil::init(handle, pforest, cat_sets_h.accessor(), vector_leaf, nullptr, nodes.data(), &fil_ps);
-  }
-};
-
 template <typename fil_node_t>
-class BasePredictSparseFilTest : public BaseFilTest {
+class BasePredictFilTest : public BaseFilTest {
  protected:
   void dense2sparse_node(const fil::dense_node* dense_root,
                          int i_dense,
@@ -717,36 +694,48 @@ class BasePredictSparseFilTest : public BaseFilTest {
 
   void init_forest(fil::forest_t* pforest) override
   {
+    constexpr bool IS_DENSE = node_traits<fil_node_t>::IS_DENSE;
+    std::vector<fil_node_t> init_nodes;
+    if constexpr (!IS_DENSE) {
+      dense2sparse();
+      init_nodes = sparse_nodes;
+    } else {
+      init_nodes = nodes;
+    }
+    ASSERT(init_nodes.size() < std::size_t(INT_MAX), "generated too many nodes");
+
     // init FIL model
-    fil::forest_params_t fil_params;
-    fil_params.num_trees        = ps.num_trees;
-    fil_params.num_cols         = ps.num_cols;
-    fil_params.algo             = ps.algo;
-    fil_params.output           = ps.output;
-    fil_params.threshold        = ps.threshold;
-    fil_params.global_bias      = ps.global_bias;
-    fil_params.leaf_algo        = ps.leaf_algo;
-    fil_params.num_classes      = ps.num_classes;
-    fil_params.blocks_per_sm    = ps.blocks_per_sm;
-    fil_params.threads_per_tree = ps.threads_per_tree;
-    fil_params.n_items          = ps.n_items;
-
-    dense2sparse();
-    fil_params.num_nodes = sparse_nodes.size();
+    fil::forest_params_t fil_params = {
+      .num_nodes        = static_cast<int>(init_nodes.size()),
+      .depth            = ps.depth,
+      .num_trees        = ps.num_trees,
+      .num_cols         = ps.num_cols,
+      .leaf_algo        = ps.leaf_algo,
+      .algo             = ps.algo,
+      .output           = ps.output,
+      .threshold        = ps.threshold,
+      .global_bias      = ps.global_bias,
+      .num_classes      = ps.num_classes,
+      .blocks_per_sm    = ps.blocks_per_sm,
+      .threads_per_tree = ps.threads_per_tree,
+      .n_items          = ps.n_items,
+    };
+
     fil::init(handle,
               pforest,
               cat_sets_h.accessor(),
               vector_leaf,
               trees.data(),
-              sparse_nodes.data(),
+              init_nodes.data(),
               &fil_params);
   }
   std::vector<fil_node_t> sparse_nodes;
   std::vector<int> trees;
 };
 
-typedef BasePredictSparseFilTest<fil::sparse_node16> PredictSparse16FilTest;
-typedef BasePredictSparseFilTest<fil::sparse_node8> PredictSparse8FilTest;
+typedef BasePredictFilTest<fil::dense_node> PredictDenseFilTest;
+typedef BasePredictFilTest<fil::sparse_node16> PredictSparse16FilTest;
+typedef BasePredictFilTest<fil::sparse_node8> PredictSparse8FilTest;
 
 class TreeliteFilTest : public BaseFilTest {
  protected: