message: Adding ROCm support for the map_stage_op
diff: @@ -543,7 +543,7 @@ REGISTER_KERNEL_BUILDER(Name("MapStage").Device(DEVICE_CPU), MapStageOp<false>);
  REGISTER_KERNEL_BUILDER(Name("OrderedMapStage").Device(DEVICE_CPU),
  MapStageOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(
  Name("MapStage").HostMemory("key").HostMemory("indices").Device(DEVICE_GPU),
  MapStageOp<false>);
  @@ -552,7 +552,7 @@ REGISTER_KERNEL_BUILDER(Name("OrderedMapStage")
  .HostMemory("indices")
  .Device(DEVICE_GPU),
  MapStageOp<true>);
  -#endif // GOOGLE_CUDA
  +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  #ifdef TENSORFLOW_USE_SYCL
  REGISTER_KERNEL_BUILDER(Name("MapStage")
  @@ -603,7 +603,7 @@ REGISTER_KERNEL_BUILDER(Name("MapUnstage").Device(DEVICE_CPU),
  REGISTER_KERNEL_BUILDER(Name("OrderedMapUnstage").Device(DEVICE_CPU),
  MapUnstageOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("MapUnstage")
  .HostMemory("key")
  .HostMemory("indices")
  @@ -663,7 +663,7 @@ REGISTER_KERNEL_BUILDER(Name("MapPeek").Device(DEVICE_CPU), MapPeekOp<false>);
  REGISTER_KERNEL_BUILDER(Name("OrderedMapPeek").Device(DEVICE_CPU),
  MapPeekOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(
  Name("MapPeek").HostMemory("key").HostMemory("indices").Device(DEVICE_GPU),
  MapPeekOp<false>);
  @@ -726,7 +726,7 @@ REGISTER_KERNEL_BUILDER(Name("MapUnstageNoKey").Device(DEVICE_CPU),
  REGISTER_KERNEL_BUILDER(Name("OrderedMapUnstageNoKey").Device(DEVICE_CPU),
  MapUnstageNoKeyOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("MapUnstageNoKey")
  .HostMemory("key")
  .HostMemory("indices")
  @@ -775,7 +775,7 @@ REGISTER_KERNEL_BUILDER(Name("MapSize").Device(DEVICE_CPU), MapSizeOp<false>);
  REGISTER_KERNEL_BUILDER(Name("OrderedMapSize").Device(DEVICE_CPU),
  MapSizeOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("MapSize").Device(DEVICE_GPU).HostMemory("size"),
  MapSizeOp<false>);
  REGISTER_KERNEL_BUILDER(
  @@ -814,7 +814,7 @@ REGISTER_KERNEL_BUILDER(Name("MapIncompleteSize").Device(DEVICE_CPU),
  REGISTER_KERNEL_BUILDER(Name("OrderedMapIncompleteSize").Device(DEVICE_CPU),
  MapIncompleteSizeOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(
  Name("MapIncompleteSize").Device(DEVICE_GPU).HostMemory("size"),
  MapIncompleteSizeOp<false>);
  @@ -849,7 +849,7 @@ REGISTER_KERNEL_BUILDER(Name("MapClear").Device(DEVICE_CPU), MapClearOp<false>);
  REGISTER_KERNEL_BUILDER(Name("OrderedMapClear").Device(DEVICE_CPU),
  MapClearOp<true>);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("MapClear").Device(DEVICE_GPU), MapClearOp<false>);
  REGISTER_KERNEL_BUILDER(Name("OrderedMapClear").Device(DEVICE_GPU),
  MapClearOp<true>);
  