| 203 | } |
| 204 | |
| 205 | void Compute(OpKernelContext* context) override { |
| 206 | const Tensor& logits_in_ = context->input(0); |
| 207 | OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(logits_in_.shape()), |
| 208 | errors::InvalidArgument("logits must have >= 1 dimension, got ", |
| 209 | logits_in_.shape().DebugString())); |
| 210 | auto logits_in = logits_in_.flat_inner_dims<T>(); |
| 211 | const int rows = logits_in.dimension(0); |
| 212 | const int cols = logits_in.dimension(1); |
| 213 | Tensor* softmax_out = nullptr; |
| 214 | OP_REQUIRES_OK(context, context->forward_input_or_allocate_output( |
| 215 | {0}, 0, logits_in_.shape(), &softmax_out)); |
| 216 | |
| 217 | const auto& cu_stream = GetGpuStream(context); |
| 218 | if (logits_in_.NumElements() > 0) { |
| 219 | Tensor max_logits; |
| 220 | Tensor sum_probs; |
| 221 | OP_REQUIRES_OK(context, |
| 222 | context->allocate_temp(DataTypeToEnum<T>::value, |
| 223 | softmax_out->shape(), &max_logits)); |
| 224 | |
| 225 | typedef typename softmax_traits<T>::accumulator_type acc_type; |
| 226 | OP_REQUIRES_OK(context, |
| 227 | context->allocate_temp(DataTypeToEnum<acc_type>::value, |
| 228 | softmax_out->shape(), &sum_probs)); |
| 229 | |
| 230 | DoRowReduction<T, gpuprim::Max, const T*>( |
| 231 | context, const_cast<T*>(max_logits.flat<T>().data()), |
| 232 | reinterpret_cast<const T*>(logits_in_.flat<T>().data()), rows, cols); |
| 233 | |
| 234 | |
| 235 | gpuprim::CountingInputIterator<int> counting_iterator(0); |
| 236 | using InputIterType = |
| 237 | gpuprim::TransformInputIterator<acc_type, |
| 238 | SubtractAndExpFunctor<T, acc_type>, |
| 239 | gpuprim::CountingInputIterator<int>>; |
| 240 | |
| 241 | InputIterType input_itr( |
| 242 | counting_iterator, |
| 243 | SubtractAndExpFunctor<T, acc_type>( |
| 244 | reinterpret_cast<const T*>(logits_in_.flat<T>().data()), |
| 245 | reinterpret_cast<const T*>(max_logits.flat<T>().data()), cols)); |
| 246 | |
| 247 | DoRowReduction<acc_type, gpuprim::Sum, InputIterType>( |
| 248 | context, const_cast<acc_type*>(sum_probs.flat<acc_type>().data()), |
| 249 | input_itr, rows, cols); |
| 250 | |
| 251 | auto in_ptr = reinterpret_cast<uintptr_t>(logits_in_.flat<T>().data()); |
| 252 | auto out_ptr = reinterpret_cast<uintptr_t>(softmax_out->flat<T>().data()); |
| 253 | bool aligned = in_ptr % 16 == 0 && out_ptr % 16 == 0; |
| 254 | |
| 255 | const int numThreadsPerBlock = 128; |
| 256 | if (DataTypeToEnum<T>::value == DT_HALF && aligned) { |
| 257 | const int kUnroll = 8; |
| 258 | const int numBlocks = Eigen::divup(rows * cols, |
| 259 | numThreadsPerBlock * kUnroll); |
| 260 | TF_CHECK_OK(GpuLaunchKernel( |
| 261 | GenerateNormalizedProb<T, acc_type, kUnroll>, numBlocks, |
| 262 | numThreadsPerBlock, 0, cu_stream, |
nothing calls this directly
no test coverage detected