| 810 | |
| 811 | template <typename Batch, typename IsValid> |
| 812 | std::vector<bst_feature_t> LoadBalance(Batch const &batch, size_t nnz, bst_feature_t n_columns, |
| 813 | size_t const nthreads, IsValid &&is_valid) { |
| 814 | /* Some sparse datasets have their mass concentrating on small number of features. To |
| 815 | * avoid waiting for a few threads running forever, we here distribute different number |
| 816 | * of columns to different threads according to number of entries. |
| 817 | */ |
| 818 | size_t const total_entries = nnz; |
| 819 | size_t const entries_per_thread = DivRoundUp(total_entries, nthreads); |
| 820 | |
| 821 | // Need to calculate the size for each batch. |
| 822 | std::vector<bst_idx_t> entries_per_columns = CalcColumnSize(batch, n_columns, nthreads, is_valid); |
| 823 | std::vector<bst_feature_t> cols_ptr(nthreads + 1, 0); |
| 824 | size_t count{0}; |
| 825 | size_t current_thread{1}; |
| 826 | |
| 827 | for (auto col : entries_per_columns) { |
| 828 | cols_ptr.at(current_thread)++; // add one column to thread |
| 829 | count += col; |
| 830 | CHECK_LE(count, total_entries); |
| 831 | if (count > entries_per_thread) { |
| 832 | current_thread++; |
| 833 | count = 0; |
| 834 | cols_ptr.at(current_thread) = cols_ptr[current_thread - 1]; |
| 835 | } |
| 836 | } |
| 837 | // Idle threads. |
| 838 | for (; current_thread < cols_ptr.size() - 1; ++current_thread) { |
| 839 | cols_ptr[current_thread + 1] = cols_ptr[current_thread]; |
| 840 | } |
| 841 | return cols_ptr; |
| 842 | } |
| 843 | |
| 844 | /*! |
| 845 | * A sketch matrix storing sketches for each feature. |