TensorReductionGpu.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0
 
#ifndef EIGEN_TENSOR_TENSOR_REDUCTION_GPU_H
#define EIGEN_TENSOR_TENSOR_REDUCTION_GPU_H
 
// IWYU pragma: private
#include "./InternalHeaderCheck.h"
 
namespace Eigen {
namespace internal {
 
#if defined(EIGEN_USE_GPU) && defined(EIGEN_GPUCC)
// Full reducers for GPU, don't vectorize for now
 
// Reducer function that enables multiple gpu threads to safely accumulate at the same
// output address. It basically reads the current value of the output variable, and
// attempts to update it with the new value. If in the meantime another gpu thread
// updated the content of the output address it will try again.
template <typename T, typename R>
__device__ EIGEN_ALWAYS_INLINE void atomicReduce(T* output, T accum, R& reducer) {
  EIGEN_IF_CONSTEXPR (sizeof(T) == 4) {
    unsigned int oldval = *reinterpret_cast<unsigned int*>(output);
    unsigned int newval = oldval;
    reducer.reduce(accum, reinterpret_cast<T*>(&newval));
    if (newval == oldval) {
      return;
    }
    unsigned int readback;
    while ((readback = atomicCAS((unsigned int*)output, oldval, newval)) != oldval) {
      oldval = readback;
      newval = oldval;
      reducer.reduce(accum, reinterpret_cast<T*>(&newval));
      if (newval == oldval) {
        return;
      }
    }
  } else EIGEN_IF_CONSTEXPR (sizeof(T) == 8) {
    unsigned long long oldval = *reinterpret_cast<unsigned long long*>(output);
    unsigned long long newval = oldval;
    reducer.reduce(accum, reinterpret_cast<T*>(&newval));
    if (newval == oldval) {
      return;
    }
    unsigned long long readback;
    while ((readback = atomicCAS(reinterpret_cast<unsigned long long*>(output), oldval, newval)) != oldval) {
      oldval = readback;
      newval = oldval;
      reducer.reduce(accum, reinterpret_cast<T*>(&newval));
      if (newval == oldval) {
        return;
      }
    }
  } else {
    gpu_assert(0 && "Wordsize not supported");
  }
}
 
// We extend atomicExch to support extra data types
template <typename Type>
__device__ inline Type atomicExchCustom(Type* address, Type val) {
  return atomicExch(address, val);
}
 
template <typename T>
EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR auto reduction_shuffle_mask() {
#if defined(EIGEN_HIP_DEVICE_COMPILE)
  return 0xFFFFFFFFFFFFFFFFull;
#else
  return 0xFFFFFFFFu;
#endif
}
 
template <typename T>
__device__ EIGEN_ALWAYS_INLINE T reduction_shuffle_down(T value, int offset) {
#if defined(EIGEN_HIPCC)
  return __shfl_down(value, offset, warpSize);
#else
  return __shfl_down_sync(reduction_shuffle_mask<T>(), value, offset, warpSize);
#endif
}
 
template <>
__device__ EIGEN_ALWAYS_INLINE int reduction_shuffle_down<int>(int value, int offset) {
#if defined(EIGEN_HIPCC)
  return __shfl_down(value, offset, warpSize);
#else
  return __shfl_down_sync(reduction_shuffle_mask<int>(), value, offset, warpSize);
#endif
}
 
template <>
__device__ EIGEN_ALWAYS_INLINE float reduction_shuffle_down<float>(float value, int offset) {
#if defined(EIGEN_HIPCC)
  return __shfl_down(value, offset, warpSize);
#else
  return __shfl_down_sync(reduction_shuffle_mask<float>(), value, offset, warpSize);
#endif
}
 
template <>
__device__ EIGEN_ALWAYS_INLINE double reduction_shuffle_down<double>(double value, int offset) {
#if defined(EIGEN_HIPCC)
  return __shfl_down(value, offset, warpSize);
#else
  return __shfl_down_sync(reduction_shuffle_mask<double>(), value, offset, warpSize);
#endif
}
 
template <>
__device__ inline double atomicExchCustom(double* address, double val) {
  unsigned long long int* address_as_ull = reinterpret_cast<unsigned long long int*>(address);
  return __longlong_as_double(atomicExch(address_as_ull, __double_as_longlong(val)));
}
 
// Half-float reduction specializations.
template <typename R>
__device__ inline void atomicReduce(half2* output, half2 accum, R& reducer) {
  unsigned int oldval = *reinterpret_cast<unsigned int*>(output);
  unsigned int newval = oldval;
  reducer.reducePacket(accum, reinterpret_cast<half2*>(&newval));
  if (newval == oldval) {
    return;
  }
  unsigned int readback;
  while ((readback = atomicCAS((unsigned int*)output, oldval, newval)) != oldval) {
    oldval = readback;
    newval = oldval;
    reducer.reducePacket(accum, reinterpret_cast<half2*>(&newval));
    if (newval == oldval) {
      return;
    }
  }
}
#ifdef EIGEN_GPU_COMPILE_PHASE
// reduction should be associative since reduction is not atomic in wide vector but atomic in half2 operations
template <typename R>
__device__ inline void atomicReduce(Packet4h2* output, Packet4h2 accum, R& reducer) {
  half2* houtput = reinterpret_cast<half2*>(output);
  half2* haccum = reinterpret_cast<half2*>(&accum);
  for (int i = 0; i < 4; ++i) {
    atomicReduce(houtput + i, *(haccum + i), reducer);
  }
}
#endif  // EIGEN_GPU_COMPILE_PHASE
 
template <>
__device__ inline void atomicReduce(float* output, float accum, SumReducer<float>&) {
  atomicAdd(output, accum);
}
 
template <>
__device__ inline void atomicReduce(double* output, double accum, SumReducer<double>&) {
  atomicAdd(output, accum);
}
 
template <typename CoeffType, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionInitKernel(const CoeffType val, Index num_preserved_coeffs,
                                                                 CoeffType* output) {
  const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
  const Index num_threads = blockDim.x * gridDim.x;
  for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) {
    output[i] = val;
  }
}
 
template <int BlockSize, int NumPerThread, typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void FullReductionKernel(Reducer reducer, const Self input, Index num_coeffs,
                                                                 typename Self::CoeffReturnType* output,
                                                                 unsigned int* semaphore) {
  // Initialize the output value
  const Index first_index = blockIdx.x * BlockSize * NumPerThread + threadIdx.x;
  if (gridDim.x == 1) {
    if (first_index == 0) {
      *output = reducer.initialize();
    }
  } else {
    if (threadIdx.x == 0) {
      unsigned int block = atomicCAS(semaphore, 0u, 1u);
      if (block == 0) {
        // We're the first block to run, initialize the output value
        atomicExchCustom(output, reducer.initialize());
        __threadfence();
        atomicExch(semaphore, 2u);
      } else {
        // Wait for the first block to initialize the output value.
        // Use atomicCAS here to ensure that the reads aren't cached
        unsigned int val;
        do {
          val = atomicCAS(semaphore, 2u, 2u);
        } while (val < 2u);
      }
    }
  }
 
  __syncthreads();
 
  eigen_assert(gridDim.x == 1 || *semaphore >= 2u);
 
  typename Self::CoeffReturnType accum = reducer.initialize();
  Index max_iter = numext::mini<Index>(num_coeffs - first_index, NumPerThread * BlockSize);
  for (Index i = 0; i < max_iter; i += BlockSize) {
    const Index index = first_index + i;
    eigen_assert(index < num_coeffs);
    typename Self::CoeffReturnType val = input.m_impl.coeff(index);
    reducer.reduce(val, &accum);
  }
 
#pragma unroll
  for (int offset = warpSize / 2; offset > 0; offset /= 2) {
    reducer.reduce(reduction_shuffle_down(accum, offset), &accum);
  }
 
  if ((threadIdx.x & (warpSize - 1)) == 0) {
    atomicReduce(output, accum, reducer);
  }
 
  if (gridDim.x > 1 && threadIdx.x == 0) {
    // Let the last block reset the semaphore
    atomicInc(semaphore, gridDim.x + 1);
#if defined(EIGEN_HIPCC)
    __threadfence_system();
#endif
  }
}
 
// Half-float reduction specializations.
template <typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionInitFullReduxKernelHalfFloat(Reducer reducer, const Self input,
                                                                                   Index num_coeffs, half* scratch) {
  eigen_assert(blockDim.x == 1);
  eigen_assert(gridDim.x == 1);
  typedef packet_traits<Eigen::half>::type packet_type;
  Index packet_remainder = num_coeffs % Index(unpacket_traits<packet_type>::size);
  if (packet_remainder != 0) {
    half2* h2scratch = reinterpret_cast<half2*>(scratch);
    for (Index i = num_coeffs - packet_remainder; i + 2 <= num_coeffs; i += 2) {
      *h2scratch = __halves2half2(input.coeff(i), input.coeff(i + 1));
      h2scratch++;
    }
    if ((num_coeffs & 1) != 0) {
      half lastCoeff = input.coeff(num_coeffs - 1);
      *h2scratch = __halves2half2(lastCoeff, reducer.initialize());
    }
  } else {
    packet_type reduce = reducer.template initializePacket<packet_type>();
    internal::pstoreu(scratch, reduce);
  }
}
 
template <typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionInitKernelHalfFloat(Reducer reducer, const Self /*input*/,
                                                                          Index num_coeffs, half* output) {
  const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
  const Index num_threads = blockDim.x * gridDim.x;
  typedef typename packet_traits<Eigen::half>::type PacketType;
 
  const Index num_packets = num_coeffs / Index(unpacket_traits<PacketType>::size);
  PacketType* p_output = reinterpret_cast<PacketType*>(output);
  for (Index i = thread_id; i < num_packets; i += num_threads) {
    p_output[i] = reducer.template initializePacket<PacketType>();
  }
  Index packet_remainder = num_coeffs % Index(unpacket_traits<PacketType>::size);
  if (thread_id < packet_remainder) {
    output[num_coeffs - packet_remainder + thread_id] = reducer.initialize();
  }
}
 
template <int BlockSize, int NumPerThread, typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void FullReductionKernelHalfFloat(Reducer reducer, const Self input,
                                                                          Index num_coeffs, half* output,
                                                                          half* scratch) {
  typedef typename packet_traits<Eigen::half>::type PacketType;
  const int packet_width = unpacket_traits<PacketType>::size;
  eigen_assert(NumPerThread % packet_width == 0);
  const Index first_index = blockIdx.x * BlockSize * NumPerThread + packet_width * threadIdx.x;
 
  // Initialize the output value if it wasn't initialized by the ReductionInitKernel
 
  if (gridDim.x == 1) {
    if (first_index == 0) {
      int rem = num_coeffs % packet_width;
      if (rem != 0) {
        half2* p_scratch = reinterpret_cast<half2*>(scratch);
        pstoreu(scratch, reducer.template initializePacket<PacketType>());
        for (int i = 0; i < rem / 2; i++) {
          *p_scratch = __halves2half2(input.coeff(num_coeffs - packet_width + 2 * i),
                                      input.coeff(num_coeffs - packet_width + 2 * i + 1));
          p_scratch++;
        }
        if ((num_coeffs & 1) != 0) {
          half last = input.coeff(num_coeffs - 1);
          *p_scratch = __halves2half2(last, reducer.initialize());
        }
      } else {
        PacketType reduce = reducer.template initializePacket<PacketType>();
        pstoreu(scratch, reduce);
      }
    }
    __syncthreads();
  }
 
  PacketType accum = reducer.template initializePacket<PacketType>();
  const Index max_iter =
      numext::mini<Index>((num_coeffs - first_index) / packet_width, NumPerThread * BlockSize / packet_width);
  for (Index i = 0; i < max_iter; i += BlockSize) {
    const Index index = first_index + packet_width * i;
    eigen_assert(index + packet_width < num_coeffs);
    PacketType val = input.template packet<Unaligned>(index);
    reducer.reducePacket(val, &accum);
  }
 
#pragma unroll
  for (int offset = warpSize / 2; offset > 0; offset /= 2) {
#if defined(EIGEN_HIPCC)
    PacketType r1;
    half2* hr = reinterpret_cast<half2*>(&r1);
    half2* hacc = reinterpret_cast<half2*>(&accum);
    for (int i = 0; i < packet_width / 2; i++) {
      // FIXME : remove this workaround once we have native half/half2 support for __shfl_down
      union {
        int i;
        half2 h;
      } wka_in, wka_out;
      wka_in.h = hacc[i];
      wka_out.i = __shfl_down(wka_in.i, offset, warpSize);
      hr[i] = wka_out.h;
    }
    reducer.reducePacket(r1, &accum);
#else
    PacketType r1;
    half2* hr = reinterpret_cast<half2*>(&r1);
    half2* hacc = reinterpret_cast<half2*>(&accum);
    for (int i = 0; i < packet_width / 2; i++) {
      hr[i] = __shfl_down_sync(0xFFFFFFFF, hacc[i], (unsigned)offset, warpSize);
    }
    reducer.reducePacket(r1, &accum);
 
#endif
  }
 
  if ((threadIdx.x & (warpSize - 1)) == 0) {
    atomicReduce(reinterpret_cast<PacketType*>(scratch), accum, reducer);
  }
 
  __syncthreads();
  half2* rv1 = reinterpret_cast<half2*>(scratch);
  if (packet_width > 2) {
    reducer.reducePacket(rv1[2], rv1);
    reducer.reducePacket(rv1[3], rv1 + 1);
    reducer.reducePacket(rv1[1], rv1);
  }
  if (gridDim.x == 1) {
    if (first_index == 0) {
      half tmp = __low2half(*rv1);
      reducer.reduce(__high2half(*rv1), &tmp);
      *output = tmp;
    }
  }
}
 
template <typename Op>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionCleanupKernelHalfFloat(Op reducer, half* output, half* scratch) {
  eigen_assert(threadIdx.x == 1);
  typedef packet_traits<Eigen::half>::type packet_type;
  EIGEN_IF_CONSTEXPR (unpacket_traits<packet_type>::size == 1) {
    *output = *scratch;
  } else {
    half2* pscratch = reinterpret_cast<half2*>(scratch);
    half tmp = __float2half(0.f);
    for (int i = 0; i < unpacket_traits<packet_type>::size; i += 2) {
      reducer.reduce(__low2half(*pscratch), &tmp);
      reducer.reduce(__high2half(*pscratch), &tmp);
      pscratch++;
    }
    *output = tmp;
  }
}
 
template <typename Self, typename Op, typename OutputType, bool PacketAccess, typename Enabled = void>
struct FullReductionLauncher {
  static void run(const Self&, Op&, const GpuDevice&, OutputType*, typename Self::Index) {
    gpu_assert(false && "Should only be called on doubles, floats and half floats");
  }
};
 
// Specialization for float and double
template <typename Self, typename Op, typename OutputType, bool PacketAccess>
struct FullReductionLauncher<
    Self, Op, OutputType, PacketAccess,
    std::enable_if_t<std::is_same<float, OutputType>::value || std::is_same<double, OutputType>::value, void>> {
  static void run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output,
                  typename Self::Index num_coeffs) {
    typedef typename Self::Index Index;
    const int block_size = 256;
    const int num_per_thread = 128;
    const int num_blocks = numext::div_ceil<int>(num_coeffs, block_size * num_per_thread);
 
    unsigned int* semaphore = nullptr;
    if (num_blocks > 1) {
      semaphore = device.semaphore();
    }
 
    LAUNCH_GPU_KERNEL((FullReductionKernel<block_size, num_per_thread, Self, Op, Index>), num_blocks, block_size, 0,
                      device, reducer, self, num_coeffs, output, semaphore);
  }
};
 
