src/third_party/llvm-project/openmp/libomptarget/deviceRTLs/nvptx/docs/ReductionDesign.txt

**Design document for OpenMP reductions on the GPU** 

//Abstract: //In this document we summarize the new design for an OpenMP
implementation of reductions on NVIDIA GPUs.  This document comprises
* a succinct background review,
* an introduction to the decoupling of reduction algorithm and
    data-structure-specific processing routines,
* detailed illustrations of reduction algorithms used and
* a brief overview of steps we have made beyond the last implementation.

**Problem Review**

Consider a typical OpenMP program with reduction pragma.

```
    double foo, bar;
    #pragma omp parallel for reduction(+:foo, bar)
    for (int i = 0; i < N; i++) {
      foo+=A[i]; bar+=B[i];
    }
```
where 'foo' and 'bar' are reduced across all threads in the parallel region.
Our primary goal is to efficiently aggregate the values of foo and bar in
such manner that
* makes the compiler logically concise.
* efficiently reduces within warps, threads, blocks and the device.

**Introduction to Decoupling**
In this section we address the problem of making the compiler
//logically concise// by partitioning the task of reduction into two broad
categories: data-structure specific routines and algorithmic routines.

The previous reduction implementation was highly coupled with
the specificity of the reduction element data structures (e.g., sizes, data
types) and operators of the reduction (e.g., addition, multiplication). In
our implementation we strive to decouple them. In our final implementations,
we could remove all template functions in our runtime system.

The (simplified) pseudo code generated by LLVM is as follows:

```
    1. Create private copies of variables: foo_p, bar_p
    2. Each thread reduces the chunk of A and B assigned to it and writes
       to foo_p and bar_p respectively.
    3. ret = kmpc_nvptx_reduce_nowait(..., reduceData, shuffleReduceFn, 
               interWarpCpyFn)
        where:
        struct ReduceData {
          double *foo;
          double *bar;
        } reduceData
        reduceData.foo = &foo_p
        reduceData.bar = &bar_p

        shuffleReduceFn and interWarpCpyFn are two auxiliary functions
        generated to aid the runtime performing algorithmic steps
        while being data-structure agnostic about ReduceData.

        In particular, shuffleReduceFn is a function that takes the following
        inputs:
        a. local copy of ReduceData
        b. its lane_id
        c. the offset of the lane_id which hosts a remote ReduceData
                relative to the current one
        d. an algorithm version paramter determining which reduction
                algorithm to use.
        This shuffleReduceFn retrieves the remote ReduceData through shuffle
        intrinsics and  reduces, using the algorithm specified by the 4th
        parameter, the local ReduceData and with the remote ReduceData element
        wise, and places the resultant values into the local ReduceData.

        Different reduction algorithms are implemented with different runtime
        functions, but they all make calls to this same shuffleReduceFn to
        perform the essential reduction step. Therefore, based on the 4th
        parameter, this shuffleReduceFn will behave slightly differently to
        cooperate with the runtime function to ensure correctness under
        different circumstances.

        InterWarpCpyFn, as the name suggests, is a function that copies data
        across warps. Its function is to tunnel all the thread private
        ReduceData that is already reduced within a warp to a lane in the first
        warp with minimal shared memory footprint. This is an essential step to
        prepare for the last step of a block reduction.

        (Warp, block, device level reduction routines that utilize these
        auxiliary functions will be discussed in the next section.)

    4. if ret == 1:
        The master thread stores the reduced result in the globals.
        foo += reduceData.foo; bar += reduceData.bar
```

**Reduction Algorithms**

On the warp level, we have three versions of the algorithms:

1. Full Warp Reduction

```
gpu_regular_warp_reduce(void *reduce_data,
                        kmp_ShuffleReductFctPtr ShuffleReduceFn) {
  for (int offset = WARPSIZE/2; offset > 0; offset /= 2)
    ShuffleReduceFn(reduce_data, 0, offset, 0);
}
```
ShuffleReduceFn is used here with lane_id set to 0 because it is not used
therefore we save instructions by not retrieving lane_id from the corresponding
special registers. The 4th parameters, which represents the version of the
algorithm being used here, is set to 0 to signify full warp reduction.

In this version specified (=0), the ShuffleReduceFn behaves, per element, as
follows:

```
//reduce_elem refers to an element in the local ReduceData
//remote_elem is retrieved from a remote lane
remote_elem = shuffle_down(reduce_elem, offset, 32);
reduce_elem = reduce_elem @ remote_elem;

```

An illustration of this algorithm operating on a hypothetical 8-lane full-warp
would be:
{F74}
The coloring invariant follows that elements with the same color will be
combined and reduced in the next reduction step. As can be observed, no overhead
is present, exactly log(2, N) steps are needed.

