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

Split slice_op into multiple files to improve compilation times for CPU.

Browse Source 
Change: 130469255
This commit is contained in:
A. Unique TensorFlower 2016-08-16 16:45:09 -08:00 committed by TensorFlower Gardener
parent 014cc05a11
commit b07575c710
10 changed files with 187 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/contrib/makefile/tf_op_files.txt
View File

@ -26,6 +26,12 @@ tensorflow/core/kernels/sparse_to_dense_op.cc
tensorflow/core/kernels/softsign_op.cc
tensorflow/core/kernels/softplus_op.cc
tensorflow/core/kernels/softmax_op.cc
tensorflow/core/kernels/slice_op_cpu_impl_1.cc
tensorflow/core/kernels/slice_op_cpu_impl_2.cc
tensorflow/core/kernels/slice_op_cpu_impl_3.cc
tensorflow/core/kernels/slice_op_cpu_impl_4.cc
tensorflow/core/kernels/slice_op_cpu_impl_5.cc
tensorflow/core/kernels/slice_op_cpu_impl_6.cc
tensorflow/core/kernels/slice_op.cc
tensorflow/core/kernels/shape_ops.cc
tensorflow/core/kernels/session_ops.cc

7
tensorflow/core/kernels/BUILD
View File

@ -1883,6 +1883,13 @@ filegroup(
"shape_ops.cc",
"slice_op.cc",
"slice_op.h",
"slice_op_cpu_impl.h",
"slice_op_cpu_impl_1.cc",
"slice_op_cpu_impl_2.cc",
"slice_op_cpu_impl_3.cc",
"slice_op_cpu_impl_4.cc",
"slice_op_cpu_impl_5.cc",
"slice_op_cpu_impl_6.cc",
"softmax_op.cc",
"softmax_op.h",
"softmax_op_functor.h",

27
tensorflow/core/kernels/slice_op.cc
View File

@ -201,6 +201,33 @@ class SliceOp : public OpKernel {
}
};
// Forward declarations of the functor specializations for declared in the
// sharded source files.
namespace functor {
#define DECLARE_CPU_SPEC(T, NDIM) \
template <> \
void Slice<CPUDevice, T, NDIM>::operator()( \
const CPUDevice& d, typename TTypes<T, NDIM>::Tensor output, \
typename TTypes<T, NDIM>::ConstTensor input, \
const Eigen::DSizes<Eigen::DenseIndex, NDIM>& indices, \
const Eigen::DSizes<Eigen::DenseIndex, NDIM>& sizes); \
extern template struct Slice<CPUDevice, T, NDIM>;
#define DECLARE_FOR_N(T) \
DECLARE_CPU_SPEC(T, 1); \
DECLARE_CPU_SPEC(T, 2); \
DECLARE_CPU_SPEC(T, 3); \
DECLARE_CPU_SPEC(T, 4); \
DECLARE_CPU_SPEC(T, 5); \
DECLARE_CPU_SPEC(T, 6);
TF_CALL_ALL_TYPES(DECLARE_FOR_N);
DECLARE_FOR_N(bfloat16);
#undef DECLARE_FOR_N
#undef DECLARE_CPU_SPEC
} // namespace functor
#define REGISTER_SLICE(type) \
REGISTER_KERNEL_BUILDER(Name("Slice") \
.Device(DEVICE_CPU) \

39
tensorflow/core/kernels/slice_op_cpu_impl.h Normal file
View File

@ -0,0 +1,39 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef THIRD_PARTY_TENSORFLOW_CORE_KERNELS_SLICE_OP_CPU_IMPL_H_
#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_SLICE_OP_CPU_IMPL_H_
#define EIGEN_USE_THREADS
#include "tensorflow/core/framework/bfloat16.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/kernels/slice_op.h"
namespace tensorflow {
using CpuDevice = Eigen::ThreadPoolDevice;
#define DEFINE_CPU_KERNELS(T) \
template struct functor::Slice<CpuDevice, T, CPU_PROVIDED_IXDIM>;
TF_CALL_ALL_TYPES(DEFINE_CPU_KERNELS);
DEFINE_CPU_KERNELS(bfloat16);
#undef DEFINE_GPU_KERNELS
} // namespace tensorflow
#endif // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_SLICE_OP_CPU_IMPL_H_

18
tensorflow/core/kernels/slice_op_cpu_impl_1.cc Normal file
View File

@ -0,0 +1,18 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define CPU_PROVIDED_IXDIM 1
#include "tensorflow/core/kernels/slice_op_cpu_impl.h"
#undef CPU_PROVIDED_IXDIM

18
tensorflow/core/kernels/slice_op_cpu_impl_2.cc Normal file
View File

@ -0,0 +1,18 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define CPU_PROVIDED_IXDIM 2
#include "tensorflow/core/kernels/slice_op_cpu_impl.h"
#undef CPU_PROVIDED_IXDIM

18
tensorflow/core/kernels/slice_op_cpu_impl_3.cc Normal file
View File

@ -0,0 +1,18 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define CPU_PROVIDED_IXDIM 3
#include "tensorflow/core/kernels/slice_op_cpu_impl.h"
#undef CPU_PROVIDED_IXDIM

18
tensorflow/core/kernels/slice_op_cpu_impl_4.cc Normal file
View File

@ -0,0 +1,18 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define CPU_PROVIDED_IXDIM 4
#include "tensorflow/core/kernels/slice_op_cpu_impl.h"
#undef CPU_PROVIDED_IXDIM

18
tensorflow/core/kernels/slice_op_cpu_impl_5.cc Normal file
View File

@ -0,0 +1,18 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define CPU_PROVIDED_IXDIM 5
#include "tensorflow/core/kernels/slice_op_cpu_impl.h"
#undef CPU_PROVIDED_IXDIM

18
tensorflow/core/kernels/slice_op_cpu_impl_6.cc Normal file
View File

@ -0,0 +1,18 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define CPU_PROVIDED_IXDIM 6
#include "tensorflow/core/kernels/slice_op_cpu_impl.h"
#undef CPU_PROVIDED_IXDIM