// Half-float reduction specializations.
template <typename Self, typename Op>
struct FullReductionLauncher<Self, Op, Eigen::half, false> {
  static void run(const Self&, Op&, const GpuDevice&, half*, typename Self::Index) {
    gpu_assert(false && "Should not be called since there is no packet accessor");
  }
};
 
template <typename Self, typename Op>
struct FullReductionLauncher<Self, Op, Eigen::half, true> {
  static void run(const Self& self, Op& reducer, const GpuDevice& device, half* output,
                  typename Self::Index num_coeffs) {
    typedef typename Self::Index Index;
 
    const int block_size = 256;
    const int num_per_thread = 128;
    const int num_blocks = numext::div_ceil<int>(num_coeffs, block_size * num_per_thread);
    half* scratch = static_cast<half*>(device.scratchpad());
 
    if (num_blocks > 1) {
      // We initialize the output and the scratchpad outside the reduction kernel when we can't be sure that there
      // won't be race conditions between multiple thread blocks.
      LAUNCH_GPU_KERNEL((ReductionInitFullReduxKernelHalfFloat<Self, Op, Index>), 1, 1, 0, device, reducer, self,
                        num_coeffs, scratch);
    }
 
    LAUNCH_GPU_KERNEL((FullReductionKernelHalfFloat<block_size, num_per_thread, Self, Op, Index>), num_blocks,
                      block_size, 0, device, reducer, self, num_coeffs, output, scratch);
 
    if (num_blocks > 1) {
      LAUNCH_GPU_KERNEL((ReductionCleanupKernelHalfFloat<Op>), 1, 1, 0, device, reducer, output, scratch);
    }
  }
};
 
