Upgrade Treelite module

conda/environments/cuml_dev_cuda10.1.yml

@@ -21,7 +21,7 @@ dependencies:
 - faiss-proc=*=cuda
 - umap-learn
 - scikit-learn=0.23.1
-- treelite=0.93
+- treelite=1.0.0rc1
 - pip
 - pip:
     - sphinx_markdown_tables

conda/environments/cuml_dev_cuda10.2.yml

@@ -21,7 +21,7 @@ dependencies:
 - faiss-proc=*=cuda
 - umap-learn
 - scikit-learn=0.23.1
-- treelite=0.93
+- treelite=1.0.0rc1
 - pip
 - pip:
     - sphinx_markdown_tables

conda/environments/cuml_dev_cuda11.0.yml

@@ -21,7 +21,7 @@ dependencies:
 - faiss-proc=*=cuda
 - umap-learn
 - scikit-learn=0.23.1
-- treelite=0.93
+- treelite=1.0.0rc1
 - pip
 - pip:
     - sphinx_markdown_tables

conda/recipes/cuml/meta.yaml

@@ -28,7 +28,7 @@ requirements:
     - setuptools
     - cython>=0.29,<0.30
     - cmake>=3.14
-    - treelite=0.93
+    - treelite=1.0.0rc1
     - cudf {{ minor_version }}
     - libcuml={{ version }}
     - libcumlprims {{ minor_version }}
@@ -40,7 +40,7 @@ requirements:
     - dask-cudf {{ minor_version }}
     - libcuml={{ version }}
     - libcumlprims {{ minor_version }}
-    - treelite=0.93
+    - treelite=1.0.0rc1
     - cupy>7.1.0,<9.0.0a0
     - nccl>=2.5
     - ucx-py {{ minor_version }}

conda/recipes/libcuml/meta.yaml

@@ -37,7 +37,7 @@ requirements:
     - ucx-py {{ minor_version }}
     - libcumlprims {{ minor_version }}
     - lapack
-    - treelite=0.93
+    - treelite=1.0.0rc1
     - faiss-proc=*=cuda
     - gtest=1.10.0
     - libfaiss=1.6.3
@@ -47,7 +47,7 @@ requirements:
     - nccl>=2.5
     - ucx-py {{ minor_version }}
     - {{ pin_compatible('cudatoolkit', max_pin='x.x') }}
-    - treelite=0.93
+    - treelite=1.0.0rc1
     - faiss-proc=*=cuda
     - libfaiss=1.6.3
 

cpp/cmake/Dependencies.cmake

@@ -185,7 +185,7 @@ endif(BUILD_STATIC_FAISS)
 ##############################################################################
 # - treelite build -----------------------------------------------------------
 
-find_package(Treelite 0.93 REQUIRED)
+find_package(Treelite 1.0.0 REQUIRED)
 
 ##############################################################################
 # - googletest build -----------------------------------------------------------

cpp/src/decisiontree/decisiontree_impl.cuh

@@ -30,6 +30,7 @@
 #include "levelalgo/metric.cuh"
 #include "memory.cuh"
 #include "quantile/quantile.cuh"
