NVIDIA / NVIDIA/cccl

[FEA]: For Hopper provide a barrier completion function that performs an expect-tx

Open
#420 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

### Is this a duplicate?

- [x] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)

### Area

libcu++

### Is your feature request related to a problem? Please describe.

In Hopper barriers often expect the same amount of bytes on every barrier phase, but the user has to manually unbalance the count which is an error prone operation.

### Describe the solution you'd like

I think we should provide a completion function that automatically unbalances the count:

```c++
cuda::barrier bar(expected_count, cuda::expect_tx_fn cf(nbytes));
```

That way users don't need to manually unbalance the count. For example, if the barrier has an `expected_count == 1`, and the producer never arrives at the barrier, the consumer only needs to arrive normally, not with an `arrive_tx` expecting a free parameter:

```c++
__shared__ cuda::barrier bar;
if (threadIdx.x == some thread during init)
init(bar, expected_count, cuda::expect_tx_fn cf(nbytes));

// producer threads during a phase
for (...) {
other_bar.wait(...);
cuda::memcpy_async_tx(bar, ...); // requires manual re-balance
cuda::memcpy_async_tx(bar, ...); // requires manual re-balance
}

// consumer threads during a phase
for (...) {
bar.arrive_and_wait(); // automatically re-balance the phase
use_data();
other_bar.arrive();
}
```

Without this, the consumer threads need to:
- pick one thread to perform an expect tx
- carefully orchestrate arrivals with that expect tx

```c++
for (...) {
if (consumer_thread == x) bar.wait(cuda::arrive_tx(bar, 1, nbytes));
else bar.arrive_and_wait();
use_data();
other_bar.arrive();
}
```

This is quite error prone, not only can `nbytes` change within a phase, but it is very easy for users to incorrectly do:

```c++
for (...) {
bar.wait(cuda::arrive_tx(bar, 1, nbytes)); // nconsumers * nbytes now --> program hangs
use_data();
other_bar.arrive();
}
```

or

```c++
for (...) {
bar.wait(cuda::arrive_tx(bar, 1, nbytes / nconsumers)); // nbytes % nconsumers != 0 --> program hangs
use_data();
other_bar.arrive();
}
```

### Describe alternatives you've considered

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.