template <typename Self, typename Op, bool Vectorizable>
struct FullReducer<Self, Op, GpuDevice, Vectorizable> {
  // Unfortunately nvidia doesn't support well exotic types such as complex,
  // so reduce the scope of the optimized version of the code to the simple cases
  // of doubles, floats and half floats
  // Half-float reduction specializations.
  static constexpr bool HasOptimizedImplementation =
      !Self::ReducerTraits::IsStateful && (std::is_same<typename Self::CoeffReturnType, float>::value ||
                                           std::is_same<typename Self::CoeffReturnType, double>::value ||
                                           (std::is_same<typename Self::CoeffReturnType, Eigen::half>::value &&
                                            reducer_traits<Op, GpuDevice>::PacketAccess));
 
  template <typename OutputType>
  static void run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output) {
    gpu_assert(HasOptimizedImplementation && "Should only be called on doubles, floats or half floats");
    const Index num_coeffs = array_prod(self.m_impl.dimensions());
    // Don't crash when we're called with an input tensor of size 0.
    if (num_coeffs == 0) {
      return;
    }
 
    FullReductionLauncher<Self, Op, OutputType, reducer_traits<Op, GpuDevice>::PacketAccess>::run(self, reducer, device,
                                                                                                  output, num_coeffs);
  }
};
 
