Reduce MKL overheads on small shapes by not rewriting node to use MKL
Hi @milpuz01 Can you please resolve conflicts? Thank you! > Some more comments. I'll post more tomorrow. Sorry for the delay!

Thanks @penpornk. I have tried to address all your comments in https://github.com/tensorflow/tensorflow/pull/60026/commits/450f4feac896baa6b7d75d82b02973f73cf990b6 @TensorFlow-MKL Any feedback on this PR? Thank you! > @TensorFlow-MKL Any feedback on this PR? Thank you!

@agramesh1, @mahmoud-abuzaina, @ashiqimranintel, @sachinmuradi  do you maybe have any comments about the PR? Many thanks. Hi, we have couple main concerns :
1. Moving the fusion to mkl_layout_pass is not a good idea. From design perspective, we should preferably do all fusions in remapper with proper #ifdef INTEL_MKL and  IsMKLEnabled() checks. 
2. oneDNN would be the right place to handle optimizations for processor (or instruction set) variations instead of adding in mkl_layout_pass.

Both these changes applied for other fusions and primitives, can easily explode mkl_layout_pass code.  Hi @gaurides, thank you very much for your reply and the concerns that you have raised. I have couple of questions/comments that I have put below:

> 1. Moving the fusion to mkl_layout_pass is not a good idea. From design perspective, we should preferably do all fusions in remapper with proper #ifdef INTEL_MKL and  IsMKLEnabled() checks.

