[oneDNN] Add oneDNN version of SparseMatrixMatMul
Nice  cool > Could you please help make more changes? There are some more errors.

Sure. Sorry, didn't see those errors in my local tests, and I didn't see which target in the CI threw those errors. I think I've made all of the appropriate changes; @cantonios @penpornk does this look OK to you? @cantonios Made that change, and I think the current "Presubmit" test failure is unrelated to this PR. @cantonios I can see that some internal checks failed, but can't see what they are. > @cantonios I can see that some internal checks failed, but can't see what they are.

```
ERROR: tensorflow/core/kernels/mkl/BUILD:104:22: in cc_library rule //tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op: Visibility error:
target '//tensorflow/core/kernels/sparse:mat_mul_op.h' is not visible from
target '//tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op'
``` > > @cantonios I can see that some internal checks failed, but can't see what they are.
> 
> ```
> ERROR: tensorflow/core/kernels/mkl/BUILD:104:22: in cc_library rule //tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op: Visibility error:
> target '//tensorflow/core/kernels/sparse:mat_mul_op.h' is not visible from
> target '//tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op'
> ```

Is it possible to get a commandline reproducer for this error? Can't reproduce locally, so fixing it might be trial-and-error that spams this PR with commits and spams you with emails! > > > @cantonios I can see that some internal checks failed, but can't see what they are.
> > 
> > 
> > ```
> > ERROR: tensorflow/core/kernels/mkl/BUILD:104:22: in cc_library rule //tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op: Visibility error:
> > target '//tensorflow/core/kernels/sparse:mat_mul_op.h' is not visible from
> > target '//tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op'
> > ```
> 
> Is it possible to get a commandline reproducer for this error? Can't reproduce locally, so fixing it might be trial-and-error that spams this PR with commits and spams you with emails!

I'm not sure... I'm surprised bazel doesn't find the dependency issue.  Maybe Google's internal tooling is more strict? > > > > @cantonios I can see that some internal checks failed, but can't see what they are.
> > > 
> > > 
> > > ```
> > > ERROR: tensorflow/core/kernels/mkl/BUILD:104:22: in cc_library rule //tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op: Visibility error:
> > > target '//tensorflow/core/kernels/sparse:mat_mul_op.h' is not visible from
> > > target '//tensorflow/core/kernels/mkl:mkl_sparse_matrix_matmul_op'
> > > ```
> > 
> > 
> > Is it possible to get a commandline reproducer for this error? Can't reproduce locally, so fixing it might be trial-and-error that spams this PR with commits and spams you with emails!
> 
> I'm not sure... I'm surprised bazel doesn't find the dependency issue. Maybe Google's internal tooling is more strict?

All right, trying some things. Thank you for your reviews, by the way! @cantonios Can you get the internal checks running to test the latest version out? > Thank you! Sorry for so many change requests! (It's from our internal checks.)

Hey, I get it -- totally fine! Thanks for bearing with me while I made them. Any update on this? Is the next step in the process for @qqfish to review? Looks like the Copybara checks are passing now. > Any update on this? Is the next step in the process for @qqfish to review? Looks like the Copybara checks are passing now.

It's going through the internal battery of tests now. > > Any update on this? Is the next step in the process for @qqfish to review? Looks like the Copybara checks are passing now.
> 
> It's going through the internal battery of tests now.

Oh, nice, sounds good. Unfortunately I had to rollback this PR since it broke a oneDNN build on arm64. The error was:

```
In file included from tensorflow/core/kernels/mkl/mkl_einsum_op.cc:19:
In file included from tensorflow/core/kernels/mkl/mkl_batch_matmul_helper.h:25:
./tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h:927:38: error: no member named 'csr' in 'dnnl::memory::desc'
  927 |       const auto tmp = memory::desc::csr(
      |                        ~~~~~~~~~~~~~~^
```

@matthew-olson-intel, can you create a new PR with this issue fixed?

 @reedwm Sure. The new PR is in #63030.
Adds `_MklNativeSparseMatrixMatMul` and its accompanying kernel, which uses oneDNN to multiply a CSR sparse matrix by a dense tensor. The op is enabled with an environment variable (`TF_ENABLE_ONEDNN_SPMM`), so is entirely opt-in. It also includes tests and a benchmark, which we've used below to measure its performance against the existing kernel.

The performance looks promising particularly for larger shapes, and is optimized to use the AVX2 and AVX512 ISAs. These results were collected using the new benchmark in `tensorflow/core/kernels/mkl/mkl_sparse_matrix_matmul_op_benchmark.cc` on an Intel Xeon Platinum 8480 with hyperthreading enabled. To minimize NUMA effects, we bound it to the first socket.

Configuration (NNZ_M_K_N) | Eigen Time (ns) | oneDNN Time (ns) | Ratio
-- | -- | -- | --
128_8_512_1 | 15300 | 18598 | 0.82
128_16_512_1 | 14778 | 18551 | 0.80
128_128_512_1 | 18613 | 19165 | 0.97
128_4096_4096_1 | 151525 | 35904 | 4.22
1024_4096_4096_1 | 161593 | 37041 | 4.36
16384_4096_4096_1 | 163055 | 49355 | 3.30
128_8_1024_16 | 17307 | 18420 | 0.94
128_16_1024_16 | 17765 | 18678 | 0.95
128_128_1024_16 | 19247 | 19200 | 1.00
128_4096_4096_128 | 160341 | 140894 | 1.14
128_4096_4096_1024 | 181590 | 156980 | 1.16
1024_8_1024_16 | 24265 | 19502 | 1.24
1024_16_1024_16 | 24223 | 20448 | 1.18
1024_128_1024_16 | 26013 | 20396 | 1.28
1024_4096_4096_128 | 157612 | 139688 | 1.13
1024_4096_4096_1024 | 177549 | 160973 | 1.10
16384_8_1024_16 | 153005 | 36643 | 4.18
16384_16_1024_16 | 152853 | 36597 | 4.18
16384_128_1024_16 | 153600 | 31928 | 4.81
16384_4096_4096_128 | 166061 | 142494 | 1.17
16384_4096_4096_1024 | 244243 | 194536 | 1.26
16384_4096_4096_4096 | 654950 | 536546 | 1.22
100_1_1000000_100 | 18230 | 19991 | 0.91
200_1_2000000_100 | 20509 | 20818 | 0.99
400_1_4000000_100 | 23023 | 21638 | 1.06
400_4_1000000_100 | 22164 | 21349 | 1.04
800_4_2000000_100 | 26180 | 23374 | 1.12
1600_4_4000000_100 | 40538 | 28622 | 1.42
800_8_1000000_100 | 27015 | 23231 | 1.16
1600_8_2000000_100 | 39953 | 27894 | 1.43
3200_8_4000000_100 | 77708 | 42709 | 1.82