[oneDNN] Make cross_entropy calculations parallel on CPU
Hi @penpornk Can you please review this PR ? Thank you! Hi @AnetaKaczynska Can you please check @cantonios's comments and keep us posted? Thank you! This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you. Hi @AnetaKaczynska Any update on this PR? Please. Thank you! Hi @cantonios, I just ran benchmarks again and it still has better performance. Can you please give me more details to reproduce the worse performance case? Was it on CPU or GPU? > Hi @cantonios, I just ran benchmarks again and it still has better performance. Can you please give me more details to reproduce the worse performance case? Was it on CPU or GPU?

CPU.  It's only the CPU kernel that's modified, and the benchmark numbers in the revert commit list the CPU benchmark numbers. @cantonios One more thing, can you please share system specifications that were used to run this benchmark? What CPU did you use? > @cantonios One more thing, can you please share system specifications that were used to run this benchmark? What CPU did you use?

Similar results on both haswell and skylake. > > @cantonios One more thing, can you please share system specifications that were used to run this benchmark? What CPU did you use?
> 
> Similar results on both haswell and skylake.

@cantonios I ran it on skylake machine and got better results after optimization:

Benchmark | Items with optimization | Items without optimization
-- | -- | --
BM_Xent_16_10000_cpu_float/real_time | 182.368M/s | 103.502M/s
BM_Xent_32_10000_cpu_float/real_time | 289.048M/s | 176.654M/s
BM_Xent_64_10000_cpu_float/real_time | 428.617M/s | 298.246M/s
BM_Xent_16_10000_cpu_bfloat16/real_time | 151.684M/s | 86.4855M/s
BM_Xent_32_10000_cpu_bfloat16/real_time | 358.028M/s | 149.516M/s
BM_Xent_64_10000_cpu_bfloat16/real_time | 390.276M/s | 254.699M/s

I'm trying to get down to the bottom of why we're getting better performance. Can you please provide me more details about the system that was used for benchmarking and significant run options such as cores used? If possible it would be great if I could use the machine you ran it on. 



























































 > > > @cantonios One more thing, can you please share system specifications that were used to run this benchmark? What CPU did you use?
> > 
> > 
> > Similar results on both haswell and skylake.
> 
> @cantonios I ran it on skylake machine and got better results after optimization:
> 
> Benchmark	Items with optimization	Items without optimization
> BM_Xent_16_10000_cpu_float/real_time	182.368M/s	103.502M/s
> BM_Xent_32_10000_cpu_float/real_time	289.048M/s	176.654M/s
> BM_Xent_64_10000_cpu_float/real_time	428.617M/s	298.246M/s
> BM_Xent_16_10000_cpu_bfloat16/real_time	151.684M/s	86.4855M/s
> BM_Xent_32_10000_cpu_bfloat16/real_time	358.028M/s	149.516M/s
> BM_Xent_64_10000_cpu_bfloat16/real_time	390.276M/s	254.699M/s
> I'm trying to get down to the bottom of why we're getting better performance. Can you please provide me more details about the system that was used for benchmarking and significant run options such as cores used? If possible it would be great if I could use the machine you ran it on.

It's on our Google perflab machines.  You cannot access them.  This is without oneDNN, using TF's defaults.

Plus, this is only half the battle.  We also received reports from several teams within Google that their models no longer converged after this change - so there are numerical differences as well.
This optimization maximizes performance on CPU by implementing specialized CPU  [XentEigenImpl](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/xent_op.h#L52) class, which makes cross entropy parallel, while possibly using oneDNN for softmax calculations.
Original implementation calculates softmax_cross_entropy for a whole batch on a single thread. After optimization, each batch item is processed separately in parallel.

### Microbenchmarks
| Tensor Dims | Numeric Type | Before (ns) | After (ns) | Speed up |
|:----------:|:------------:|:-----------:|:----------:|:--------:|
|    1x10000 |    FLOAT32   |     593,294 |    641,580 |    0.92x |
|    2x10000 |    FLOAT32   |   1,027,291 |  1,027,994 |    1.00x |
|    4x10000 |    FLOAT32   |   1,363,172 |  1,225,472 |    1.11x |
|    8x10000 |    FLOAT32   |   2,871,685 |  1,177,349 |    2.44x |
|   16x10000 |    FLOAT32   |   3,395,105 |  1,468,886 |    2.31x |
|   32x10000 |    FLOAT32   |   4,318,577 |  1,643,422 |    2.63x |
|   64x10000 |    FLOAT32   |   5,095,850 |  1,786,184 |    2.85x |
|  128x10000 |    FLOAT32   |   6,130,740 |  2,179,011 |    2.81x |
|  256x10000 |    FLOAT32   |   6,216,732 |  2,858,630 |    2.17x |
|  512x10000 |    FLOAT32   |   7,033,319 |  3,207,006 |    2.19x |
| 1024x10000 |    FLOAT32   |  25,139,020 |  3,905,827 |    6.44x |
|    1x10000 |   BFLOAT16   |   1,449,381 |  1,655,900 |    0.88x |
|    2x10000 |   BFLOAT16   |   2,225,481 |  2,906,942 |    0.77x |
|    4x10000 |   BFLOAT16   |   3,543,601 |  2,921,440 |    1.21x |
|    8x10000 |   BFLOAT16   |   4,690,691 |  2,354,173 |    1.99x |
|   16x10000 |   BFLOAT16   |   6,033,001 |  2,568,764 |    2.35x |
|   32x10000 |   BFLOAT16   |   7,195,912 |  2,782,907 |    2.59x |
|   64x10000 |   BFLOAT16   |   7,901,942 |  3,114,311 |    2.54x |
|  128x10000 |   BFLOAT16   |   8,063,111 |  3,522,027 |    2.29x |
|  256x10000 |   BFLOAT16   |  10,001,815 |  4,475,105 |    2.23x |
|  512x10000 |   BFLOAT16   |  12,569,523 |  4,468,942 |    2.81x |
| 1024x10000 |   BFLOAT16   |  15,845,288 |  7,066,050 |    2.24x |
