| 69 | |
| 70 | template <typename T> |
| 71 | void ConcatGPU( |
| 72 | OpKernelContext* c, |
| 73 | const std::vector<std::unique_ptr<typename TTypes<T, 2>::ConstMatrix>>& |
| 74 | inputs_flat, |
| 75 | Tensor* output, typename TTypes<T, 2>::Tensor* output_flat) { |
| 76 | if (inputs_flat.size() < 16) { |
| 77 | if (output->NumElements() < std::numeric_limits<int32>::max()) { |
| 78 | ConcatGPUSlice<T, int32>(c->eigen_gpu_device(), inputs_flat, output_flat); |
| 79 | } else { |
| 80 | ConcatGPUSlice<T, int64>(c->eigen_gpu_device(), inputs_flat, output_flat); |
| 81 | } |
| 82 | } else { |
| 83 | // Switching indexing to int64 might cause performance issues. |
| 84 | // Hence, we keep int32 indexing in the GPU kernel unless we need to |
| 85 | // switch to int64. |
| 86 | if (output->NumElements() < std::numeric_limits<int32>::max()) { |
| 87 | ConcatGPUCall<T, int32>(c, inputs_flat, output_flat); |
| 88 | } else { |
| 89 | ConcatGPUCall<T, int64>(c, inputs_flat, output_flat); |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | #define REGISTER(T) \ |
| 95 | template void ConcatGPU<T>( \ |
nothing calls this directly
no test coverage detected