message: Adding ROCm support for the softmax op
diff: @@ -14,7 +14,7 @@ limitations under the License.
  ==============================================================================*/
  #include "tensorflow/core/lib/strings/str_util.h"
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  #define EIGEN_USE_GPU
  @@ -28,6 +28,12 @@ limitations under the License.
  #include "tensorflow/core/platform/types.h"
  #include "tensorflow/core/util/gpu_kernel_helper.h"
  +#if GOOGLE_CUDA
  +namespace gpuprim = ::cub;
  +#elif TENSORFLOW_USE_ROCM
  +namespace gpuprim = ::hipcub;
  +#endif
  +
  namespace tensorflow {
  namespace {
  @@ -142,7 +148,7 @@ class SoftmaxOpGPU : public OpKernel {
  OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
  {0}, 0, logits_in_.shape(), &softmax_out));
  - const cudaStream_t& cu_stream = GetCudaStream(context);
  + const auto& cu_stream = GetGpuStream(context);
  if (logits_in_.NumElements() > 0) {
  Tensor max_logits;
  Tensor sum_probs;
  @@ -155,18 +161,18 @@ class SoftmaxOpGPU : public OpKernel {
  context->allocate_temp(DataTypeToEnum<acc_type>::value,
  softmax_out->shape(), &sum_probs));
  - DoRowReduction<T, cub::Max, const T*>(
  + DoRowReduction<T, gpuprim::Max, const T*>(
  context, const_cast<T*>(max_logits.flat<T>().data()),
  reinterpret_cast<const T*>(logits_in_.flat<T>().data()), rows, cols);
  const int numThreads = 128;
  const int numBlocks = Eigen::divup(rows * cols, numThreads);
  - cub::CountingInputIterator<int> counting_iterator(0);
  - typedef cub::TransformInputIterator<acc_type,
  + gpuprim::CountingInputIterator<int> counting_iterator(0);
  + using InputIterType =
  + gpuprim::TransformInputIterator<acc_type,
  SubtractAndExpFunctor<T, acc_type>,
  - cub::CountingInputIterator<int>>
  - InputIterType;
  + gpuprim::CountingInputIterator<int>>;
  InputIterType input_itr(
  counting_iterator,
  @@ -174,11 +180,11 @@ class SoftmaxOpGPU : public OpKernel {
  reinterpret_cast<const T*>(logits_in_.flat<T>().data()),
  reinterpret_cast<const T*>(max_logits.flat<T>().data()), cols));
  - DoRowReduction<acc_type, cub::Sum, InputIterType>(
  + DoRowReduction<acc_type, gpuprim::Sum, InputIterType>(
  context, const_cast<acc_type*>(sum_probs.flat<acc_type>().data()),
  input_itr, rows, cols);
  - TF_CHECK_OK(CudaLaunchKernel(
  + TF_CHECK_OK(GpuLaunchKernel(
  GenerateNormalizedProb<T, acc_type>, numBlocks, numThreads, 0,
  cu_stream, reinterpret_cast<const T*>(logits_in_.flat<T>().data()),
  reinterpret_cast<const acc_type*>(sum_probs.flat<acc_type>().data()),
  @@ -209,4 +215,4 @@ REGISTER_KERNEL_BUILDER(
  } // end namespace tensorflow
  -#endif // GOOGLE_CUDA
  +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  