What are the steps to build Tensorflow with a custom oneDNN library implementation? - tensorflow

Question

I am using a custom oneDNN library implementation which I need Tensorflow (v2.4.0) to build against.
However, I noticed that there are no build options to use a system-provided OneDNN libary when building Tensorflow.
I would like to know what are the steps to support a Tensorflow build that uses a oneDNN library provided by the system.
Some specifics...
The oneDNN library version is 1.6.4 and is already installed in the system (Linux).
This version corresponds to the one Tensforflow uses when compiling with the "--config=mkl_opensource_only" Bazel flag.
I have access to the library source code, but it would be best to use the compiled library.
The target architecture is RISC-V and the OS is Linux.

Answer 1

There is no easy way of telling bazel to link against a custom library, but if you have the modified source directory of oneDNN you can edit tensorflow/workspace.bzl file and replace mkl_dnn_v1 repository definition with a new_local_repository rule to point to your modified source directory. I.e. replace this block:
tf_http_archive(
name = "mkl_dnn_v1",
build_file = clean_dep("//third_party/mkl_dnn:mkldnn_v1.BUILD"),
sha256 = "5369f7b2f0b52b40890da50c0632c3a5d1082d98325d0f2bff125d19d0dcaa1d",
strip_prefix = "oneDNN-1.6.4",
urls = [
"https://storage.googleapis.com/mirror.tensorflow.org/github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
"https://github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
],
)
With something like this:
native.new_local_repository(
name = "mkl_dnn_v1",
build_file = clean_dep("//third_party/mkl_dnn:mkldnn_v1.BUILD"),
path = "/path/to/your/modified/oneDNN/sources",
)
You might also want to modify third_party/mkl_dnn/mkldnn_v1.BUILD file if you have added any new source files.
Addendum:
--config=mkl_opensource_only seems to be broken now, you might have a better luck using just --config=mkl instead.