aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/kernels/eigen_pooling.h
diff options
context:
space:
mode:
authorGravatar Benoit Steiner <bsteiner@google.com>2016-07-14 13:59:03 -0800
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2016-07-14 15:03:48 -0700
commitdd410bc164d4026a2feb5baae26e6df7a2005d89 (patch)
treec493b1dac80b942fef606aa0a8c8d1023c086e81 /tensorflow/core/kernels/eigen_pooling.h
parent7520953ac6ee65f7cca8ab3f124fb2cb1629918e (diff)
Backported some changes to the reduction code from upstream Eigen to keep the code in sync.
Change: 127477364
Diffstat (limited to 'tensorflow/core/kernels/eigen_pooling.h')
-rw-r--r--tensorflow/core/kernels/eigen_pooling.h18
1 files changed, 18 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/eigen_pooling.h b/tensorflow/core/kernels/eigen_pooling.h
index aa3b274893..96acbe824a 100644
--- a/tensorflow/core/kernels/eigen_pooling.h
+++ b/tensorflow/core/kernels/eigen_pooling.h
@@ -376,6 +376,24 @@ struct AvgPoolMeanReducer {
Packet packetCount_;
};
+template <typename Device>
+struct reducer_traits<AvgPoolMeanReducer<float>, Device> {
+ enum {
+ Cost = 1,
+#if (EIGEN_ARCH_i386 || EIGEN_ARCH_x86_64) && !defined(__CUDACC__)
+ // We only support packet access for floats.
+ PacketAccess = true
+#else
+ PacketAccess = false
+#endif
+ };
+};
+
+template <>
+struct reducer_traits<AvgPoolMeanReducer<float>, GpuDevice> {
+ enum { Cost = 1, PacketAccess = false };
+};
+
} // namespace internal
#if !defined(EIGEN_HAS_INDEX_LIST)