[feature request] Faster Split (data partition)

[guolinke]

Besides of histogram construction, the data partition is another time-consuming part in LightGBM. This process is to update the row set for each leaf.

The code is in the following:

LightGBM/src/treelearner/data_partition.hpp

Lines 113 to 172 in 997a469

	void Split(int leaf, const Dataset* dataset, int feature,
	const uint32_t* threshold, int num_threshold, bool default_left,
	int right_leaf) {
	Common::FunctionTimer fun_timer("DataPartition::Split", global_timer);
	const data_size_t min_inner_size = 512;
	// get leaf boundary
	const data_size_t begin = leaf_begin_[leaf];
	const data_size_t cnt = leaf_count_[leaf];
	const int nblock = std::min(num_threads_, (cnt + min_inner_size - 1) / min_inner_size);
	data_size_t inner_size = SIZE_ALIGNED((cnt + nblock - 1) / nblock);
	auto left_start = indices_.data() + begin;
	global_timer.Start("DataPartition::Split.MT");
	// split data multi-threading
	OMP_INIT_EX();
	#pragma omp parallel for schedule(static, 1)
	for (int i = 0; i < nblock; ++i) {
	OMP_LOOP_EX_BEGIN();
	data_size_t cur_start = i * inner_size;
	data_size_t cur_cnt = std::min(inner_size, cnt - cur_start);
	if (cur_cnt <= 0) {
	left_cnts_buf_[i] = 0;
	right_cnts_buf_[i] = 0;
	continue;
	}
	// split data inner, reduce the times of function called
	data_size_t cur_left_count = dataset->Split(feature, threshold, num_threshold, default_left,
	left_start + cur_start, cur_cnt,
	temp_left_indices_.data() + cur_start,
	temp_right_indices_.data() + cur_start);
	offsets_buf_[i] = cur_start;
	left_cnts_buf_[i] = cur_left_count;
	right_cnts_buf_[i] = cur_cnt - cur_left_count;
	OMP_LOOP_EX_END();
	}
	OMP_THROW_EX();
	global_timer.Stop("DataPartition::Split.MT");
	global_timer.Start("DataPartition::Split.Merge");
	left_write_pos_buf_[0] = 0;
	right_write_pos_buf_[0] = 0;
	for (int i = 1; i < nblock; ++i) {
	left_write_pos_buf_[i] = left_write_pos_buf_[i - 1] + left_cnts_buf_[i - 1];
	right_write_pos_buf_[i] = right_write_pos_buf_[i - 1] + right_cnts_buf_[i - 1];
	}
	data_size_t left_cnt = left_write_pos_buf_[nblock - 1] + left_cnts_buf_[nblock - 1];
	auto right_start = left_start + left_cnt;
	#pragma omp parallel for schedule(static)
	for (int i = 0; i < nblock; ++i) {
	std::copy_n(temp_left_indices_.data() + offsets_buf_[i],
	left_cnts_buf_[i], left_start + left_write_pos_buf_[i]);
	std::copy_n(temp_right_indices_.data() + offsets_buf_[i],
	right_cnts_buf_[i], right_start + right_write_pos_buf_[i]);
	}
	// update leaf boundary
	leaf_count_[leaf] = left_cnt;
	leaf_begin_[right_leaf] = left_cnt + begin;
	leaf_count_[right_leaf] = cnt - left_cnt;
	global_timer.Stop("DataPartition::Split.Merge");
	}

[StrikerRUS]

Closed in favor of being in #2302. We decided to keep all feature requests in one place.

Welcome to contribute this feature! Please re-open this issue (or post a comment if you are not a topic starter) if you are actively working on implementing this feature.