template <int NumPerThread, typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void InnerReductionKernel(Reducer reducer, const Self input,
                                                                  Index num_coeffs_to_reduce,
                                                                  Index num_preserved_coeffs,
                                                                  typename Self::CoeffReturnType* output) {
  typedef typename Self::CoeffReturnType Type;
  eigen_assert(blockDim.y == 1);
  eigen_assert(blockDim.z == 1);
  eigen_assert(gridDim.y == 1);
  eigen_assert(gridDim.z == 1);
 
  const int unroll_times = 16;
  eigen_assert(NumPerThread % unroll_times == 0);
 
  const Index input_col_blocks = numext::div_ceil<Index>(num_coeffs_to_reduce, blockDim.x * NumPerThread);
  const Index num_input_blocks = input_col_blocks * num_preserved_coeffs;
 
  const Index num_threads = blockDim.x * gridDim.x;
  const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
 
  // Initialize the output values if they weren't initialized by the ReductionInitKernel
  if (gridDim.x == 1) {
    for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) {
      output[i] = reducer.initialize();
    }
    __syncthreads();
  }
 
  for (Index i = blockIdx.x; i < num_input_blocks; i += gridDim.x) {
    const Index row = i / input_col_blocks;
 
    // `i`, and so `row`, is uniform across the block: a warp enters this branch as a whole, which the full-mask
    // shuffle tree below relies on.
    if (row < num_preserved_coeffs) {
      const Index col_block = i % input_col_blocks;
      const Index col_begin = col_block * blockDim.x * NumPerThread + threadIdx.x;
 
      Type reduced_val = reducer.initialize();
 
      for (Index j = 0; j < NumPerThread; j += unroll_times) {
        const Index last_col = col_begin + blockDim.x * (j + unroll_times - 1);
        if (last_col >= num_coeffs_to_reduce) {
          for (Index col = col_begin + blockDim.x * j; col < num_coeffs_to_reduce; col += blockDim.x) {
            const Type val = input.m_impl.coeff(row * num_coeffs_to_reduce + col);
            reducer.reduce(val, &reduced_val);
          }
          break;
        } else {
          // Faster version of the loop with no branches after unrolling.
#pragma unroll
          for (int k = 0; k < unroll_times; ++k) {
            const Index col = col_begin + blockDim.x * (j + k);
            reducer.reduce(input.m_impl.coeff(row * num_coeffs_to_reduce + col), &reduced_val);
          }
        }
      }
 
#pragma unroll
      for (int offset = warpSize / 2; offset > 0; offset /= 2) {
        reducer.reduce(reduction_shuffle_down(reduced_val, offset), &reduced_val);
      }
 
      if ((threadIdx.x & (warpSize - 1)) == 0) {
        atomicReduce(&(output[row]), reduced_val, reducer);
      }
    }
  }
}
 
