1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_KERNELS_SPLIT_LIB_GPU_H_ 17 #define TENSORFLOW_CORE_KERNELS_SPLIT_LIB_GPU_H_ 18 19 #define EIGEN_USE_THREADS 20 #define EIGEN_USE_GPU 21 22 #include <memory> 23 #include <vector> 24 25 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 26 #include "tensorflow/core/framework/register_types.h" 27 #include "tensorflow/core/kernels/gpu_device_array_gpu.h" 28 #include "tensorflow/core/kernels/split_lib.h" 29 30 namespace tensorflow { 31 32 template <typename T> 33 struct SplitOpGPULaunch { 34 void Run(const Eigen::GpuDevice& d, const T* input, int32_t prefix_dim_size, 35 int32_t split_dim_size, int32_t suffix_dim_size, 36 const GpuDeviceArrayStruct<T*>& output_ptr_data); 37 }; 38 39 template <typename T, typename IntType> 40 struct SplitVOpGPULaunch { 41 void Run(const Eigen::GpuDevice& d, bool fixed, const T* input, 42 int total_cols, int total_rows, 43 const GpuDeviceArrayStruct<IntType>& output_scan, 44 const GpuDeviceArrayStruct<T*>& output_ptr_data); 45 }; 46 47 // Explicit instantiations in split_lib_gpu.cu.cc. 48 #define REGISTER_GPU_KERNEL(T) \ 49 extern template struct SplitOpGPULaunch<T>; \ 50 extern template struct SplitVOpGPULaunch<T, int32>; \ 51 extern template struct SplitVOpGPULaunch<T, int64>; 52 53 TF_CALL_bfloat16(REGISTER_GPU_KERNEL); 54 TF_CALL_uint8(REGISTER_GPU_KERNEL); 55 TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNEL); 56 #undef REGISTER_GPU_KERNEL 57 58 } // namespace tensorflow 59 60 #endif // TENSORFLOW_CORE_KERNELS_SPLIT_LIB_GPU_H_ 61