| 136 | |
| 137 | template <typename T, typename IntType> |
| 138 | void ConcatGPUImpl(const Eigen::GpuDevice& gpu_device, |
| 139 | const GpuDeviceArrayStruct<const T*>& input_ptrs, |
| 140 | const GpuDeviceArrayStruct<IntType>& output_scan, |
| 141 | bool fixed_size, int split_size, |
| 142 | typename TTypes<T, 2>::Matrix* output) { |
| 143 | auto config = GetGpu2DLaunchConfig(output->dimension(1), output->dimension(0), |
| 144 | gpu_device); |
| 145 | |
| 146 | if (fixed_size) { |
| 147 | TF_CHECK_OK(GpuLaunchKernel( |
| 148 | concat_fixed_kernel<T, IntType>, config.block_count, |
| 149 | config.thread_per_block, 0, gpu_device.stream(), input_ptrs, split_size, |
| 150 | static_cast<int>(output->dimension(0)), |
| 151 | static_cast<int>(output->dimension(1)), output->data())); |
| 152 | } else { |
| 153 | IntType smem_max = gpu_device.sharedMemPerBlock(); |
| 154 | IntType smem_usage = output_scan.size * sizeof(IntType); |
| 155 | // performance crossover is less than using maximum available shared memory |
| 156 | // on most processors |
| 157 | // possibly due to decreasing occupancy |
| 158 | // 4096 inputs is a lot, most code will take the smem path |
| 159 | const int32 kMaxSmemBytesPerformance = 16384; |
| 160 | if (smem_usage < smem_max && smem_usage < kMaxSmemBytesPerformance) { |
| 161 | TF_CHECK_OK(GpuLaunchKernel( |
| 162 | concat_variable_kernel<T, IntType, true>, config.block_count, |
| 163 | config.thread_per_block, smem_usage, gpu_device.stream(), input_ptrs, |
| 164 | output_scan, static_cast<IntType>(output->dimension(0)), |
| 165 | static_cast<IntType>(output->dimension(1)), output->data())); |
| 166 | } else { |
| 167 | TF_CHECK_OK(GpuLaunchKernel( |
| 168 | concat_variable_kernel<T, IntType, false>, config.block_count, |
| 169 | config.thread_per_block, 0, gpu_device.stream(), input_ptrs, |
| 170 | output_scan, static_cast<IntType>(output->dimension(0)), |
| 171 | static_cast<IntType>(output->dimension(1)), output->data())); |
| 172 | } |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | #define REGISTER_GPUCONCAT32(T) \ |
| 177 | template void ConcatGPUSlice<T, int32>( \ |
nothing calls this directly
no test coverage detected