// Half-float reduction specializations.
 
template <int NumPerThread, typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void InnerReductionKernelHalfFloat(Reducer reducer, const Self input,
                                                                           Index num_coeffs_to_reduce,
                                                                           Index num_preserved_coeffs, half* output) {
  eigen_assert(blockDim.y == 1);
  eigen_assert(blockDim.z == 1);
  eigen_assert(gridDim.y == 1);
  eigen_assert(gridDim.z == 1);
 
  typedef typename packet_traits<Eigen::half>::type PacketType;
  const int packet_width = unpacket_traits<PacketType>::size;
  const int unroll_times = 16 / packet_width;
  eigen_assert(NumPerThread % unroll_times == 0);
  eigen_assert(unroll_times % 2 == 0);
 
  const Index input_col_blocks = numext::div_ceil<Index>(num_coeffs_to_reduce, blockDim.x * NumPerThread * 2);
  const Index num_input_blocks = numext::div_ceil<Index>(input_col_blocks * num_preserved_coeffs, 2);
 
  const Index num_threads = blockDim.x * gridDim.x;
  const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
 
  // Initialize the output values if they weren't initialized by the ReductionInitKernel
  if (gridDim.x == 1) {
    Index i = packet_width * thread_id;
    for (; i + packet_width <= num_preserved_coeffs; i += packet_width * num_threads) {
      PacketType* poutput = reinterpret_cast<PacketType*>(output + i);
      *poutput = reducer.template initializePacket<PacketType>();
    }
    if (i < num_preserved_coeffs) {
      output[i] = reducer.initialize();
    }
    __syncthreads();
  }
 
  for (Index i = blockIdx.x; i < num_input_blocks; i += gridDim.x) {
    const Index row = 2 * (i / input_col_blocks);  // everybody takes 2 rows
 
    // Block-uniform, as in InnerReductionKernel: the full-mask shuffles below are safe.
    if (row + 1 < num_preserved_coeffs) {
      const Index col_block = i % input_col_blocks;
      const Index col_begin = packet_width * (col_block * blockDim.x * NumPerThread + threadIdx.x);
 
      PacketType reduced_val1 = reducer.template initializePacket<PacketType>();
      PacketType reduced_val2 = reducer.template initializePacket<PacketType>();
 
      for (Index j = 0; j < NumPerThread; j += unroll_times) {
        const Index last_col = col_begin + blockDim.x * (j + unroll_times - 1) * packet_width;
        if (last_col >= num_coeffs_to_reduce) {
          Index col = col_begin + blockDim.x * j;
          for (; col + packet_width <= num_coeffs_to_reduce; col += blockDim.x) {
            const PacketType val1 = input.m_impl.template packet<Unaligned>(row * num_coeffs_to_reduce + col);
            reducer.reducePacket(val1, &reduced_val1);
            const PacketType val2 = input.m_impl.template packet<Unaligned>((row + 1) * num_coeffs_to_reduce + col);
            reducer.reducePacket(val2, &reduced_val2);
          }
          if (col < num_coeffs_to_reduce) {
            PacketType r1 = reducer.template initializePacket<PacketType>();
            PacketType r2 = reducer.template initializePacket<PacketType>();
            half2* hr1 = reinterpret_cast<half2*>(&r1);
            half2* hr2 = reinterpret_cast<half2*>(&r2);
            while (col + 1 < num_coeffs_to_reduce) {
              *hr1 = __halves2half2(input.m_impl.coeff(row * num_coeffs_to_reduce + col),
                                    input.m_impl.coeff(row * num_coeffs_to_reduce + col + 1));
              *hr2 = __halves2half2(input.m_impl.coeff((row + 1) * num_coeffs_to_reduce + col),
                                    input.m_impl.coeff((row + 1) * num_coeffs_to_reduce + col + 1));
              hr1++;
              hr2++;
              col += 2;
            }
            if (col < num_coeffs_to_reduce) {
              // Peel;
              const half last1 = input.m_impl.coeff(row * num_coeffs_to_reduce + col);
              *hr1 = __halves2half2(last1, reducer.initialize());
              const half last2 = input.m_impl.coeff((row + 1) * num_coeffs_to_reduce + col);
              *hr2 = __halves2half2(last2, reducer.initialize());
            }
            reducer.reducePacket(r1, &reduced_val1);
            reducer.reducePacket(r2, &reduced_val2);
          }
          break;
        } else {
          // Faster version of the loop with no branches after unrolling.
#pragma unroll
          for (int k = 0; k < unroll_times; ++k) {
            const Index col = col_begin + blockDim.x * (j + k) * packet_width;
            reducer.reducePacket(input.m_impl.template packet<Unaligned>(row * num_coeffs_to_reduce + col),
                                 &reduced_val1);
            reducer.reducePacket(input.m_impl.template packet<Unaligned>((row + 1) * num_coeffs_to_reduce + col),
                                 &reduced_val2);
          }
        }
      }
 
#pragma unroll
      for (int offset = warpSize / 2; offset > 0; offset /= 2) {
#if defined(EIGEN_HIPCC)
        PacketType r1;
        PacketType r2;
        half2* hr1 = reinterpret_cast<half2*>(&r1);
        half2* hr2 = reinterpret_cast<half2*>(&r2);
        half2* rv1 = reinterpret_cast<half2*>(&reduced_val1);
        half2* rv2 = reinterpret_cast<half2*>(&reduced_val2);
        for (int i = 0; i < packet_width / 2; i++) {
          // FIXME : remove this workaround once we have native half/half2 support for __shfl_down
          union {
            int i;
            half2 h;
          } wka_in1, wka_out1;
          wka_in1.h = rv1[i];
          wka_out1.i = __shfl_down(wka_in1.i, offset, warpSize);
          hr1[i] = wka_out1.h;
 
          union {
            int i;
            half2 h;
          } wka_in2, wka_out2;
          wka_in2.h = rv2[i];
          wka_out2.i = __shfl_down(wka_in2.i, offset, warpSize);
          hr2[i] = wka_out2.h;
        }
        reducer.reducePacket(r1, &reduced_val1);
        reducer.reducePacket(r2, &reduced_val2);
#else
        PacketType r1;
        PacketType r2;
        half2* hr1 = reinterpret_cast<half2*>(&r1);
        half2* hr2 = reinterpret_cast<half2*>(&r2);
        half2* rr1 = reinterpret_cast<half2*>(&reduced_val1);
        half2* rr2 = reinterpret_cast<half2*>(&reduced_val2);
        for (int j = 0; j < packet_width / 2; j++) {
          hr1[j] = __shfl_down_sync(0xFFFFFFFF, rr1[j], (unsigned)offset, warpSize);
          hr2[j] = __shfl_down_sync(0xFFFFFFFF, rr2[j], (unsigned)offset, warpSize);
        }
        reducer.reducePacket(r1, &reduced_val1);
        reducer.reducePacket(r2, &reduced_val2);
 
#endif
      }
      half2* rv1 = reinterpret_cast<half2*>(&reduced_val1);
      half2* rv2 = reinterpret_cast<half2*>(&reduced_val2);
      half2 val;
      if (packet_width > 2) {
        reducer.reducePacket(rv1[2], rv1);
        reducer.reducePacket(rv1[3], rv1 + 1);
        reducer.reducePacket(rv1[1], rv1);
        reducer.reducePacket(rv2[2], rv2);
        reducer.reducePacket(rv2[3], rv2 + 1);
        reducer.reducePacket(rv2[1], rv2);
      }
      half val1 = __low2half(*rv1);
      reducer.reduce(__high2half(*rv1), &val1);
      half val2 = __low2half(*rv2);
      reducer.reduce(__high2half(*rv2), &val2);
      val = __halves2half2(val1, val2);
      if ((threadIdx.x & (warpSize - 1)) == 0) {
        half* loc = output + row;
        atomicReduce(reinterpret_cast<half2*>(loc), val, reducer);
      }
    }
  }
}
 
