diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h index 26fc6f6006..3d016f62c6 100644 --- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h +++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h @@ -29,8 +29,9 @@ class TreeEnsembleCommonAttributes { int64_t n_trees_; bool same_mode_; bool has_missing_tracks_; - int parallel_tree_; // starts parallelizing the computing if n_tree >= parallel_tree_ and n_rows == 1 - int parallel_N_; // starts parallelizing the computing if n_rows >= parallel_N_ + int parallel_tree_; // starts parallelizing the computing by trees if n_tree >= parallel_tree_ + int parallel_tree_N_; // batch size if parallelizing by trees + int parallel_N_; // starts parallelizing the computing by rows if n_rows <= parallel_N_ }; // TI: input type @@ -50,6 +51,7 @@ class TreeEnsembleCommon : public TreeEnsembleCommonAttributes { virtual Status compute(OpKernelContext* ctx, const Tensor* X, Tensor* Y, Tensor* label) const; Status Init(int parallel_tree, + int parallel_tree_N, int parallel_N, const std::string& aggregate_function, const std::vector& base_values, @@ -84,7 +86,7 @@ class TreeEnsembleCommon : public TreeEnsembleCommonAttributes { template Status TreeEnsembleCommon::Init(const OpKernelInfo& info) { std::vector base_values_as_tensor, nodes_hitrates_as_tensor, - nodes_values_as_tensor, target_weights_as_tensor; + nodes_values_as_tensor, target_weights_as_tensor; #if !defined(ORT_MINIMAL_BUILD) ORT_THROW_IF_ERROR(GetVectorAttrsOrDefault(info, "base_values_as_tensor", base_values_as_tensor)); ORT_THROW_IF_ERROR(GetVectorAttrsOrDefault(info, "nodes_hitrates_as_tensor", nodes_hitrates_as_tensor)); @@ -94,6 +96,7 @@ Status TreeEnsembleCommon::Init(const OpKe return Init( 80, + 128, 50, info.GetAttrOrDefault("aggregate_function", "SUM"), info.GetAttrsOrDefault("base_values"), @@ -119,29 +122,33 @@ Status TreeEnsembleCommon::Init(const OpKe } template -Status TreeEnsembleCommon::Init(int parallel_tree, int parallel_N, - const std::string& aggregate_function, - const std::vector& base_values, - const std::vector& base_values_as_tensor, - int64_t n_targets_or_classes, - const std::vector& nodes_falsenodeids, - const std::vector& nodes_featureids, - const std::vector& nodes_hitrates, - const std::vector& nodes_hitrates_as_tensor, - const std::vector& nodes_missing_value_tracks_true, - const std::vector& nodes_modes, - const std::vector& nodes_nodeids, - const std::vector& nodes_treeids, - const std::vector& nodes_truenodeids, - const std::vector& nodes_values, - const std::vector& nodes_values_as_tensor, - const std::string& post_transform, - const std::vector& target_class_ids, - const std::vector& target_class_nodeids, - const std::vector& target_class_treeids, - const std::vector& target_class_weights, - const std::vector& target_class_weights_as_tensor) { +Status TreeEnsembleCommon::Init( + int parallel_tree, + int parallel_tree_N, + int parallel_N, + const std::string& aggregate_function, + const std::vector& base_values, + const std::vector& base_values_as_tensor, + int64_t n_targets_or_classes, + const std::vector& nodes_falsenodeids, + const std::vector& nodes_featureids, + const std::vector& nodes_hitrates, + const std::vector& nodes_hitrates_as_tensor, + const std::vector& nodes_missing_value_tracks_true, + const std::vector& nodes_modes, + const std::vector& nodes_nodeids, + const std::vector& nodes_treeids, + const std::vector& nodes_truenodeids, + const std::vector& nodes_values, + const std::vector& nodes_values_as_tensor, + const std::string& post_transform, + const std::vector& target_class_ids, + const std::vector& target_class_nodeids, + const std::vector& target_class_treeids, + const std::vector& target_class_weights, + const std::vector& target_class_weights_as_tensor) { parallel_tree_ = parallel_tree; + parallel_tree_N_ = parallel_tree_N; parallel_N_ = parallel_N; ORT_ENFORCE(n_targets_or_classes > 0); @@ -367,7 +374,7 @@ void TreeEnsembleCommon::ComputeAgg(concur if (n_targets_or_classes_ == 1) { if (N == 1) { ScoreValue score = {0, 0}; - if (n_trees_ <= parallel_tree_) { /* section A: 1 output, 1 row and not enough trees to parallelize */ + if (n_trees_ <= parallel_tree_ || max_num_threads == 1) { /* section A: 1 output, 1 row and not enough trees to parallelize */ for (int64_t j = 0; j < n_trees_; ++j) { agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[onnxruntime::narrow(j)], x_data)); } @@ -379,45 +386,66 @@ void TreeEnsembleCommon::ComputeAgg(concur [this, &scores, &agg, x_data](ptrdiff_t j) { agg.ProcessTreeNodePrediction1(scores[j], *ProcessTreeNodeLeave(roots_[j], x_data)); }, - 0); + max_num_threads); for (auto it = scores.cbegin(); it != scores.cend(); ++it) { agg.MergePrediction1(score, *it); } } agg.FinalizeScores1(z_data, score, label_data); - } else if (N <= parallel_N_) { /* section C: 1 output, 2+ rows but not enough rows to parallelize */ - ScoreValue score; + } else if (N <= parallel_N_ || max_num_threads == 1) { /* section C: 1 output, 2+ rows but not enough rows to parallelize */ + // Not enough data to parallelize but the computation is split into batches of 128 rows, + // and then loop on trees to evaluate every tree on this batch. + // This change was introduced by PR: https://github.com/microsoft/onnxruntime/pull/13835. + // The input tensor (2D) is stored in a contiguous array. Therefore, it is faster + // to loop on tree first and inside that loop evaluate a tree on the input tensor (inner loop). + // The processor is faster when it has to move chunks of a contiguous array (branching). + // However, if the input tensor is too big, the data does not hold on caches (L1, L2, L3). + // In that case, looping first on tree or on data is almost the same. That's why the first loop + // split into batch so that every batch holds on caches, then loop on trees and finally loop + // on the batch rows. + std::vector> scores(parallel_tree_N_); size_t j; + int64_t i, batch, batch_end; - for (int64_t i = 0; i < N; ++i) { - score = {0, 0}; - for (j = 0; j < static_cast(n_trees_); ++j) { - agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + for (batch = 0; batch < N; batch += parallel_tree_N_) { + batch_end = std::min(N, batch + parallel_tree_N_); + for (i = batch; i < batch_end; ++i) { + scores[SafeInt(i - batch)] = {0, 0}; + } + for (j = 0; j < static_cast(n_trees_); ++j) { + for (i = batch; i < batch_end; ++i) { + agg.ProcessTreeNodePrediction1(scores[SafeInt(i - batch)], *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + } + for (i = batch; i < batch_end; ++i) { + agg.FinalizeScores1(z_data + i, scores[SafeInt(i - batch)], + label_data == nullptr ? nullptr : (label_data + i)); } - - agg.FinalizeScores1(z_data + i, score, - label_data == nullptr ? nullptr : (label_data + i)); } } else if (n_trees_ > max_num_threads) { /* section D: 1 output, 2+ rows and enough trees to parallelize */ auto num_threads = std::min(max_num_threads, SafeInt(n_trees_)); std::vector> scores(SafeInt(num_threads) * N); - concurrency::ThreadPool::TrySimpleParallelFor( - ttp, - num_threads, - [this, &agg, &scores, num_threads, x_data, N, stride](ptrdiff_t batch_num) { - auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, onnxruntime::narrow(this->n_trees_)); - for (int64_t i = 0; i < N; ++i) { - scores[batch_num * SafeInt(N) + i] = {0, 0}; - } - for (auto j = work.start; j < work.end; ++j) { - for (int64_t i = 0; i < N; ++i) { - agg.ProcessTreeNodePrediction1(scores[batch_num * SafeInt(N) + i], - *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + int64_t end_n, begin_n = 0; + while (begin_n < N) { + end_n = std::min(N, begin_n + parallel_tree_N_); + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, &agg, &scores, num_threads, x_data, N, begin_n, end_n, stride](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, onnxruntime::narrow(this->n_trees_)); + for (int64_t i = begin_n; i < end_n; ++i) { + scores[batch_num * SafeInt(N) + i] = {0, 0}; } - } - }); - + for (auto j = work.start; j < work.end; ++j) { + for (int64_t i = begin_n; i < end_n; ++i) { + agg.ProcessTreeNodePrediction1(scores[batch_num * SafeInt(N) + i], + *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + } + }); + begin_n = end_n; + } concurrency::ThreadPool::TrySimpleParallelFor( ttp, num_threads, @@ -444,11 +472,11 @@ void TreeEnsembleCommon::ComputeAgg(concur agg.FinalizeScores1(z_data + i, score, label_data == nullptr ? nullptr : (label_data + i)); }, - 0); + max_num_threads); } } else { - if (N == 1) { /* section A2: 2+ outputs, 1 row, not enough trees to parallelize */ - if (n_trees_ <= parallel_tree_) { /* section A2 */ + if (N == 1) { /* section A2: 2+ outputs, 1 row, not enough trees to parallelize */ + if (n_trees_ <= parallel_tree_ || max_num_threads == 1) { /* section A2 */ InlinedVector> scores(onnxruntime::narrow(n_targets_or_classes_), {0, 0}); for (int64_t j = 0; j < n_trees_; ++j) { agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[onnxruntime::narrow(j)], x_data)); @@ -472,38 +500,53 @@ void TreeEnsembleCommon::ComputeAgg(concur } agg.FinalizeScores(scores[0], z_data, -1, label_data); } - } else if (N <= parallel_N_) { /* section C2: 2+ outputs, 2+ rows, not enough rows to parallelize */ - InlinedVector> scores(onnxruntime::narrow(n_targets_or_classes_)); + } else if (N <= parallel_N_ || max_num_threads == 1) { /* section C2: 2+ outputs, 2+ rows, not enough rows to parallelize */ + std::vector>> scores(parallel_tree_N_); size_t j, limit; - - for (int64_t i = 0; i < N; ++i) { - std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); - for (j = 0, limit = roots_.size(); j < limit; ++j) { - agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); - } - - agg.FinalizeScores(scores, z_data + i * n_targets_or_classes_, -1, - label_data == nullptr ? nullptr : (label_data + i)); + int64_t i, batch, batch_end; + batch_end = std::min(N, static_cast(parallel_tree_N_)); + for (i = 0; i < batch_end; ++i) { + scores[SafeInt(i)].resize(onnxruntime::narrow(n_targets_or_classes_)); } + for (batch = 0; batch < N; batch += parallel_tree_N_) { + batch_end = std::min(N, batch + parallel_tree_N_); + for (i = batch; i < batch_end; ++i) { + std::fill(scores[SafeInt(i - batch)].begin(), scores[SafeInt(i - batch)].end(), ScoreValue({0, 0})); + } + for (j = 0, limit = roots_.size(); j < limit; ++j) { + for (i = batch; i < batch_end; ++i) { + agg.ProcessTreeNodePrediction(scores[SafeInt(i - batch)], *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + } + for (i = batch; i < batch_end; ++i) { + agg.FinalizeScores(scores[SafeInt(i - batch)], z_data + i * n_targets_or_classes_, -1, + label_data == nullptr ? nullptr : (label_data + i)); + } + } + } else if (n_trees_ >= max_num_threads) { /* section: D2: 2+ outputs, 2+ rows, enough trees to parallelize*/ auto num_threads = std::min(max_num_threads, SafeInt(n_trees_)); std::vector>> scores(SafeInt(num_threads) * N); - concurrency::ThreadPool::TrySimpleParallelFor( - ttp, - num_threads, - [this, &agg, &scores, num_threads, x_data, N, stride](ptrdiff_t batch_num) { - auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, onnxruntime::narrow(this->n_trees_)); - for (int64_t i = 0; i < N; ++i) { - scores[batch_num * SafeInt(N) + i].resize(onnxruntime::narrow(n_targets_or_classes_), {0, 0}); - } - for (auto j = work.start; j < work.end; ++j) { - for (int64_t i = 0; i < N; ++i) { - agg.ProcessTreeNodePrediction(scores[batch_num * SafeInt(N) + i], - *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + int64_t end_n, begin_n = 0; + while (begin_n < N) { + end_n = std::min(N, begin_n + parallel_tree_N_); + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, &agg, &scores, num_threads, x_data, N, stride, begin_n, end_n](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, onnxruntime::narrow(this->n_trees_)); + for (int64_t i = begin_n; i < end_n; ++i) { + scores[batch_num * SafeInt(N) + i].resize(onnxruntime::narrow(n_targets_or_classes_), {0, 0}); } - } - }); - + for (auto j = work.start; j < work.end; ++j) { + for (int64_t i = begin_n; i < end_n; ++i) { + agg.ProcessTreeNodePrediction(scores[batch_num * SafeInt(N) + i], + *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + } + }); + begin_n = end_n; + } concurrency::ThreadPool::TrySimpleParallelFor( ttp, num_threads, @@ -661,11 +704,11 @@ class TreeEnsembleCommonClassifier : public TreeEnsembleCommon class_labels_; public: - virtual Status Init(const OpKernelInfo& info); virtual Status compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z, Tensor* label) const; Status Init(int parallel_tree, + int parallel_tree_N, int parallel_N, const std::string& aggregate_function, const std::vector& base_values, @@ -691,11 +734,10 @@ class TreeEnsembleCommonClassifier : public TreeEnsembleCommon& classlabels_int64s); }; - template Status TreeEnsembleCommonClassifier::Init(const OpKernelInfo& info) { std::vector base_values_as_tensor, nodes_hitrates_as_tensor, - nodes_values_as_tensor, class_weights_as_tensor; + nodes_values_as_tensor, class_weights_as_tensor; #if !defined(ORT_MINIMAL_BUILD) ORT_THROW_IF_ERROR(GetVectorAttrsOrDefault(info, "base_values_as_tensor", base_values_as_tensor)); ORT_THROW_IF_ERROR(GetVectorAttrsOrDefault(info, "nodes_hitrates_as_tensor", nodes_hitrates_as_tensor)); @@ -705,6 +747,7 @@ Status TreeEnsembleCommonClassifier::Init( return Init( 80, + 128, 50, info.GetAttrOrDefault("aggregate_function", "SUM"), info.GetAttrsOrDefault("base_values"), @@ -731,54 +774,58 @@ Status TreeEnsembleCommonClassifier::Init( } template -Status TreeEnsembleCommonClassifier::Init(int parallel_tree, - int parallel_N, - const std::string& aggregate_function, - const std::vector& base_values, - const std::vector& base_values_as_tensor, - const std::vector& nodes_falsenodeids, - const std::vector& nodes_featureids, - const std::vector& nodes_hitrates, - const std::vector& nodes_hitrates_as_tensor, - const std::vector& nodes_missing_value_tracks_true, - const std::vector& nodes_modes, - const std::vector& nodes_nodeids, - const std::vector& nodes_treeids, - const std::vector& nodes_truenodeids, - const std::vector& nodes_values, - const std::vector& nodes_values_as_tensor, - const std::string& post_transform, - const std::vector& class_ids, - const std::vector& class_nodeids, - const std::vector& class_treeids, - const std::vector& class_weights, - const std::vector& class_weights_as_tensor, - const std::vector& classlabels_strings, - const std::vector& classlabels_int64s) { - auto status = TreeEnsembleCommon::Init(parallel_tree, - parallel_N, - aggregate_function, - base_values, - base_values_as_tensor, - classlabels_strings.empty() ? classlabels_int64s.size() - : classlabels_strings.size(), - nodes_falsenodeids, - nodes_featureids, - nodes_hitrates, - nodes_hitrates_as_tensor, - nodes_missing_value_tracks_true, - nodes_modes, - nodes_nodeids, - nodes_treeids, - nodes_truenodeids, - nodes_values, - nodes_values_as_tensor, - post_transform, - class_ids, - class_nodeids, - class_treeids, - class_weights, - class_weights_as_tensor); +Status TreeEnsembleCommonClassifier::Init( + int parallel_tree, + int parallel_tree_N, + int parallel_N, + const std::string& aggregate_function, + const std::vector& base_values, + const std::vector& base_values_as_tensor, + const std::vector& nodes_falsenodeids, + const std::vector& nodes_featureids, + const std::vector& nodes_hitrates, + const std::vector& nodes_hitrates_as_tensor, + const std::vector& nodes_missing_value_tracks_true, + const std::vector& nodes_modes, + const std::vector& nodes_nodeids, + const std::vector& nodes_treeids, + const std::vector& nodes_truenodeids, + const std::vector& nodes_values, + const std::vector& nodes_values_as_tensor, + const std::string& post_transform, + const std::vector& class_ids, + const std::vector& class_nodeids, + const std::vector& class_treeids, + const std::vector& class_weights, + const std::vector& class_weights_as_tensor, + const std::vector& classlabels_strings, + const std::vector& classlabels_int64s) { + auto status = TreeEnsembleCommon::Init( + parallel_tree, + parallel_tree_N, + parallel_N, + aggregate_function, + base_values, + base_values_as_tensor, + classlabels_strings.empty() ? classlabels_int64s.size() + : classlabels_strings.size(), + nodes_falsenodeids, + nodes_featureids, + nodes_hitrates, + nodes_hitrates_as_tensor, + nodes_missing_value_tracks_true, + nodes_modes, + nodes_nodeids, + nodes_treeids, + nodes_truenodeids, + nodes_values, + nodes_values_as_tensor, + post_transform, + class_ids, + class_nodeids, + class_treeids, + class_weights, + class_weights_as_tensor); ORT_RETURN_IF_ERROR(status); classlabels_strings_ = classlabels_strings; diff --git a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc index 3879222812..e27f300642 100644 --- a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc @@ -274,6 +274,7 @@ TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeC2) { std::vector base_values{0.f, 0.f}; GenTreeAndRunTest(1, X, base_values, results, "AVERAGE", false, 200, 130); // section C2 GenTreeAndRunTest(3, X, base_values, results, "AVERAGE", false, 200, 130); // section C2 + GenTreeAndRunTest(3, X, base_values, results, "AVERAGE", false, 400, 130); // section C2 } TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeD2) { @@ -312,6 +313,7 @@ TEST(MLOpTest, TreeRegressorMultiTargetMin) { GenTreeAndRunTest(3, X, base_values, results, "MIN", false); GenTreeAndRunTest(1, X, base_values, results, "MIN", true); GenTreeAndRunTest(3, X, base_values, results, "MIN", true); + GenTreeAndRunTest(3, X, base_values, results, "MIN", false, 109 * 8); } TEST(MLOpTest, TreeRegressorMultiTargetMax) { @@ -550,6 +552,7 @@ TEST(MLOpTest, TreeRegressorSingleTargetSum) { GenTreeAndRunTest1(3, "SUM", false); GenTreeAndRunTest1(1, "SUM", true); GenTreeAndRunTest1(3, "SUM", true); + GenTreeAndRunTest1(3, "SUM", false, 1023); } TEST(MLOpTest, TreeRegressorSingleTargetSum_as_tensor) {