vec_log_softmax_lastdim_chunk_size Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/LogSoftmaxKernelImpl.h lines 18–28

template <typename scalar_t>
int64_t vec_log_softmax_lastdim_chunk_size(int64_t grain_size, int64_t outer_size, int64_t dim_size) {
  // Coincidentally, at::internal::GRAIN_SIZE is 32768, which is equal to the
  // size of L1D cache on many processors. Some processors have 48 KB L1D cache
  // nowadays, so maybe in the future, we can leverage the knowledge of a
  // machine's L1D cache size.
  int64_t MAX_CHUNK_SIZE = std::max<int64_t>(
      1,
      grain_size / (sizeof(scalar_t) * dim_size));
  return std::min<int64_t>(MAX_CHUNK_SIZE, outer_size);
}