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

Add ROCm macro for building avgpooling_op

Browse Source
This commit is contained in:
jerryyin 2019-03-18 16:06:22 +00:00
parent 5bc5a4b562
commit 0dd6f96419
3 changed files with 10 additions and 10 deletions

12
tensorflow/core/kernels/avgpooling_op.cc
View File

@ -36,10 +36,10 @@ limitations under the License.
#include "tensorflow/core/util/padding.h"
#include "tensorflow/core/util/tensor_format.h"
#if GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "tensorflow/core/kernels/maxpooling_op_gpu.h"
#include "tensorflow/core/kernels/pooling_ops_common_gpu.h"
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
namespace tensorflow {
@ -112,7 +112,7 @@ REGISTER_KERNEL_BUILDER(
Name("AvgPool").Device(DEVICE_CPU).TypeConstraint<Eigen::half>("T"),
AvgPoolingOp<CPUDevice, Eigen::half>);
#if GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename T>
class AvgPoolingOp<GPUDevice, T> : public UnaryOp<T> {
public:
@ -205,7 +205,7 @@ REGISTER_KERNEL_BUILDER(
REGISTER_KERNEL_BUILDER(
Name("AvgPool").Device(DEVICE_GPU).TypeConstraint<double>("T"),
AvgPoolingOp<GPUDevice, double>);
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// The operation to compute AvgPool gradients.
// It takes two inputs:
@ -368,7 +368,7 @@ TF_CALL_float(REGISTER_CPU_KERNEL);
TF_CALL_double(REGISTER_CPU_KERNEL);
TF_CALL_half(REGISTER_CPU_KERNEL);
#if GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// A CUDNN based AvgPoolingGrad implementation. It includes the padding as the
// candidates for the pooling operation.
@ -577,6 +577,6 @@ REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad")
.HostMemory("orig_input_shape"),
AvgPoolingGradOpCustomGPUKernel<Eigen::half>);
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow

4
tensorflow/core/kernels/maxpooling_op_gpu.h
View File

@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#if !GOOGLE_CUDA
#error This file must only be included when building with Cuda support
#if !GOOGLE_CUDA && !TENSORFLOW_USE_ROCM
#error This file must only be included when building with Cuda or ROCm support
#endif
#ifndef TENSORFLOW_CORE_KERNELS_MAXPOOLING_OP_GPU_H_

4
tensorflow/core/kernels/pooling_ops_common_gpu.h
View File

@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#if !GOOGLE_CUDA
#error This file must only be included when building with Cuda support
#if !GOOGLE_CUDA && !TENSORFLOW_USE_ROCM
#error This file must only be included when building with Cuda or ROCm support
#endif
#ifndef TENSORFLOW_CORE_KERNELS_POOLING_OPS_COMMON_GPU_H_