[ROCm] Enabling vectorization/packet support for some cwise ops on the ROCm platform

tensorflow/core/kernels/cwise_ops.h

@@ -177,11 +177,7 @@ template <typename T>
 struct functor_traits<div_no_nan_op<T>> {
   enum {
     Cost = functor_traits<scalar_quotient_op<T>>::Cost + NumTraits<T>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = true,
-#endif  // TENSORFLOW_USE_ROCM
   };
 };
 
@@ -194,11 +190,7 @@ template <typename T>
 struct functor_traits<mul_no_nan_op<T>> {
   enum {
     Cost = functor_traits<scalar_product_op<T>>::Cost + NumTraits<T>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = true,
-#endif  // TENSORFLOW_USE_ROCM
   };
 };
 
@@ -257,11 +249,7 @@ struct functor_traits<
     scalar_left<Tout, Tin, Binary, is_scalar_in_host_memory>> {
   enum {
     Cost = functor_traits<Binary>::Cost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = functor_traits<Binary>::PacketAccess,
-#endif  // TENSORFLOW_USE_ROCM
   };
 };
 
@@ -312,11 +300,7 @@ struct functor_traits<
     scalar_right<Tout, Tin, Binary, is_scalar_in_host_memory>> {
   enum {
     Cost = functor_traits<Binary>::Cost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = functor_traits<Binary>::PacketAccess,
-#endif  // TENSORFLOW_USE_ROCM
   };
 };
 
@@ -458,11 +442,7 @@ struct functor_traits<google_floor_div<Scalar>> {
     Cost = 2 * Eigen::internal::scalar_div_cost<
                    Scalar, packet_traits<Scalar>::HasDiv>::value +
            NumTraits<Scalar>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<Scalar>::HasDiv
-#endif
   };
 };
 
@@ -485,12 +465,8 @@ struct functor_traits<google_floor_div_real<Scalar>> {
     Cost = 2 * Eigen::internal::scalar_div_cost<
                    Scalar, packet_traits<Scalar>::HasDiv>::value +
            2 * NumTraits<Scalar>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess =
         packet_traits<Scalar>::HasDiv && packet_traits<Scalar>::HasFloor
-#endif
   };
 };
 
@@ -596,13 +572,9 @@ struct functor_traits<scalar_round_half_to_even_op<Scalar>> {
   enum {
     Cost = Eigen::NumTraits<Scalar>::IsInteger ? 0
                                                : 4 * NumTraits<Scalar>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<Scalar>::HasFloor &&
                    packet_traits<Scalar>::HasAdd &&
                    packet_traits<Scalar>::HasMul,
-#endif
   };
 };
 
@@ -638,11 +610,7 @@ template <typename Scalar, bool IsInteger>
 struct functor_traits<scalar_round_up_op<Scalar, IsInteger>> {
   enum {
     Cost = IsInteger ? 0 : 4 * NumTraits<Scalar>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = IsInteger || packet_traits<Scalar>::HasFloor
-#endif
   };
 };
 
@@ -695,11 +663,7 @@ struct functor_traits<xlogy_op<Scalar>> {
   enum {
     Cost = functor_traits<scalar_log_op<Scalar>>::Cost +
            Eigen::NumTraits<Scalar>::MulCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = functor_traits<scalar_log_op<Scalar>>::PacketAccess
-#endif
   };
 };
 
@@ -730,11 +694,7 @@ struct functor_traits<xdivy_op<Scalar>> {
         Eigen::NumTraits<Scalar>::AddCost +
         Eigen::internal::scalar_div_cost<Scalar,
                                          packet_traits<Scalar>::HasDiv>::value,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<Scalar>::HasDiv
-#endif
   };
 };
 
@@ -760,11 +720,7 @@ template <typename T>
 struct functor_traits<scalar_erfinv_op<T>> {
   enum {
     Cost = functor_traits<scalar_ndtri_op<T>>::Cost + NumTraits<T>::AddCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<T>::HasNdtri,
-#endif
   };
 };
 

tensorflow/core/kernels/cwise_ops_gradients.h

@@ -94,11 +94,7 @@ template <typename T>
 struct functor_traits<scalar_inverse_gradient_op<T>> {
   enum {
     Cost = NumTraits<T>::AddCost + 2 * NumTraits<T>::MulCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<T>::HasMul,
-#endif
   };
 };
 
@@ -127,11 +123,7 @@ struct scalar_sqrt_gradient_op {
 template <typename T>
 struct functor_traits<scalar_sqrt_gradient_op<T>> {
   enum {
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<T>::HasMul & packet_traits<T>::HasDiv,
-#endif
     Cost = NumTraits<T>::MulCost + scalar_div_cost<T, PacketAccess>::value,
   };
 };
@@ -166,11 +158,7 @@ template <typename T>
 struct functor_traits<scalar_rsqrt_gradient_op<T>> {
   enum {
     Cost = 4 * NumTraits<T>::MulCost,
-#if TENSORFLOW_USE_ROCM
-    PacketAccess = false,
-#else
     PacketAccess = packet_traits<T>::HasMul,
-#endif
   };
 };