`thrust::all_of` is slower than a naive reduction
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
In a `thrust::all_of`, when the first element that violates the predicate is discovered, the computation can be aborted, i.e., an "early exit".
For example, imagine you are given a `thrust::device_vector` and want to check if any of the values are negative. You could do this with a `thrust::all_of` or with a `thrust::count_if`:
```
thrust::device_vector values(...);
bool all_positive = thrust::all_of(values.begin(), values.end(), [](auto v){return v > 0;});
bool all_positive = values.size() == thrust::count_if(values.begin(), values.end(), [](auto v){return v > 0;});
```
`count_if` must read everything in `values`, whereas `all_of` can shortcut if an early exit exists. Therefore, I would expect `all_of` to out perform `count_if` when one or more negative values exist. If no negative values are present, then both `all_of` and `count_if` must read everything in `values` and I would expect their performance to be roughly equivalent.
However, this is not the case. I have found that the performance of `thrust::all_of` is extremely erratic with a 10x difference between the best and worst performance. Furthermore, an `all_of` is _always_ slower than a naive reduction as in `count_if`.
Here are the results of performing 100 trials of the example I described above on an input size of 100,000,000 million `int64_t` elements on a GV100.
### No Early Exit
| | mean (us) | min (us) | max (us) |
|:----------:|:---------:|:--------:|:--------:|
| `all_of` | 75269 | 56403 | 104922 |
| `count_if` | 3124 | 1686 | 4413 |
### Single Early Exit
| | mean (us) | min (us) | max (us) |
|:----------:|:---------:|:--------:|:--------:|
| `all_of` | 51620 | 9346 | 370845 |
| `count_if` | 3100 | 1703 | 5158 |
As you can see, whether or not an early exit exists, `all_of` is always significantly slower than a `count_if`.
Looking at the profile of `all_of` (attached), it appears that the reason it is so slow is because a single invocation of `all_of` results in ~50 invocations of `DeviceReduceKernel`. I suspect this is because the implementation of `all_of` does a set of batched reductions in attempt to avoid reading the entire input when an early exit exists. However, launching all of these small kernels (each with their own allocation/free) results in a significant amount of overhead. This overhead is exacerbated by the fact that each batch is executed on the same stream, meaning there is no overlap or concurrency between batches.
I suspect a better implementation would launch a single kernel, where threads occasionally poll an atomic flag to check if an early exit exists, at which point they exit the computation. Or, forgo an attempt at an early exit and just do the naive reduction like in `count_if`.

[nsys_profile.zip](https://github.com/thrust/thrust/files/3623728/nsys_profile.zip)
Reproducer code:
```
// compile with `nvcc --std=c++14 -O3 --expt-extended-lambda thrust_logical.cu -o thrust_logical -lnvToolsExt`
#include
#include
#include
#include
#include
#include
#include
template
struct time_result {
T min{std::numeric_limits::max()};
T max{std::numeric_limits::lowest()};
T mean{0};
T sum{0};
std::size_t count{0};
void add_measurement(T new_duration) {
++count;
sum += new_duration;
mean = sum / count;
min = std::min(min, new_duration);
max = std::max(max, new_duration);
}
std::string to_string() {
return std::string{
"count: " + std::to_string(count) + " mean: " + std::to_string(mean) +
" min: " + std::to_string(min) + " max: " + std::to_string(max)};
}
};
template
typename Duration::rep time_it(std::string const& name, F&& fun,
Args&&... args) {
const auto begin = std::chrono::high_resolution_clock::now();
nvtxRangePushA(name.c_str());
std::forward(fun)(std::forward(args)...);
nvtxRangePop();
const auto end = std::chrono::high_resolution_clock::now();
return std::chrono::duration_cast(end - begin).count();
}
template
auto time_trial(std::string const& name, std::size_t num_trials, InputGenerator&& generator, F&& f,
Args&&... args) {
time_result result{};
for (auto i = 0; i < num_trials; ++i) {
auto input = generator();
result.add_measurement(time_it(name, std::forward(f), input,
std::forward(args)...));
}
return result;
}
struct is_positive {
template
bool __device__ operator()(T v) {
return v > 0;
}
};
int main(void) {
constexpr std::size_t input_size{100'000'000};
constexpr std::size_t num_trials{100};
auto all_of = [](auto const& values) {
return thrust::all_of(thrust::device, values.begin(), values.end(),
is_positive{});
};
auto count_if = [](auto const& values) {
return thrust::count_if(thrust::device, values.begin(), values.end(),
is_positive{});
};
auto no_early_out = []() {
nvtxRangePushA("no early out input");
thrust::device_vector values(input_size, 1);
cudaDeviceSynchronize();
nvtxRangePop();
return values;
};
auto early_out = []() {
nvtxRangePushA("early out input");
thrust::device_vector values(input_size, 1);
thrust::default_random_engine engine(
std::chrono::high_resolution_clock::now().time_since_epoch().count());
thrust::uniform_int_distribution distribution{0, input_size};
auto random_location = distribution(engine);
values[random_location] = -1;
cudaDeviceSynchronize();
nvtxRangePop();
return values;
};
std::cout << "No early out(us):\n";
std::cout << "all of: "
<< time_trial("all_of", num_trials, no_early_out, all_of).to_string()
<< std::endl;
std::cout << "count if: "
<< time_trial("count_if", num_trials, no_early_out, count_if).to_string()
<< std::endl;
std::cout << std::endl << "With early out(us):\n";
std::cout << "all of: "
<< time_trial("all_of", num_trials, early_out, all_of).to_string()
<< std::endl;
std::cout << "count if: "
<< time_trial("count_if", num_trials, early_out, count_if).to_string()
<< std::endl;
return 0;
}
```
### Tasks
- [ ] https://github.com/NVIDIA/cccl/issues/2113
Contributor guide
Assessment
This issue has not been assessed yet.