message: Use parallelFor in GatherNdSlice. Saw speedups from 3-12% on Intel Skylake.
diff: @@ -101,40 +101,21 @@ struct GatherNdSlice<CPUDevice, T, Index, IXDIM> {
  typename TTypes<Index>::ConstMatrix Tindices,
  typename TTypes<T>::Matrix Tout) {
  std::atomic<Index> error_loc(-1);
  -
  - const Eigen::DenseIndex batch_size = Tindices.dimension(0);
  -#if !defined(EIGEN_HAS_INDEX_LIST)
  - Eigen::Tensor<Eigen::DenseIndex, 1>::Dimensions reshape_dims{{ 1 }};
  - Eigen::array<Eigen::DenseIndex, 1> broadcast_dims{{ batch_size }};
  -#else
  -#if !defined(INTEL_MKL) || !defined(ENABLE_MKL)
  - Eigen::IndexList<Eigen::type2index<1> > reshape_dims;
  -#endif // defined(INTEL_MKL) && defined(ENABLE_MKL)
  - Eigen::IndexList<Eigen::DenseIndex> broadcast_dims;
  - broadcast_dims.set(0, batch_size);
  -#endif // !defined(EIGEN_HAS_INDEX_LIST)
  + const Eigen::Index batch_size = Tindices.dimension(0);
  generator::GatherNdSliceGenerator<T, Index, IXDIM> gather_nd_generator(
  slice_size, Tindices, Tparams, Tout, &error_loc);
  -// TODO(b/137289929): Parallelize this with ParallelFor and remove OpenMP call.
  -#if defined(INTEL_MKL) && defined(ENABLE_MKL)
  -// Eigen implementation below is not highly performant. gather_nd_generator
  -// does not seem to be called in parallel, leading to very poor performance.
  -// Additionally, since it uses scalar (Tscratch) to invoke 'generate', it
  -// needs to go through redundant operations like 'reshape', 'broadcast' and
  -// 'sum'. OpenMP loop below essentially does same thing as Eigen code, but
  -// is considerably more efficient.
  -#pragma omp parallel for
  - for (Eigen::DenseIndex i = 0; i < batch_size; i++) {
  - const Eigen::array<Eigen::DenseIndex, 1> loc{i};
  + auto compute_shard = [&](Eigen::Index begin, Eigen::Index end) {
  + for (Eigen::Index i = begin; i < end; ++i) {
  + const Eigen::array<Eigen::Index, 1> loc{i};
  gather_nd_generator(loc);
  }
  -#else // INTEL_MKL && ENABLE_MKL
  - Tscratch.device(d) = Tscratch.reshape(reshape_dims)
  - .broadcast(broadcast_dims)
  - .generate(gather_nd_generator)
  - .sum();
  -#endif // INTEL_MKL && ENABLE_MKL
  + };
  + Eigen::Index bytes_moved = sizeof(T) * (slice_size + IXDIM);
  + auto cost = Eigen::TensorOpCost(bytes_moved /* bytes loaded */,
  + bytes_moved /* bytes stored */,
  + slice_size + IXDIM /* compute cycles */);
  + d.parallelFor(batch_size, cost, compute_shard);
  // error_loc() returns -1 if there's no out-of-bounds index,
  // otherwise it returns the location of an OOB index in Tindices.
  