I can see that there are few other fusions already present in `mkl_layout_pass` such as merging convolution with bias (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/mkl_layout_pass.cc#L3048). Is there expectations that those fusions will also move to `remapper`?

> Both these changes applied for other fusions and primitives, can easily explode mkl_layout_pass code.

I think the same argument could be applied to `remapper` too as adding fusions there it will also make it explode in terms of code size? What influences decisions where graph rewriting pass should be? My impression was that lowering op to oneDNN should be in `mkl_layout_pass`, while fusion and other (high-level?) optimisations should be in `remapper` that might benefit other low-level libraries (such as Eigen).

> 2. oneDNN would be the right place to handle optimizations for processor (or instruction set) variations instead of adding in mkl_layout_pass.

The optimisation that we are doing with heuristic here is that we are choosing between lowering operation to be executed using oneDNN (i.e. rewriting it to be _Mkl) or executing operation using Eigen library so that is not possible to do in oneDNN as that decision is made before a call is made to oneDNN.

@penpornk, do you think that it might be introducing a new layout pass for AArch64 that would be called instead of `mkl_layout_pass`. We are already seeing some divergence in terms of fusion or lowering ops to onedNN such as here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/grappler/optimizers/remapper.cc#L4603 and https://github.com/tensorflow/tensorflow/pull/60723 

Thanks,
Milos. Hi @milpuz01, right, you will see some fusions currently happening in mkl_layout_pass but they are very old (5-6 years old). They need to be moved to remapper whenever possible. All new fusions have been added to remapper. We also added a pattern_matcher in remapper, which is easy and convenient to define patterns. take a [look here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/grappler/optimizers/remapper.cc#L1444) . It's easier to match patterns than with code. So, you can try using it to do all fusions in grappler.
Using IsMKLEnabled() and #ifdef in remapper seems ok. You can probably define another functions that does the conditions checks to fuse or not. So, I recommend keeping fusions in remapper.
With respect to the logic to decide whether to use Mkl* version of the op or use Eigen, it's ok if you want to keep it in mkl_layout_pass as long as it doesn't affect Intel CPU code. But based on multiple HWs arch, it is easy to see it get complicated.  We typically try to use generic conditions for shape/size.  @milpuz01 Could you please add some unit tests. Hi @milpuz01 Any update on this PR? Please. Thank you! >  So, you can try using it to do all fusions in grappler. Using IsMKLEnabled() and #ifdef in remapper seems ok. You can probably define another functions that does the conditions checks to fuse or not. So, I recommend keeping fusions in remapper.

Thanks @gaurides for in-depth explanation. I have now reverted the change and put back fusion back that was removed from the mapper. In order to be able to decide whether to fuse or not on AArch64 I have added new function as per your suggestion in a new header file `mkl_heuristics.h` that can check now whether we want to fuse or not and use it in remapper. The same function is also used in layout pass too when decision needs to be made there. 

> With respect to the logic to decide whether to use Mkl* version of the op or use Eigen, it's ok if you want to keep it in mkl_layout_pass as long as it doesn't affect Intel CPU code. 

All the code that we are added is either guarded by `DNNL_AARCH64_USE_ACL` or reverts to default behaviour so it doesn't affect Intel's CPU code.
 > @milpuz01 Could you please add some unit tests.

Thanks @mdfaijul for the comment. I have added new tests in `mkl_herustics_test.cc` to test method for calculating FLOPS and threshold. The rest of changes do not introduce any new behaviour that is not covered already by existing tests that are in tests for `remapper_test.cc` and `mkl_remapper_test.cc`. Furthermore, in this change PR: https://github.com/tensorflow/tensorflow/pull/60160 we have extended convolution benchmark to take into account rewriting of nodes using layout pass so that tests that we get best performance when running convolutions.  > Could you please help fix [PY+CPP Ubuntu CPU](https://source.cloud.google.com/results/invocations/c09a2a5f-e1e3-410f-8b8f-d2b79592ea1d/log) build failures?

Sorry about this @penporn. I believe that commit [a1f135e](https://github.com/tensorflow/tensorflow/pull/60026/commits/a1f135efcaa05abb86ad7b1f5f91258cf9ae5e60) should fix the failures as ordering of included headers was not correct in `mkl_heuristics_test.cc` Could you please help take a look at `//tensorflow/core/common_runtime/eager:context_test` failure?
```
==================== Test output for //tensorflow/core/common_runtime/eager:context_test:
2023-07-25 12:21:47.692537: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-07-25 12:21:47.692765: W tensorflow/tsl/lib/monitoring/collection_registry.cc:81] Trying to register 2 metrics with the same name: /tensorflow/core/bfc_allocator_delay. The old value will be erased in order to register a new one. Please check if you link the metric more than once, or if the name is already used by other metrics.
[==========] Running 10 tests from 1 test suite.
[----------] Global test environment set-up.
[----------] 10 tests from EagerContextTest
[ RUN      ] EagerContextTest.CompositeDevice
[       OK ] EagerContextTest.CompositeDevice (25 ms)
[ RUN      ] EagerContextTest.CompositeDeviceWithGivenName
[       OK ] EagerContextTest.CompositeDeviceWithGivenName (0 ms)
[ RUN      ] EagerContextTest.AddFunctionDef
[       OK ] EagerContextTest.AddFunctionDef (26 ms)
[ RUN      ] EagerContextTest.AddFunctionDefRepeatSame
[       OK ] EagerContextTest.AddFunctionDefRepeatSame (0 ms)
[ RUN      ] EagerContextTest.AddFunctionDefRepeatDifferent
[       OK ] EagerContextTest.AddFunctionDefRepeatDifferent (1 ms)
[ RUN      ] EagerContextTest.FunctionErrorRecovery
2023-07-25 12:21:47.827207: F ./tensorflow/core/framework/device_base.h:146] Check failed: cpu_worker_threads_ != nullptr 
*** Received signal 6 ***
*** BEGIN MANGLED STACK TRACE ***
```
I believe this comes from the changes in //tensorflow/core/common_runtime/optimize_function_graph_utils.cc.
https://github.com/tensorflow/tensorflow/pull/60026/files#diff-cf5de83fc2a1e4038beb6f3994daae6bb9bc321f8eb05a6e125980ab4b598680R556-R562 > Could you please help take a look at `//tensorflow/core/common_runtime/eager:context_test` failure?

Thanks @penpornk. Commit 2701cc4 should make the test to pass. I have restricted to only check for number of threads for intra op if it is local device because that is common path that is used when doing MKL layout pass. I have also moved to set it inside of `PreprocessAndPartitionGraph()` instead of `OptimizeFunctionGraph()` as MKL layout pass is called from former.
We observed that when you run networks that have layers with smaller input and overall lower compute density overhead into MKL path dominates useful compute so in that case it is better to call default, Eigen, path.
 
This patch introduces changes to MKL layout pass where a liner analytical model is used to decide based on compute density and number of threads to either rewrite node to use MKL code path or leave it as it is. It propagates input shapes of the operation by decorating node with _input_shapes attribute number of threads used to parallelise operation by reading it from device on which operation is expected to execute (in this case CPUDevice) which, are then retrived in MKL layout pass and passed to the analytical model.
 
With this approach we have observed speedup for models available on TF-Hub as shown in the two attached graphs when running with 8 and 16 cores on Neoverse-V1 platform. Before shows results without this patch and after shows results with the patch relative when same models are ran with TF_ENABLE_ONEDNN_OPTS=0. Overall for 8 cores we see ~1.2x average speedup over TF_ENABLE_ONEDNN_OPTS=0 for 8 cores and ~1.12x for 16 cores.

![8cores_results](https://user-images.githubusercontent.com/79916358/225777461-0962b80c-4923-4b50-bb55-b65874ee3da9.png)

![16cores_results](https://user-images.githubusercontent.com/79916358/225777470-3fea315f-d78c-447b-9bbc-4f09202400da.png)

In addition to analytical model improvements for convolution, fused convolution, fused batch normalisation and swish, this patch also moves code for merging sigmoid and multiplication from remapper.cc to mkl_layout_pass.cc as it is MKL merge operation.

cc: @penpornk @cantonios 