Repository: Project-MONAI/MONAI
File: monai/_extensions/gmm/gmm_cuda.cu (CUDA)

Code before:
      if (my_alpha != -1) {
        if (gaussian_index == (my_alpha & 15) + (my_alpha >> 4) * MIXTURE_COUNT) {
          float feature[CHANNEL_COUNT + 1];

          feature[0] = 1;

          for (int i = 0; i < CHANNEL_COUNT; i++) {
            feature[i + 1] = g_batch_image[global_index + i * element_count];
          }

          for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) {
            for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) {
              matrix[index] += feature[i] * feature[j];
            }
          }
        }
      }
    }
  }

  __syncthreads();

  for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) {
    float matrix_component = matrix[i];
#ifdef __HIP_PLATFORM_AMD__
    matrix_component += __shfl_down(matrix_component, 16);
    matrix_component += __shfl_down(matrix_component, 8);
    matrix_component += __shfl_down(matrix_component, 4);
    matrix_component += __shfl_down(matrix_component, 2);
    matrix_component += __shfl_down(matrix_component, 1);
#else
    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16);
    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8);
    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4);
    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2);
    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1);
#endif
    if (lane_index == 0) {
      s_matrix_component[warp_index] = matrix_component;
    }

    __syncthreads();

    if (warp_index == 0) {
      matrix_component = s_matrix_component[lane_index];
#ifdef __HIP_PLATFORM_AMD__
      if (warp_count >= 32) {
        matrix_component += __shfl_down(matrix_component, 16);
      }
      if (warp_count >= 16) {
        matrix_component += __shfl_down(matrix_component, 8);

Review comment (suggestion):
To reduce the amount of code, and remove near-duplication, perhaps defining macros for unsupported functions would make more sense, eg.

```cpp
#ifdef __HIP_PLATFORM_AMD__
  #define __SHFL_DOWN __shfl_down
  #define __SHFL_XOR __shfl_xor
#else
  #define __SHFL_DOWN __shfl_down_sync
  #define __SHFL_XOR __shfl_xor_sync
#endif
```
Wherever `__shfl_down_sync` or `__shfl_xor_sync` appears replace with this macro, the same can be done for the other unsupported functions. If support for these comes about in the future we can remove the macro.

Code after:

  for (int load = 0; load < load_count; load++) {
    global_index += load * block_size;

    if (global_index < element_count) {
      int my_alpha = g_batch_alpha[global_index];

      if (my_alpha != -1) {
        if (gaussian_index == (my_alpha & 15) + (my_alpha >> 4) * MIXTURE_COUNT) {
          float feature[CHANNEL_COUNT + 1];

          feature[0] = 1;

          for (int i = 0; i < CHANNEL_COUNT; i++) {
            feature[i + 1] = g_batch_image[global_index + i * element_count];
          }

          for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) {
            for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) {
              matrix[index] += feature[i] * feature[j];
            }
          }
        }
      }
    }
  }

  __syncthreads();

  for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) {
    float matrix_component = matrix[i];
    matrix_component += __SHFL_DOWN(matrix_component, 16);
    matrix_component += __SHFL_DOWN(matrix_component, 8);
    matrix_component += __SHFL_DOWN(matrix_component, 4);
    matrix_component += __SHFL_DOWN(matrix_component, 2);
    matrix_component += __SHFL_DOWN(matrix_component, 1);
    if (lane_index == 0) {
      s_matrix_component[warp_index] = matrix_component;
    }

    __syncthreads();

    if (warp_index == 0) {
      matrix_component = s_matrix_component[lane_index];
      if (warp_count >= 32) {
        matrix_component += __SHFL_DOWN(matrix_component, 16);
      }
      if (warp_count >= 16) {
        matrix_component += __SHFL_DOWN(matrix_component, 8);
      }
      if (warp_count >= 8) {