experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Added ROCm support for the pad op

Browse Source
This commit is contained in:
Jeffrey Poznanovic 2019-03-19 02:08:24 +00:00
parent 6e9cb400d1
commit 42621d70df
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/kernels/pad_op.cc
View File

@ -294,7 +294,7 @@ TF_CALL_POD_TYPES(REGISTER_KERNEL);
TF_CALL_string(REGISTER_KERNEL); TF_CALL_string(REGISTER_KERNEL);
#undef REGISTER_KERNEL #undef REGISTER_KERNEL
#if GOOGLE_CUDA #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Forward declarations of the functor specializations for GPU. // Forward declarations of the functor specializations for GPU.
namespace functor { namespace functor {
#define DECLARE_GPU_SPEC(T, Dims) \ #define DECLARE_GPU_SPEC(T, Dims) \
@ -395,7 +395,7 @@ REGISTER_KERNEL_BUILDER(Name("PadV2")
.HostMemory("constant_values") .HostMemory("constant_values")
.HostMemory("output"), .HostMemory("output"),
PadOp<CPUDevice, int32, int64>); PadOp<CPUDevice, int32, int64>);
#endif #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#ifdef TENSORFLOW_USE_SYCL #ifdef TENSORFLOW_USE_SYCL
// Registration of the GPU implementations. // Registration of the GPU implementations.

4
tensorflow/core/kernels/pad_op_gpu.cu.cc
View File

@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
limitations under the License. limitations under the License.
==============================================================================*/ ==============================================================================*/
#if GOOGLE_CUDA #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU #define EIGEN_USE_GPU
@ -45,4 +45,4 @@ TF_CALL_uint8(DEFINE_GPU_SPECS);
} // namespace tensorflow } // namespace tensorflow
#endif // GOOGLE_CUDA #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM