| 1069 | template <typename T, typename Op, typename OUT_T, typename IN_T, |
| 1070 | typename ReductionAxes> |
| 1071 | void ReduceImpl(OpKernelContext* ctx, OUT_T out, IN_T in, int in_rank, |
| 1072 | int in_dim0, int in_dim1, int in_dim2, int out_rank, |
| 1073 | const ReductionAxes& reduction_axes, Op op) { |
| 1074 | T init = reduction_op_helper::IdentityValue<T, Op>()(); |
| 1075 | const gpuStream_t& cu_stream = GetGpuStream(ctx); |
| 1076 | if (out_rank == 0) { |
| 1077 | const int in_size = in_dim0 * in_dim1 * in_dim2; |
| 1078 | LaunchScalarReduction(ctx, out, in, in_size, op, init, cu_stream); |
| 1079 | } else if (in_rank == 2 && out_rank == 1 && |
| 1080 | reduction_axes[0] == 1) { // row reduction |
| 1081 | LaunchRowReduction(ctx, out, in, in_dim0, in_dim1, op, init, cu_stream); |
| 1082 | } else if (in_rank == 2 && out_rank == 1 && |
| 1083 | reduction_axes[0] == 0) { // column reduction |
| 1084 | LaunchColumnReduction(ctx, out, in, in_dim0, in_dim1, op, init, cu_stream); |
| 1085 | } else if (in_rank == 3 && out_rank == 2 && reduction_axes[0] == 1) { |
| 1086 | int elems_per_thread = in_dim1 / (in_dim0 * in_dim2); |
| 1087 | if (elems_per_thread >= 16) { |
| 1088 | Launch3DYReduction(ctx, out, in, in_dim0, in_dim1, in_dim2, op, init, |
| 1089 | cu_stream); |
| 1090 | } else { |
| 1091 | Launch3DYReductionSimple(ctx, out, in, in_dim0, in_dim1, in_dim2, op, |
| 1092 | init, cu_stream); |
| 1093 | } |
| 1094 | } else if (in_rank == 3 && out_rank == 1 && reduction_axes[0] == 0 && |
| 1095 | reduction_axes[1] == 2) { |
| 1096 | Launch3DXZReduction(ctx, out, in, in_dim0, in_dim1, in_dim2, op, init, |
| 1097 | cu_stream); |
| 1098 | } else { |
| 1099 | std::stringstream ss; |
| 1100 | ss << "Invalid reduction requested: in_rank, out_rank, axes " << in_rank |
| 1101 | << " " << out_rank; |
| 1102 | if (out_rank == 1) ss << " " << reduction_axes[0]; |
| 1103 | if (out_rank == 2) ss << " " << reduction_axes[1]; |
| 1104 | LOG(FATAL) << ss.str(); |
| 1105 | } |
| 1106 | } |
| 1107 | |
| 1108 | template <typename Reducer> |
| 1109 | struct ReduceFunctor<GPUDevice, Reducer> { |
nothing calls this directly
no test coverage detected