| 887 | |
| 888 | template <typename T, typename Op, typename OUT_T, typename IN_T> |
| 889 | void Launch3DYReduction(OpKernelContext* ctx, OUT_T out, IN_T in, int extent_x, |
| 890 | int extent_y, int extent_z, Op op, T init, |
| 891 | const gpuStream_t& cu_stream) { |
| 892 | int threads_per_block = 128; |
| 893 | |
| 894 | int n_group_in = extent_y; |
| 895 | int n_size = extent_z; |
| 896 | |
| 897 | // Calculate and allocate temporary space |
| 898 | std::size_t temp_storage_bytes = 0; |
| 899 | // A plane's size is n_group_in * n_size. We make sure no single plane crosses |
| 900 | // more than one thread block, meaning a thread block will handle one whole |
| 901 | // plane or multiple planes in the second stage. Also, It may handle a partial |
| 902 | // plane when n_size is too large and the while-loop will stop at |
| 903 | // n_group_in = 1, where we directly copy the temp to output in the next |
| 904 | // stage. |
| 905 | while (n_group_in >= 2 && n_group_in * n_size > threads_per_block) { |
| 906 | int n_group_out = std::max(1, n_group_in / (2 * kUnroll)); |
| 907 | temp_storage_bytes += n_group_out * n_size; |
| 908 | n_group_in = n_group_out; |
| 909 | } |
| 910 | temp_storage_bytes *= extent_x * sizeof(T); |
| 911 | Tensor temp_storage; |
| 912 | OP_REQUIRES_OK( |
| 913 | ctx, ctx->allocate_temp( |
| 914 | DT_INT8, TensorShape({static_cast<int64>(temp_storage_bytes)}), |
| 915 | &temp_storage)); |
| 916 | |
| 917 | // Reduction |
| 918 | n_group_in = extent_y; |
| 919 | int temp_in_offset = -1; |
| 920 | int temp_out_offset = 0; |
| 921 | int num_blocks; |
| 922 | while (n_group_in >= 2 && n_group_in * n_size > threads_per_block) { |
| 923 | int n_group_out = std::max(1, n_group_in / (2 * kUnroll)); |
| 924 | num_blocks = |
| 925 | Eigen::divup(extent_x * n_group_out * n_size, threads_per_block); |
| 926 | TF_CHECK_OK(GpuLaunchKernel( |
| 927 | ColumnReduceInToTempKernel<IN_T, Op>, num_blocks, threads_per_block, 0, |
| 928 | cu_stream, (void*)(temp_storage.flat<int8_t>().data()), temp_in_offset, |
| 929 | temp_out_offset, in, extent_x, n_group_in, extent_z, op)); |
| 930 | |
| 931 | n_group_in = n_group_out; |
| 932 | temp_in_offset = temp_out_offset; |
| 933 | temp_out_offset = temp_in_offset + extent_x * n_group_out * n_size; |
| 934 | } |
| 935 | |
| 936 | if (n_group_in * n_size <= threads_per_block) { |
| 937 | num_blocks = extent_x; |
| 938 | } else { |
| 939 | DCHECK_EQ(1, n_group_in); |
| 940 | num_blocks = Eigen::divup(extent_x * n_size, threads_per_block); |
| 941 | } |
| 942 | |
| 943 | TF_CHECK_OK(GpuLaunchKernel( |
| 944 | ColumnReduceTempToOutKernel<IN_T, OUT_T, Op>, num_blocks, |
| 945 | threads_per_block, 2 * sizeof(T) * threads_per_block, cu_stream, |
| 946 | (void*)(temp_storage.flat<int8_t>().data()), temp_in_offset, in, out, |
no test coverage detected