aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2018-08-09 03:24:10 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-08-09 03:28:03 -0700
commitcfdc46956574de703df2251ac52d30e593d8d417 (patch)
tree1f5cefb5b4fe8b5790aaad9d2f931cf137f1b4fb
parente2d9cd3fa7612a712a7adeb750a1f2c9bd05ae17 (diff)
Enable batch matmul for half types again which was accidentally disabled in CL 197137612.
PiperOrigin-RevId: 208022865
-rw-r--r--tensorflow/core/kernels/BUILD2
-rw-r--r--tensorflow/core/kernels/batch_matmul_op_real.cc3
2 files changed, 3 insertions, 2 deletions
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index e66e9a10e7..d142e36772 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -2854,6 +2854,8 @@ tf_kernel_library(
srcs = [] + if_mkl([
"mkl_batch_matmul_op.cc",
]),
+ # <prefix>*impl.h are excluded by default from the CPU build, add explicitly.
+ hdrs = ["batch_matmul_op_impl.h"],
# Override EIGEN_STRONG_INLINE to inline when --define=override_eigen_strong_inline=true,
# to avoid long compiling time. See https://github.com/tensorflow/tensorflow/issues/10521
copts = if_override_eigen_strong_inline(["/DEIGEN_STRONG_INLINE=inline"]),
diff --git a/tensorflow/core/kernels/batch_matmul_op_real.cc b/tensorflow/core/kernels/batch_matmul_op_real.cc
index fe259c1634..aa7a2752e8 100644
--- a/tensorflow/core/kernels/batch_matmul_op_real.cc
+++ b/tensorflow/core/kernels/batch_matmul_op_real.cc
@@ -31,8 +31,7 @@ TF_CALL_int32(REGISTER_BATCH_MATMUL_CPU);
#if GOOGLE_CUDA
TF_CALL_float(REGISTER_BATCH_MATMUL_GPU);
TF_CALL_double(REGISTER_BATCH_MATMUL_GPU);
-// TODO(csigg): Implement Stream::ThenBlasGemv for Eigen::half and uncomment.
-// TF_CALL_half(REGISTER_BATCH_MATMUL_GPU);
+TF_CALL_half(REGISTER_BATCH_MATMUL_GPU);
#endif // GOOGLE_CUDA
#ifdef TENSORFLOW_USE_SYCL