template <typename Self, typename Op, typename OutputType, bool PacketAccess, typename Enabled = void>
struct InnerReductionLauncher {
  static EIGEN_DEVICE_FUNC bool run(const Self&, Op&, const GpuDevice&, OutputType*, typename Self::Index,
                                    typename Self::Index) {
    gpu_assert(false && "Should only be called to reduce doubles, floats and half floats on a gpu device");
    return true;
  }
};
 
// Specialization for float and double
template <typename Self, typename Op, typename OutputType, bool PacketAccess>
struct InnerReductionLauncher<
    Self, Op, OutputType, PacketAccess,
    std::enable_if_t<std::is_same<float, OutputType>::value || std::is_same<double, OutputType>::value, void>> {
  static bool run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output,
                  typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) {
    typedef typename Self::Index Index;
 
    const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals;
    const int block_size = 256;
    const int num_per_thread = 128;
    const int dyn_blocks = numext::div_ceil<int>(num_coeffs, block_size * num_per_thread);
    const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / block_size;
    const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);
 
    if (num_blocks > 1) {
      // We initialize the outputs outside the reduction kernel when we can't be sure that there
      // won't be race conditions between multiple thread blocks.
      const int dyn_blocks2 = numext::div_ceil<int>(num_preserved_vals, 1024);
      const int max_blocks2 = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / 1024;
      const int num_blocks2 = numext::mini<int>(max_blocks2, dyn_blocks2);
      LAUNCH_GPU_KERNEL((ReductionInitKernel<OutputType, Index>), num_blocks2, 1024, 0, device, reducer.initialize(),
                        num_preserved_vals, output);
    }
 
    LAUNCH_GPU_KERNEL((InnerReductionKernel<num_per_thread, Self, Op, Index>), num_blocks, block_size, 0, device,
                      reducer, self, num_coeffs_to_reduce, num_preserved_vals, output);
 
    return false;
  }
};
 