+#include "treelite_util.h"
 
 namespace ML {
 
@@ -140,7 +141,7 @@ struct Node_ID_info {
 template <class T, class L>
 void build_treelite_tree(TreeBuilderHandle tree_builder,
                          DecisionTree::TreeMetaDataNode<T, L> *tree_ptr,
-                         int num_output_group) {
+                         int num_class) {
   int node_id = 0;
   TREELITE_CHECK(TreeliteTreeBuilderCreateNode(tree_builder, node_id));
 
@@ -174,28 +175,54 @@ void build_treelite_tree(TreeBuilderHandle tree_builder,
           TreeliteTreeBuilderCreateNode(tree_builder, node_id + 2));
 
         // Set node from current level as numerical node. Children IDs known.
+        ValueHandle threshold;
+        TREELITE_CHECK(TreeliteTreeBuilderCreateValue(
+          &q_node.node.quesval, TreeliteType<T>::value, &threshold));
         TREELITE_CHECK(TreeliteTreeBuilderSetNumericalTestNode(
           tree_builder, q_node.unique_node_id, q_node.node.colid,
-          "<=", q_node.node.quesval, 1, node_id + 1, node_id + 2));
+          "<=", threshold, 1, node_id + 1, node_id + 2));
+        TREELITE_CHECK(TreeliteTreeBuilderDeleteValue(threshold));
 
         node_id += 2;
       } else {
-        if (num_output_group == 1) {
+        if (num_class == 1) {
+          ValueHandle leaf_value;
+          if (std::is_same<L, int>::value) {
+            // Integer output is not yet supported in Treelite codegen
+            float prediction = static_cast<float>(q_node.node.prediction);
+            TREELITE_CHECK(TreeliteTreeBuilderCreateValue(
+              &prediction, TreeliteType<float>::value, &leaf_value));
+          } else {
+            TREELITE_CHECK(TreeliteTreeBuilderCreateValue(
+              &q_node.node.prediction, TreeliteType<L>::value, &leaf_value));
+          }
           TREELITE_CHECK(TreeliteTreeBuilderSetLeafNode(
-            tree_builder, q_node.unique_node_id, q_node.node.prediction));
+            tree_builder, q_node.unique_node_id, leaf_value));
+          TREELITE_CHECK(TreeliteTreeBuilderDeleteValue(leaf_value));
         } else {
-          std::vector<float> leaf_vector(num_output_group);
-          for (int j = 0; j < num_output_group; j++) {
+          std::vector<float> leaf_vector(num_class);
+          std::vector<ValueHandle> leaf_vector_handle(num_class, nullptr);
+          for (int j = 0; j < num_class; j++) {
             if (q_node.node.prediction == j) {
               leaf_vector[j] = 1;
             } else {
               leaf_vector[j] = 0;
             }
           }
+          for (int j = 0; j < num_class; j++) {
+            TREELITE_CHECK(TreeliteTreeBuilderCreateValue(
+              &leaf_vector[j], TreeliteType<float>::value,
+              &leaf_vector_handle[j]));
+          }
           TREELITE_CHECK(TreeliteTreeBuilderSetLeafVectorNode(
-            tree_builder, q_node.unique_node_id, leaf_vector.data(),
-            num_output_group));
+            tree_builder, q_node.unique_node_id, leaf_vector_handle.data(),
+            num_class));
+          for (int j = 0; j < num_class; j++) {
+            TREELITE_CHECK(
+              TreeliteTreeBuilderDeleteValue(leaf_vector_handle[j]));
+          }
           leaf_vector.clear();
+          leaf_vector_handle.clear();
         }
       }
     }
@@ -521,17 +548,16 @@ template class DecisionTreeRegressor<double>;
 
 template void build_treelite_tree<float, int>(
   TreeBuilderHandle tree_builder,
-  DecisionTree::TreeMetaDataNode<float, int> *tree_ptr, int num_output_group);
+  DecisionTree::TreeMetaDataNode<float, int> *tree_ptr, int num_class);
 template void build_treelite_tree<double, int>(
   TreeBuilderHandle tree_builder,
-  DecisionTree::TreeMetaDataNode<double, int> *tree_ptr, int num_output_group);
+  DecisionTree::TreeMetaDataNode<double, int> *tree_ptr, int num_class);
 template void build_treelite_tree<float, float>(
   TreeBuilderHandle tree_builder,
-  DecisionTree::TreeMetaDataNode<float, float> *tree_ptr, int num_output_group);
+  DecisionTree::TreeMetaDataNode<float, float> *tree_ptr, int num_class);
 template void build_treelite_tree<double, double>(
   TreeBuilderHandle tree_builder,
-  DecisionTree::TreeMetaDataNode<double, double> *tree_ptr,
-  int num_output_group);
+  DecisionTree::TreeMetaDataNode<double, double> *tree_ptr, int num_class);
 }  //End namespace DecisionTree
 
 }  //End namespace ML

cpp/src/decisiontree/decisiontree_impl.h

@@ -57,7 +57,7 @@ std::string dump_node_as_json(
 template <class T, class L>
 void build_treelite_tree(TreeBuilderHandle tree_builder,
                          DecisionTree::TreeMetaDataNode<T, L> *tree_ptr,
-                         int num_output_group);
+                         int num_class);
 
 struct DataInfo {
   unsigned int NLocalrows;

cpp/src/decisiontree/treelite_util.h

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2020, 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 <cstdint>
+
+namespace ML {
+namespace DecisionTree {
+
+template <typename T>
+class TreeliteType;
+
+template <>
+class TreeliteType<float> {
+ public:
+  static constexpr const char* value = "float32";
+};
+
+template <>
+class TreeliteType<double> {
+ public:
+  static constexpr const char* value = "float64";
+};
+
+template <>
+class TreeliteType<uint32_t> {
+ public:
+  static constexpr const char* value = "uint32";
+};
+
+template <>
+class TreeliteType<int> {
+ public:
+  static_assert(sizeof(int) == sizeof(uint32_t), "int must be 32-bit");
+  static constexpr const char* value = "uint32";
+};
+
+}  //End namespace DecisionTree
+
+}  //End namespace ML