Upgrade Treelite module(#3316)

Depends on dmlc/treelite#235, dmlc/treelite#236, dmlc/treelite#237, dmlc/treelite#240, dmlc/treelite#241, dmlc/treelite#242, dmlc/treelite#243 Depends on rapidsai/integration#201 Closes #2160 Closes #2795 Closes #2817 TODOs - [x] Get code review for dmlc/treelite#235 and dmlc/treelite#236 and merge them - [x] Get code review for dmlc/treelite#237 and merge them - [x] Merge dmlc/treelite#240, dmlc/treelite#241, dmlc/treelite#242, dmlc/treelite#243 - [x] Run tests locally - [x] Tag a new version of Treelite - [x] Release a new Conda package with the new version tag - [x] Update `meta.yml` and the integration package to use the new Treelite - [x] Run tests (this time build with new Conda package) Authors: - Hyunsu Cho <chohyu01@cs.washington.edu> Approvers: - AJ Schmidt (@ajschmidt8) - William Hicks (@wphicks) - John Zedlewski (@JohnZed) URL: #3316
rapidsai · Jan 7, 2021 · b8e71ca · b8e71ca
1 parent ae7e444
commit b8e71ca
Show file tree

Hide file tree

Showing 17 changed files with 293 additions and 146 deletions.
diff --git a/conda/environments/cuml_dev_cuda10.1.yml b/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

diff --git a/conda/environments/cuml_dev_cuda10.2.yml b/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

diff --git a/conda/environments/cuml_dev_cuda11.0.yml b/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

diff --git a/conda/recipes/cuml/meta.yaml b/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 }}

diff --git a/conda/recipes/libcuml/meta.yaml b/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
 

diff --git a/cpp/cmake/Dependencies.cmake b/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 -----------------------------------------------------------

diff --git a/cpp/src/decisiontree/decisiontree_impl.cuh b/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
diff --git a/cpp/src/decisiontree/decisiontree_impl.h b/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;

diff --git a/cpp/src/decisiontree/treelite_util.h b/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