// Half-float reduction specializations.
template <typename Self, typename Op>
struct InnerReductionLauncher<Self, Op, Eigen::half, false> {
  static bool run(const Self&, Op&, const GpuDevice&, half*, typename Self::Index, typename Self::Index) {
    gpu_assert(false && "Should not be called since there is no packet accessor");
    return true;
  }
};
 
template <typename Self, typename Op>
struct InnerReductionLauncher<Self, Op, Eigen::half, true> {
  static bool run(const Self& self, Op& reducer, const GpuDevice& device, half* output,
                  typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) {
    typedef typename Self::Index Index;
 
    if (num_preserved_vals % 2 != 0) {
      // Not supported yet, revert to the slower code path
      return true;
    }
 
    const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals;
    const int block_size = 128;
    const int num_per_thread = 64;
    const int dyn_blocks = numext::div_ceil<int>(num_coeffs, block_size * num_per_thread);
    const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / block_size;
    const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);
 
    if (num_blocks > 1) {
      // We initialize the outputs outside the reduction kernel when we can't be sure that there
      // won't be race conditions between multiple thread blocks.
      LAUNCH_GPU_KERNEL((ReductionInitKernelHalfFloat<Self, Op, Index>), 1, 1, 0, device, reducer, self,
                        num_preserved_vals, output);
    }
 
    LAUNCH_GPU_KERNEL((InnerReductionKernelHalfFloat<num_per_thread, Self, Op, Index>), num_blocks, block_size, 0,
                      device, reducer, self, num_coeffs_to_reduce, num_preserved_vals, output);
 
    return false;
  }
};
 
template <typename Self, typename Op>
struct InnerReducer<Self, Op, GpuDevice> {
  // Unfortunately nvidia doesn't support well exotic types such as complex,
  // so reduce the scope of the optimized version of the code to the simple case
  // of floats and half floats.
  // Half-float reduction specializations.
  static constexpr bool HasOptimizedImplementation =
      !Self::ReducerTraits::IsStateful && (std::is_same<typename Self::CoeffReturnType, float>::value ||
                                           std::is_same<typename Self::CoeffReturnType, double>::value ||
                                           (std::is_same<typename Self::CoeffReturnType, Eigen::half>::value &&
                                            reducer_traits<Op, GpuDevice>::PacketAccess));
 
  template <typename OutputType>
  static bool run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output,
                  typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) {
    gpu_assert(HasOptimizedImplementation && "Should only be called on doubles, floats or half floats");
    const Index num_coeffs = array_prod(self.m_impl.dimensions());
    // Don't crash when we're called with an input tensor of size 0.
    if (num_coeffs == 0) {
      return true;
    }
    // It's faster to use the usual code.
    if (num_coeffs_to_reduce <= 128) {
      return true;
    }
 
    return InnerReductionLauncher<Self, Op, OutputType, reducer_traits<Op, GpuDevice>::PacketAccess>::run(
        self, reducer, device, output, num_coeffs_to_reduce, num_preserved_vals);
  }
};
 
