_vec_log_softmax_lastdim Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/SoftMaxKernel.cpp lines 34–59

template <typename scalar_t>
inline void _vec_log_softmax_lastdim(
    const scalar_t* input_data_base,
    scalar_t* output_data_base,
    int64_t outer_size,
    int64_t dim_size) {
  const auto chunk_size = vec_log_softmax_lastdim_chunk_size<scalar_t>(
      at::internal::GRAIN_SIZE,
      outer_size,
      dim_size);
  // Note: grain_size value of 0
  // We don't change the number of OpenMP threads in the OpenMP thread-pool,
  // so some threads do useful work, while others don't.
  // We can simply use grain_size of 0 & rely upon invoke_parallel to distribute
  // work among threads in an equitable manner. We compute CHUNK_SIZE to ensure
  // each thread's computations would be efficient.
  parallel_for(0, outer_size, 0, [&](int64_t begin, int64_t end) {
    serial_vec_log_softmax_lastdim_range(
        input_data_base,
        output_data_base,
        dim_size,
        chunk_size,
        begin,
        end);
  });
}