2. Contiguous Full Warp Reduction
```
gpu_irregular_warp_reduce(void *reduce_data,
                          kmp_ShuffleReductFctPtr ShuffleReduceFn, int size,
                          int lane_id) {
  int curr_size;
  int offset;
    curr_size = size;
    mask = curr_size/2;
    while (offset>0) {
      ShuffleReduceFn(reduce_data, lane_id, offset, 1);
      curr_size = (curr_size+1)/2;
      offset = curr_size/2;
    }
}
```

In this version specified (=1), the ShuffleReduceFn behaves, per element, as
follows:
```
//reduce_elem refers to an element in the local ReduceData
//remote_elem is retrieved from a remote lane
remote_elem = shuffle_down(reduce_elem, offset, 32);
if (lane_id < offset) {
    reduce_elem = reduce_elem @ remote_elem
} else {
    reduce_elem = remote_elem
}
```

An important invariant (also a restriction on the starting state of the
reduction) is that this algorithm assumes that all unused ReduceData are
located in a contiguous subset of threads in a warp starting from lane 0.

With the presence of a trailing active lane with an odd-numbered lane
id, its value will not be aggregated with any other lane. Therefore,
in order to preserve the invariant, such ReduceData is copied to the first lane
whose thread-local ReduceData has already being used in a previous reduction
and would therefore be useless otherwise.

An illustration of this algorithm operating on a hypothetical 8-lane partial
warp woud be:
{F75}

As illustrated, this version of the algorithm introduces overhead whenever
we have odd number of participating lanes in any reduction step to
copy data between lanes.

3. Dispersed Partial Warp Reduction
```
gpu_irregular_simt_reduce(void *reduce_data,
                          kmp_ShuffleReductFctPtr ShuffleReduceFn) {
  int size, remote_id;
  int logical_lane_id = find_number_of_dispersed_active_lanes_before_me() * 2;
  do {
      remote_id = find_the_next_active_lane_id_right_after_me();
      // the above function returns 0 of no active lane
      // is present right after the current thread.
      size = get_number_of_active_lanes_in_this_warp();
      logical_lane_id /= 2;
      ShuffleReduceFn(reduce_data, logical_lane_id, remote_id-1-threadIdx.x, 2);
  } while (logical_lane_id % 2 == 0 && size > 1);
```

There is no assumption made about the initial state of the reduction.
Any number of lanes (>=1) could be active at any position. The reduction
result is kept in the first active lane.

In this version specified (=2), the ShuffleReduceFn behaves, per element, as
follows:
```
//reduce_elem refers to an element in the local ReduceData
//remote_elem is retrieved from a remote lane
remote_elem = shuffle_down(reduce_elem, offset, 32);
if (LaneId % 2 == 0 && Offset > 0) {
    reduce_elem = reduce_elem @ remote_elem
} else {
    reduce_elem = remote_elem
}
```
We will proceed with a brief explanation for some arguments passed in,
it is important to notice that, in this section, we will introduce the
concept of logical_lane_id, and it is important to distinguish it
from physical lane_id as defined by nvidia.
1. //logical_lane_id//: as the name suggests, it refers to the calculated
    lane_id (instead of the physical one defined by nvidia) that would make
    our algorithm logically concise. A thread with logical_lane_id k means
    there are (k-1) threads before it.
2. //remote_id-1-threadIdx.x//: remote_id is indeed the nvidia-defined lane
    id of the remote lane from which we will retrieve the ReduceData. We
    subtract (threadIdx+1) from it because we would like to maintain only one
    underlying shuffle intrinsic (which is used to communicate among lanes in a
    warp). This particular version of shuffle intrinsic we take accepts only
    offsets, instead of absolute lane_id. Therefore the subtraction is performed
    on the absolute lane_id we calculated to obtain the offset.

This algorithm is slightly different in 2 ways and it is not, conceptually, a
generalization of the above algorithms.
1. It reduces elements close to each other. For instance, values in the 0th lane
    is to be combined with that of the 1st lane; values in the 2nd lane is to be
    combined with that of the 3rd lane. We did not use the previous algorithm
    where the first half of the (partial) warp is reduced with the second half
    of the (partial) warp. This is because, the mapping
    f(x): logical_lane_id -> physical_lane_id;
    can be easily calculated whereas its inverse
    f^-1(x): physical_lane_id -> logical_lane_id
    cannot and performing such reduction requires the inverse to be known.
2. Because this algorithm is agnostic about the positions of the lanes that are
    active, we do not need to perform the coping step as in the second
    algorithm.
An illustrative run would look like
{F76}
As observed, overhead is high because in each and every step of reduction,
logical_lane_id is recalculated; so is the remote_id.

