diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-04-09 11:14:25 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-09 11:16:51 -0700 |
| commit | 20e1f3a852c83d9369c7d56a943fe6b8f9b88644 (patch) | |
| tree | f54acb82a9a7a5b6324ee0a3e21cb6025b84b961 | |
| parent | 1b97ea722cea69a8d9a6fe3bd515e22f356d40ab (diff) | |
Register tf.concat with uint8 data type.
PiperOrigin-RevId: 192154998
| -rw-r--r-- | tensorflow/core/kernels/concat_lib_gpu.cc | 1 | ||||
| -rw-r--r-- | tensorflow/core/kernels/concat_lib_gpu_impl.cu.cc | 4 | ||||
| -rw-r--r-- | tensorflow/core/kernels/concat_op.cc | 1 | ||||
| -rw-r--r-- | tensorflow/core/kernels/concat_op_test.cc | 4 |
4 files changed, 10 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/concat_lib_gpu.cc b/tensorflow/core/kernels/concat_lib_gpu.cc index d8643c0b2f..93e392d303 100644 --- a/tensorflow/core/kernels/concat_lib_gpu.cc +++ b/tensorflow/core/kernels/concat_lib_gpu.cc @@ -118,6 +118,7 @@ TF_CALL_complex128(REGISTER); TF_CALL_int64(REGISTER); TF_CALL_bfloat16(REGISTER); TF_CALL_bool(REGISTER); +TF_CALL_uint8(REGISTER); #undef REGISTER diff --git a/tensorflow/core/kernels/concat_lib_gpu_impl.cu.cc b/tensorflow/core/kernels/concat_lib_gpu_impl.cu.cc index 0f7adaf24a..a561d918bd 100644 --- a/tensorflow/core/kernels/concat_lib_gpu_impl.cu.cc +++ b/tensorflow/core/kernels/concat_lib_gpu_impl.cu.cc @@ -202,6 +202,7 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPUCONCAT32); TF_CALL_complex64(REGISTER_GPUCONCAT32); TF_CALL_complex128(REGISTER_GPUCONCAT32); TF_CALL_int64(REGISTER_GPUCONCAT32); +TF_CALL_uint8(REGISTER_GPUCONCAT32); REGISTER_GPUCONCAT32(bfloat16); REGISTER_GPUCONCAT32(bool); @@ -209,6 +210,7 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPUCONCAT64); TF_CALL_complex64(REGISTER_GPUCONCAT64); TF_CALL_complex128(REGISTER_GPUCONCAT64); TF_CALL_int64(REGISTER_GPUCONCAT64); +TF_CALL_uint8(REGISTER_GPUCONCAT64); REGISTER_GPUCONCAT64(bfloat16); REGISTER_GPUCONCAT64(bool); @@ -216,6 +218,7 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU32); TF_CALL_complex64(REGISTER_GPU32); TF_CALL_complex128(REGISTER_GPU32); TF_CALL_int64(REGISTER_GPU32); +TF_CALL_uint8(REGISTER_GPU32); REGISTER_GPU32(bfloat16); REGISTER_GPU32(bool); @@ -223,6 +226,7 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU64); TF_CALL_complex64(REGISTER_GPU64); TF_CALL_complex128(REGISTER_GPU64); TF_CALL_int64(REGISTER_GPU64); +TF_CALL_uint8(REGISTER_GPU64); REGISTER_GPU64(bfloat16); REGISTER_GPU64(bool); diff --git a/tensorflow/core/kernels/concat_op.cc b/tensorflow/core/kernels/concat_op.cc index f16766315f..a87b63f913 100644 --- a/tensorflow/core/kernels/concat_op.cc +++ b/tensorflow/core/kernels/concat_op.cc @@ -212,6 +212,7 @@ REGISTER_CONCAT(qint32); TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU); REGISTER_GPU(bfloat16); +TF_CALL_uint8(REGISTER_GPU); TF_CALL_complex64(REGISTER_GPU); TF_CALL_complex128(REGISTER_GPU); TF_CALL_int64(REGISTER_GPU); diff --git a/tensorflow/core/kernels/concat_op_test.cc b/tensorflow/core/kernels/concat_op_test.cc index e3ba8ae9f6..39b44b2fcc 100644 --- a/tensorflow/core/kernels/concat_op_test.cc +++ b/tensorflow/core/kernels/concat_op_test.cc @@ -78,6 +78,9 @@ static void BM_ConcatDim1Float(int iters, int dim2) { BENCHMARK(BM_ConcatDim0Float)->Arg(1000)->Arg(100000)->Arg(1000000); BENCHMARK(BM_ConcatDim1Float)->Arg(1000)->Arg(100000)->Arg(1000000); +static void BM_ConcatDim1uint8(int iters, int dim2) { + ConcatHelper<uint8>(iters, 1, dim2); +} static void BM_ConcatDim1int16(int iters, int dim2) { ConcatHelper<int16>(iters, 1, dim2); } @@ -85,6 +88,7 @@ static void BM_ConcatDim1bfloat16(int iters, int dim2) { ConcatHelper<bfloat16>(iters, 1, dim2); } +BENCHMARK(BM_ConcatDim1uint8)->Arg(1000)->Arg(100000)->Arg(1000000); BENCHMARK(BM_ConcatDim1int16)->Arg(1000)->Arg(100000)->Arg(1000000); BENCHMARK(BM_ConcatDim1bfloat16)->Arg(1000)->Arg(100000)->Arg(1000000); |