template <int NumPerThread, typename Self, typename Reducer, typename Index>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void OuterReductionKernel(Reducer reducer, const Self input,
                                                                  Index num_coeffs_to_reduce,
                                                                  Index num_preserved_coeffs,
                                                                  typename Self::CoeffReturnType* output) {
  const Index num_threads = blockDim.x * gridDim.x;
  const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
  // Initialize the output values if they weren't initialized by the ReductionInitKernel
  if (gridDim.x == 1) {
    for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) {
      output[i] = reducer.initialize();
    }
    __syncthreads();
  }
 
  // Do the reduction.
  const Index max_iter = num_preserved_coeffs * numext::div_ceil<Index>(num_coeffs_to_reduce, NumPerThread);
  for (Index i = thread_id; i < max_iter; i += num_threads) {
    const Index input_col = i % num_preserved_coeffs;
    const Index input_row = (i / num_preserved_coeffs) * NumPerThread;
    typename Self::CoeffReturnType reduced_val = reducer.initialize();
    const Index max_row = numext::mini(input_row + NumPerThread, num_coeffs_to_reduce);
    for (Index j = input_row; j < max_row; j++) {
      typename Self::CoeffReturnType val = input.m_impl.coeff(j * num_preserved_coeffs + input_col);
      reducer.reduce(val, &reduced_val);
    }
    atomicReduce(&(output[input_col]), reduced_val, reducer);
  }
}
 
template <typename Self, typename Op>
struct OuterReducer<Self, Op, GpuDevice> {
  // Unfortunately nvidia doesn't support well exotic types such as complex,
  // so reduce the scope of the optimized version of the code to the simple case
  // of floats and doubles.
  template <typename T>
  using IsFloatOrDouble = bool_constant<std::is_same<T, float>::value || std::is_same<T, double>::value>;
 
  static constexpr bool HasOptimizedImplementation =
      !Self::ReducerTraits::IsStateful && IsFloatOrDouble<typename Self::CoeffReturnType>::value;
  template <typename Device, typename OutputType>
  static
#if !defined(EIGEN_HIPCC)
      // FIXME :  leaving this EIGEN_DEVICE_FUNC in, results in the following runtime error
      //          (in the tensor_reduction_gpu test)
      //
      // terminate called after throwing an instance of 'std::runtime_error'
      //   what():  No device code available for function: _ZN5Eigen8internal20OuterReductionKernelIL...
      //
      // don't know why this happens (and why is it a runtime error instead of a compile time error)
      //
      // this will be fixed by HIP PR#457
      EIGEN_DEVICE_FUNC
#endif
      bool
      run(const Self&, Op&, const Device&, OutputType*, typename Self::Index, typename Self::Index) {
    gpu_assert(false && "Should only be called to reduce doubles or floats on a gpu device");
    return true;
  }
 
  // OuterReductionKernel is generic in CoeffReturnType, so one launch serves both types that
  // HasOptimizedImplementation admits. The enable_if keeps every other OutputType on the overload above: the
  // caller's EIGEN_IF_CONSTEXPR is a plain `if` in C++14, so the call is instantiated for those types as well.
  template <typename OutputType, EIGEN_SFINAE_ENABLE_IF(IsFloatOrDouble<OutputType>::value)>
  static bool run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output,
                  typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) {
    typedef typename Self::Index Index;
 
    // It's faster to use the usual code.
    if (num_coeffs_to_reduce <= 32) {
      return true;
    }
 
    // CUDA has native FP64 atomicAdd on every supported architecture (sm_60+). Keep the measured shape band
    // for the other double reducers and for HIP, where atomicAdd may still use compare-and-swap.
#if defined(EIGEN_CUDACC)
    constexpr bool has_native_double_sum = std::is_same<Op, SumReducer<double>>::value;
#else
    constexpr bool has_native_double_sum = false;
#endif
    if (std::is_same<OutputType, double>::value && !has_native_double_sum) {
      const Index multi_processors = device.getNumGpuMultiProcessors();
      if (num_coeffs_to_reduce < 64 || num_preserved_vals < 8 * multi_processors ||
          num_preserved_vals > 64 * multi_processors) {
        return true;
      }
    }
 
    const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals;
    const int block_size = 256;
    const int num_per_thread = 16;
    const int dyn_blocks = numext::div_ceil<int>(num_coeffs, block_size * num_per_thread);
    const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / block_size;
    const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);
 
    if (num_blocks > 1) {
      // We initialize the outputs outside the reduction kernel when we can't be sure that there
      // won't be race conditions between multiple thread blocks.
      const int dyn_blocks2 = numext::div_ceil<int>(num_preserved_vals, 1024);
      const int max_blocks2 = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / 1024;
      const int num_blocks2 = numext::mini<int>(max_blocks2, dyn_blocks2);
      LAUNCH_GPU_KERNEL((ReductionInitKernel<OutputType, Index>), num_blocks2, 1024, 0, device, reducer.initialize(),
                        num_preserved_vals, output);
    }
 
    LAUNCH_GPU_KERNEL((OuterReductionKernel<num_per_thread, Self, Op, Index>), num_blocks, block_size, 0, device,
                      reducer, self, num_coeffs_to_reduce, num_preserved_vals, output);
 
    return false;
  }
};
 
#endif  // defined(EIGEN_USE_GPU) && defined(EIGEN_GPUCC)
 
}  // end namespace internal
}  // end namespace Eigen
 
#endif  // EIGEN_TENSOR_TENSOR_REDUCTION_GPU_H