On a block level, we have implemented the following block reduce algorithm:

```
gpu_irregular_block_reduce(void *reduce_data,
              kmp_ShuffleReductFctPtr shuflReduceFn,
              kmp_InterWarpCopyFctPtr interWarpCpyFn,
              int size) {

  int wid = threadIdx.x/WARPSIZE;
  int lane_id = threadIdx.x%WARPSIZE;

  int warp_needed = (size+WARPSIZE-1)/WARPSIZE; //ceiling of division

  unsigned tnum = __ballot(1);
  int thread_num = __popc(tnum);

    //full warp reduction
    if (thread_num == WARPSIZE) {
      gpu_regular_warp_reduce(reduce_data, shuflReduceFn);
    }
    //partial warp reduction
    if (thread_num < WARPSIZE) {
        gpu_irregular_warp_reduce(reduce_data, shuflReduceFn, thread_num,
                                  lane_id);
    }
    //Gather all the reduced values from each warp
    //to the first warp
    //named_barrier inside this function to ensure
    //correctness. It is effectively a sync_thread
    //that won't deadlock.
    interWarpCpyFn(reduce_data, warp_needed);

    //This is to reduce data gathered from each "warp master".
    if (wid==0) {
        gpu_irregular_warp_reduce(reduce_data, shuflReduceFn, warp_needed,
                                  lane_id);
    }

  return;
}
```
In this function, no ShuffleReduceFn is directly called as it makes calls
to various versions of the warp-reduction functions. It first reduces
ReduceData warp by warp; in the end, we end up with the number of
ReduceData equal to the number of warps present in this thread
block. We then proceed to gather all such ReduceData to the first warp.

As observed, in this algorithm we make use of the function InterWarpCpyFn,
which copies data from each of the "warp master" (0th lane of each warp, where 
a warp-reduced ReduceData is held) to the 0th warp. This step reduces (in a
mathematical sense) the problem of reduction across warp masters in a block to
the problem of warp reduction which we already have solutions to.

We can thus completely avoid the use of atomics to reduce in a threadblock.

**Efficient Cross Block Reduce**

The next challenge is to reduce values across threadblocks.  We aim to do this
without atomics or critical sections.

Let a kernel be started with TB threadblocks.
Let the GPU have S SMs.
There can be at most N active threadblocks per SM at any time.

Consider a threadblock tb (tb < TB) running on SM s (s < SM).  'tb' is one of
at most 'N' active threadblocks on SM s.  Let each threadblock active on an SM
be given an instance identifier id (0 <= id < N).  Therefore, the tuple (s, id)
uniquely identifies an active threadblock on the GPU.

To efficiently implement cross block reduce, we first allocate an array for
each value to be reduced of size S*N (which is the maximum number of active
threadblocks at any time on the device).

Each threadblock reduces its value to slot [s][id].  This can be done without
locking since no other threadblock can write to the same slot concurrently.

As a final stage, we reduce the values in the array as follows:

```
// Compiler generated wrapper function for each target region with a reduction
clause.
target_function_wrapper(map_args, reduction_array)  <--- start with 1 team and 1
   thread.
  // Use dynamic parallelism to launch M teams, N threads as requested by the
  user to execute the target region.

  target_function<<M, N>>(map_args)

  Reduce values in reduction_array

```

**Comparison with Last Version**


The (simplified) pseudo code generated by LLVM on the host is as follows:


```
    1. Create private copies of variables: foo_p, bar_p
    2. Each thread reduces the chunk of A and B assigned to it and writes
       to foo_p and bar_p respectively.
    3. ret = kmpc_reduce_nowait(..., reduceData, reduceFn, lock)
        where:
        struct ReduceData {
          double *foo;
          double *bar;
        } reduceData
        reduceData.foo = &foo_p
        reduceData.bar = &bar_p

        reduceFn is a pointer to a function that takes in two inputs
        of type ReduceData, "reduces" them element wise, and places the
        result in the first input:
        reduceFn(ReduceData *a, ReduceData *b)
          a = a @ b

        Every thread in the parallel region calls kmpc_reduce_nowait with
        its private copy of reduceData.  The runtime reduces across the
        threads (using tree reduction on the operator 'reduceFn?) and stores
        the final result in the master thread if successful.
    4. if ret == 1:
        The master thread stores the reduced result in the globals.
        foo += reduceData.foo; bar += reduceData.bar
    5. else if ret == 2:
        In this case kmpc_reduce_nowait() could not use tree reduction,
        so use atomics instead:
        each thread atomically writes to foo
        each thread atomically writes to bar
```

On a GPU, a similar reduction may need to be performed across SIMT threads,
warps, and threadblocks.  The challenge is to do so efficiently in a fashion
that is compatible with the LLVM OpenMP implementation.

