From 621fdb44e5c91d7a9e2d440df30e57167ca159a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= <xadupre@users.noreply.github.com>
Date: Fri, 9 Oct 2020 11:26:07 +0200
Subject: [PATCH] Fixes #4688, remove CPUAllocator in TreeEnsemble (#5375)

---
 .../cpu/ml/tree_ensemble_classifier.cc        |  2 +-
 .../providers/cpu/ml/tree_ensemble_common.h   | 25 ++++++++++---------
 .../core/providers/cpu/ml/treeregressor.cc    |  2 +-
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc b/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc
index a18c759d5f..f4512e5df0 100644
--- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc
+++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc
@@ -173,7 +173,7 @@ common::Status TreeEnsembleClassifier<T>::Compute(OpKernelContext* context) cons
   Tensor* Y = context->Output(0, {N});
   Tensor* Z = context->Output(1, {N, tree_ensemble_.get_class_count()});
 
-  tree_ensemble_.compute(context->GetOperatorThreadPool(), &X, Z, Y);
+  tree_ensemble_.compute(context, &X, Z, Y);
   return Status::OK();
 }
 
diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h
index 0e8170d576..0dbdf544bc 100644
--- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h
+++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h
@@ -52,7 +52,7 @@ class TreeEnsembleCommon {
                      const std::vector<int64_t>& target_class_treeids,
                      const std::vector<OTYPE>& target_class_weights);
 
-  void compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const;
+  void compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z, Tensor* label) const;
 
  protected:
   TreeNodeElement<OTYPE>* ProcessTreeNodeLeave(
@@ -216,33 +216,33 @@ TreeEnsembleCommon<ITYPE, OTYPE>::TreeEnsembleCommon(int parallel_tree, int para
 }
 
 template <typename ITYPE, typename OTYPE>
-void TreeEnsembleCommon<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z,
+void TreeEnsembleCommon<ITYPE, OTYPE>::compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z,
                                                Tensor* label) const {
   switch (aggregate_function_) {
     case AGGREGATE_FUNCTION::AVERAGE:
       ComputeAgg(
-          ttp, X, Z, label,
+          ctx->GetOperatorThreadPool(), X, Z, label,
           TreeAggregatorAverage<ITYPE, OTYPE>(
               roots_.size(), n_targets_or_classes_,
               post_transform_, base_values_));
       return;
     case AGGREGATE_FUNCTION::SUM:
       ComputeAgg(
-          ttp, X, Z, label,
+          ctx->GetOperatorThreadPool(), X, Z, label,
           TreeAggregatorSum<ITYPE, OTYPE>(
               roots_.size(), n_targets_or_classes_,
               post_transform_, base_values_));
       return;
     case AGGREGATE_FUNCTION::MIN:
       ComputeAgg(
-          ttp, X, Z, label,
+          ctx->GetOperatorThreadPool(), X, Z, label,
           TreeAggregatorMin<ITYPE, OTYPE>(
               roots_.size(), n_targets_or_classes_,
               post_transform_, base_values_));
       return;
     case AGGREGATE_FUNCTION::MAX:
       ComputeAgg(
-          ttp, X, Z, label,
+          ctx->GetOperatorThreadPool(), X, Z, label,
           TreeAggregatorMax<ITYPE, OTYPE>(
               roots_.size(), n_targets_or_classes_,
               post_transform_, base_values_));
@@ -525,7 +525,7 @@ class TreeEnsembleCommonClassifier : TreeEnsembleCommon<ITYPE, OTYPE> {
 
   int64_t get_class_count() const { return this->n_targets_or_classes_; }
 
-  void compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const;
+  void compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z, Tensor* label) const;
 };
 
 template <typename ITYPE, typename OTYPE>
@@ -589,11 +589,11 @@ TreeEnsembleCommonClassifier<ITYPE, OTYPE>::TreeEnsembleCommonClassifier(
 }
 
 template <typename ITYPE, typename OTYPE>
-void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z,
+void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z,
                                                          Tensor* label) const {
   if (classlabels_strings_.size() == 0) {
     this->ComputeAgg(
-        ttp, X, Z, label,
+        ctx->GetOperatorThreadPool(), X, Z, label,
         TreeAggregatorClassifier<ITYPE, OTYPE>(
             this->roots_.size(), this->n_targets_or_classes_,
             this->post_transform_, this->base_values_,
@@ -601,10 +601,11 @@ void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(concurrency::ThreadPool
             weights_are_all_positive_));
   } else {
     int64_t N = X->Shape().NumDimensions() == 1 ? 1 : X->Shape()[0];
-    std::shared_ptr<IAllocator> allocator = std::make_shared<CPUAllocator>();
-    Tensor label_int64(DataTypeImpl::GetType<int64_t>(), TensorShape({N}), allocator);
+    AllocatorPtr alloc;
+    ORT_THROW_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc));
+    Tensor label_int64(DataTypeImpl::GetType<int64_t>(), TensorShape({N}), alloc);
     this->ComputeAgg(
-        ttp, X, Z, &label_int64,
+        ctx->GetOperatorThreadPool(), X, Z, &label_int64,
         TreeAggregatorClassifier<ITYPE, OTYPE>(
             this->roots_.size(), this->n_targets_or_classes_,
             this->post_transform_, this->base_values_,
diff --git a/onnxruntime/core/providers/cpu/ml/treeregressor.cc b/onnxruntime/core/providers/cpu/ml/treeregressor.cc
index 8e808dd4ab..cfee2ccae8 100644
--- a/onnxruntime/core/providers/cpu/ml/treeregressor.cc
+++ b/onnxruntime/core/providers/cpu/ml/treeregressor.cc
@@ -56,7 +56,7 @@ common::Status TreeEnsembleRegressor<T>::Compute(OpKernelContext* context) const
   int64_t N = X->Shape().NumDimensions() == 1 ? 1 : X->Shape()[0];
   Tensor* Y = context->Output(0, {N, tree_ensemble_.n_targets_or_classes_});
 
-  tree_ensemble_.compute(context->GetOperatorThreadPool(), X, Y, NULL);
+  tree_ensemble_.compute(context, X, Y, NULL);
 
   return Status::OK();
 }