Merge non-build changes from #26722

PiperOrigin-RevId: 244379269

tensorflow/core/kernels/batchtospace_op.cc

@@ -266,7 +266,7 @@ class BatchToSpaceOp : public OpKernel {
 TF_CALL_REAL_NUMBER_TYPES(REGISTER);
 #undef REGISTER
 
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if GOOGLE_CUDA
 #define REGISTER(T)                                        \
   REGISTER_KERNEL_BUILDER(Name("BatchToSpaceND")           \
                               .Device(DEVICE_GPU)          \
@@ -282,6 +282,6 @@ TF_CALL_REAL_NUMBER_TYPES(REGISTER);
 
 TF_CALL_GPU_NUMBER_TYPES(REGISTER);
 #undef REGISTER
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#endif  // GOOGLE_CUDA
 
 }  // end namespace tensorflow

tensorflow/core/kernels/fill_functor.cu.cc

@@ -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
 
@@ -116,4 +116,4 @@ TF_CALL_bool(DEFINE_SETONE_GPU);
 }  // end namespace functor
 }  // end namespace tensorflow
 
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

third_party/eigen3/gpu_packet_math.patch

@@ -4,7 +4,7 @@
    return make_double2(from, from);
  }
  
-+#if defined(EIGEN_CUDA_ARCH)
++#if defined(EIGEN_CUDA_ARCH) || defined(EIGEN_HIP_DEVICE_COMPILE)
  namespace {
  
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float bitwise_and(const float& a,
@@ -12,7 +12,7 @@
  pcmp_eq<double2>(const double2& a, const double2& b) {
    return make_double2(eq_mask(a.x, b.x), eq_mask(a.y, b.y));
  }
-+#endif  // EIGEN_CUDA_ARCH
++#endif  // EIGEN_CUDA_ARCH || defined(EIGEN_HIP_DEVICE_COMPILE)
  
  template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float4 plset<float4>(const float& a) {
    return make_float4(a, a+1, a+2, a+3);