message: Adding ROCm support for the eye functor
diff: @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
  limitations under the License.
  ==============================================================================*/
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  #define EIGEN_USE_GPU
  @@ -33,7 +33,7 @@ __global__ void EyeKernel(int num_threads, int batch_size, int m, int n,
  Scalar* output_ptr) {
  const Scalar one = Scalar(1);
  const Scalar zero = Scalar(0);
  - CUDA_1D_KERNEL_LOOP(index, num_threads) {
  + GPU_1D_KERNEL_LOOP(index, num_threads) {
  // TODO(rmlarsen): Benchmark to see if it's just as fast to use mod (%),
  // since it's easier to read.
  const int global_row = index / n;
  @@ -51,8 +51,8 @@ struct EyeFunctor<GPUDevice, Scalar> {
  const int batch_size = matrix_batch.dimension(0);
  const int m = matrix_batch.dimension(1);
  const int n = matrix_batch.dimension(2);
  - GpuLaunchConfig config = GetCudaLaunchConfig(batch_size * m * n, device);
  - TF_CHECK_OK(CudaLaunchKernel(EyeKernel<Scalar>, config.block_count,
  + GpuLaunchConfig config = GetGpuLaunchConfig(batch_size * m * n, device);
  + TF_CHECK_OK(GpuLaunchKernel(EyeKernel<Scalar>, config.block_count,
  config.thread_per_block, 0, device.stream(),
  config.virtual_thread_count, batch_size, m, n,
  matrix_batch.data()));
  @@ -67,4 +67,4 @@ template struct EyeFunctor<GPUDevice, complex128>;
  } // namespace functor
  } // namespace tensorflow
  -#endif // GOOGLE_CUDA
  +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  