In the previously released 0.1 version of the LLVM OpenMP compiler for GPUs,
the salient steps of the code generated are as follows:


```
    1. Create private copies of variables: foo_p, bar_p
    2. Each thread reduces the chunk of A and B assigned to it and writes
       to foo_p and bar_p respectively.
    3. ret = kmpc_reduce_nowait(..., reduceData, reduceFn, lock)
        status = can_block_reduce()
        if status == 1:
          reduce efficiently to thread 0 using shuffles and shared memory.
          return 1
        else
          cannot use efficient block reduction, fallback to atomics
          return 2
    4. if ret == 1:
        The master thread stores the reduced result in the globals.
        foo += reduceData.foo; bar += reduceData.bar
    5. else if ret == 2:
        In this case kmpc_reduce_nowait() could not use tree reduction,
        so use atomics instead:
        each thread atomically writes to foo
        each thread atomically writes to bar
```

The function can_block_reduce() is defined as follows:


```
int32_t can_block_reduce() {
  int tid = GetThreadIdInTeam();
  int nt = GetNumberOfOmpThreads(tid);
  if (nt != blockDim.x)
    return 0;
  unsigned tnum = __ballot(1);
  if (tnum != (~0x0)) {
    return 0;
  }
  return 1;
}
```

This function permits the use of the efficient block reduction algorithm
using shuffles and shared memory (return 1) only if (a) all SIMT threads in
a warp are active (i.e., number of threads in the parallel region is a
multiple of 32) and (b) the number of threads in the parallel region
(set by the num_threads clause) equals blockDim.x.

If either of these preconditions is not true, each thread in the threadblock
updates the global value using atomics.

Atomics and compare-and-swap operations are expensive on many threaded
architectures such as GPUs and we must avoid them completely.


**Appendix: Implementation Details**


```
// Compiler generated function.
reduceFn(ReduceData *a, ReduceData *b)
  a->foo = a->foo + b->foo
  a->bar = a->bar + b->bar

// Compiler generated function.
swapAndReduceFn(ReduceData *thread_private, int lane)
  ReduceData *remote = new ReduceData()
  remote->foo = shuffle_double(thread_private->foo, lane)
  remote->bar = shuffle_double(thread_private->bar, lane)
  reduceFn(thread_private, remote)

// OMP runtime function.
warpReduce_regular(ReduceData *thread_private, Fn *swapAndReduceFn):
  offset = 16
  while (offset > 0)
    swapAndReduceFn(thread_private, offset)
    offset /= 2

// OMP runtime function.
warpReduce_irregular():
  ...

// OMP runtime function.
kmpc_reduce_warp(reduceData, swapAndReduceFn)
  if all_lanes_active:
    warpReduce_regular(reduceData, swapAndReduceFn)
  else:
    warpReduce_irregular(reduceData, swapAndReduceFn)
  if in_simd_region:
    // all done, reduce to global in simd lane 0
    return 1
  else if in_parallel_region:
    // done reducing to one value per warp, now reduce across warps
    return 3

// OMP runtime function; one for each basic type.
kmpc_reduce_block_double(double *a)
  if lane == 0:
    shared[wid] = *a
  named_barrier(1, num_threads)
  if wid == 0
    block_reduce(shared)
  if lane == 0
    *a = shared[0]
  named_barrier(1, num_threads)
  if wid == 0 and lane == 0
    return 1  // write back reduced result
  else
    return 0  // don't do anything

```



```
// Compiler generated code.
    1. Create private copies of variables: foo_p, bar_p
    2. Each thread reduces the chunk of A and B assigned to it and writes
       to foo_p and bar_p respectively.
    3. ret = kmpc_reduce_warp(reduceData, swapAndReduceFn)
    4. if ret == 1:
        The master thread stores the reduced result in the globals.
        foo += reduceData.foo; bar += reduceData.bar
    5. else if ret == 3:
        ret = block_reduce_double(reduceData.foo)
        if ret == 1:
          foo += reduceData.foo
        ret = block_reduce_double(reduceData.bar)
        if ret == 1:
          bar += reduceData.bar
```

**Notes**

    1. This scheme requires that the CUDA OMP runtime can call llvm generated
       functions. This functionality now works.
    2. If the user inlines the CUDA OMP runtime bitcode, all of the machinery
       (including calls through function pointers) are optimized away.
    3. If we are reducing multiple to multiple variables in a parallel region,
       the reduce operations are all performed in warpReduce_[ir]regular(). This
       results in more instructions in the loop and should result in fewer
       stalls due to data dependencies.  Unfortunately we cannot do the same in
       kmpc_reduce_block_double() without increasing shared memory usage.