Merge pull request #47775 from alenik01:aarch64_build_patch

PiperOrigin-RevId: 365049267
Change-Id: I839a9fb319ccc347b50defeaabb1c3f1ec3cfb5c

.bazelrc

@@ -215,7 +215,8 @@ build:mkl_threadpool --define=tensorflow_mkldnn_contraction_kernel=0
 build:mkl_threadpool --define=build_with_mkl_opensource=true
 build:mkl_threadpool -c opt
 
-# Config setting to build with oneDNN for Arm.
+# Config setting to build oneDNN with Compute Library for the Arm Architecture (ACL).
+# This build is for the inference regime only.
 build:mkl_aarch64 --define=build_with_mkl_aarch64=true --define=enable_mkl=true
 build:mkl_aarch64 --define=tensorflow_mkldnn_contraction_kernel=0
 build:mkl_aarch64 --define=build_with_mkl_opensource=true

LICENSE

@@ -201,3 +201,48 @@ Copyright 2019 The TensorFlow Authors.  All rights reserved.
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
+
+MIT License
+
+Copyright (c) 2017-2021 Arm Limited
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+LICENSE
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+   ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+   WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+   DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+   ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+   (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+   LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+   ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+   (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

configure.py

@@ -1454,7 +1454,9 @@ def main():
         'adding "--config=<>" to your build command. See .bazelrc for more '
         'details.')
   config_info_line('mkl', 'Build with MKL support.')
-  config_info_line('mkl_aarch64', 'Build with oneDNN support for Aarch64.')
+  config_info_line(
+      'mkl_aarch64',
+      'Build with oneDNN and Compute Library for the Arm Architecture (ACL).')
   config_info_line('monolithic', 'Config for mostly static monolithic build.')
   config_info_line('numa', 'Build with NUMA support.')
   config_info_line(

tensorflow/core/kernels/mkl/mkl_conv_ops.cc

@@ -68,6 +68,9 @@ struct MklConvFwdParams {
   MklTensorFormat tf_fmt;
   bool native_format;
   string dtypes = string("");
+#ifdef DNNL_AARCH64_USE_ACL
+  void* filter_address = nullptr;
+#endif
   struct PostOpParam {
     string name;
     mkldnn::algorithm alg;
@@ -468,6 +471,9 @@ class MklConvFwdPrimitiveFactory : public MklPrimitiveFactory<float> {
     key_creator.AddAsKey(prefix);
     key_creator.AddAsKey(convFwdDims.src_dims);
     key_creator.AddAsKey(convFwdDims.filter_dims);
+#ifdef DNNL_AARCH64_USE_ACL
+    key_creator.AddAsKey(convFwdDims.filter_address);
+#endif
     key_creator.AddAsKey(convFwdDims.bias_dims);
     key_creator.AddAsKey(convFwdDims.dst_dims);
     key_creator.AddAsKey(convFwdDims.strides);
@@ -777,6 +783,11 @@ class MklConvOp : public OpKernel {
 
       // TODO(mdfaijul): Extend the basic parameters for data types and fusions
       this->ExtendConvFwdParams(context, convFwdDims);
+#ifdef DNNL_AARCH64_USE_ACL
+      // Specifics of ACL: a primitive per constant weights ptr
+      convFwdDims.filter_address = const_cast<void*>(
+          static_cast<const void*>(filter_tensor.flat<Tfilter>().data()));
+#endif
 
       conv_fwd =
           MklConvFwdPrimitiveFactory<Tinput, Tfilter, Tbias, Ttemp_output>::Get(

tensorflow/core/util/mkl_util.h

@@ -1862,6 +1862,12 @@ class FactoryKeyCreator {
     Append(StringPiece(buffer, sizeof(T)));
   }
 
+  // generalisation to handle pointers
+  void AddAsKey(const void* data) {
+    auto buffer = reinterpret_cast<const char*>(&data);
+    Append(StringPiece(buffer, sizeof(data)));
+  }
+
   string GetKey() { return key_; }
 
  private:

tensorflow/opensource_only.files

@@ -57,6 +57,8 @@ tensorflow/third_party/clang_toolchain/cc_configure_clang.bzl
 tensorflow/third_party/clang_toolchain/download_clang.bzl
 tensorflow/third_party/codegen.BUILD
 tensorflow/third_party/common.bzl
+tensorflow/third_party/compute_library/BUILD
+tensorflow/third_party/compute_library/LICENSE
 tensorflow/third_party/coremltools.BUILD
 tensorflow/third_party/cub.BUILD
 tensorflow/third_party/curl.BUILD
@@ -149,6 +151,7 @@ tensorflow/third_party/mkl/build_defs.bzl
 tensorflow/third_party/mkl_dnn/LICENSE
 tensorflow/third_party/mkl_dnn/build_defs.bzl
 tensorflow/third_party/mkl_dnn/mkldnn.BUILD
+tensorflow/third_party/mkl_dnn/mkldnn_acl.BUILD
 tensorflow/third_party/mkl_dnn/mkldnn_v1.BUILD
 tensorflow/third_party/mpi/.gitignore
 tensorflow/third_party/nanopb.BUILD

tensorflow/tensorflow.bzl

@@ -38,6 +38,7 @@ load(
 )
 load(
     "//third_party/mkl_dnn:build_defs.bzl",
+    "if_mkldnn_aarch64_acl",
     "if_mkldnn_openmp",
 )
 load("@bazel_skylib//lib:new_sets.bzl", "sets")
@@ -368,6 +369,7 @@ def tf_copts(
         # optimizations for Intel builds using oneDNN if configured
         if_enable_mkl(["-DENABLE_MKL"]) +
         if_mkldnn_openmp(["-DENABLE_ONEDNN_OPENMP"]) +
+        if_mkldnn_aarch64_acl(["-DENABLE_MKL", "-DENABLE_ONEDNN_OPENMP"]) +
         if_android_arm(["-mfpu=neon"]) +
         if_linux_x86_64(["-msse3"]) +
         if_ios_x86_64(["-msse4.1"]) +

tensorflow/workspace2.bzl

@@ -188,6 +188,29 @@ def _tf_repositories():
         ],
     )
 
+    tf_http_archive(
+        name = "mkl_dnn_acl_compatible",
+        build_file = "//third_party/mkl_dnn:mkldnn_acl.BUILD",
+        sha256 = "5f7fd92e2d0bf83580656695d4404e2cd1390ecad36496fd8ba10b5adc905f70",
+        strip_prefix = "oneDNN-2.1",
+        urls = [
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/oneapi-src/oneDNN/archive/v2.1.tar.gz",
+            "https://github.com/oneapi-src/oneDNN/archive/v2.1.tar.gz",
+        ],
+    )
+
+    tf_http_archive(
+        name = "compute_library",
+        sha256 = "cdb3d8a7ab7ea13f0df207a20657f2827ac631c24aa0e8487bacf97697237bdf",
+        strip_prefix = "ComputeLibrary-21.02",
+        build_file = "//third_party/compute_library:BUILD",
+        patch_file = "//third_party/compute_library:compute_library.patch",
+        urls = [
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/ARM-software/ComputeLibrary/archive/v21.02.tar.gz",
+            "https://github.com/ARM-software/ComputeLibrary/archive/v21.02.tar.gz",
+        ],
+    )
+
     tf_http_archive(
         name = "arm_compiler",
         build_file = "//:arm_compiler.BUILD",

third_party/compute_library/BUILD

@@ -0,0 +1,87 @@
+exports_files(["LICENSE"])
+
+cc_library(
+    name = "include",
+    hdrs = glob([
+        "include/**/*.h",
+        "include/**/*.hpp",
+    ]),
+    includes = ["include"],
+    strip_include_prefix = "include",
+)
+
+cc_library(
+    name = "arm_compute_core",
+    srcs = glob(
+        [
+            "src/core/*.cpp",
+            "src/core/helpers/*.cpp",
+            "src/core/CPP/**/*.cpp",
+            "src/core/utils/**/*.cpp",
+            "src/core/NEON/kernels/**/*.cpp",
+            "src/core/cpu/kernels/*.cpp",
+            "src/core/cpu/kernels/**/*.cpp",
+            "src/core/**/*.hpp",
+            "**/*.h",
+        ],
+        exclude = ["src/core/TracePoint.cpp"],
+    ),
+    hdrs = glob([
+        "arm_compute/core/**/*.h",
+        "**/*.inl",
+    ]) + [
+        "arm_compute_version.embed",
+    ],
+    defines = [
+        "ENABLE_FP16_KERNELS",
+        "ENABLE_FP32_KERNELS",
+        "ENABLE_QASYMM8_KERNELS",
+        "ENABLE_QASYMM8_SIGNED_KERNELS",
+        "ENABLE_QSYMM16_KERNELS",
+        "ENABLE_INTEGER_KERNELS",
+        "ENABLE_NHWC_KERNELS",
+    ],
+    includes = [
+        "src/core/NEON/kernels/assembly",
+        "src/core/NEON/kernels/convolution/common",
+        "src/core/NEON/kernels/convolution/winograd",
+    ],
+    deps = ["include"],
+)
+
+cc_library(
+    name = "arm_compute_runtime",
+    srcs = glob([
+        "src/runtime/*.cpp",
+        "src/runtime/CPP/**/*.cpp",
+        "src/runtime/NEON/**/*.cpp",
+        "src/runtime/cpu/**/*.cpp",
+        "**/*.h",
+    ]),
+    hdrs = glob(["arm_compute/runtime/**/*.h"]) + [
+        "arm_compute_version.embed",
+    ],
+    defines = ["ARM_COMPUTE_CPP_SCHEDULER"],
+    linkopts = ["-lpthread"],
+    visibility = ["//visibility:public"],
+    deps = ["arm_compute_core"],
+)
+
+cc_library(
+    name = "arm_compute_graph",
+    srcs = glob([
+        "src/graph/*.cpp",
+        "src/graph/algorithms/*.cpp",
+        "src/graph/backends/*.cpp",
+        "src/graph/detail/*.cpp",
+        "src/graph/frontend/*.cpp",
+        "src/graph/mutators/*.cpp",
+        "src/graph/nodes/*.cpp",
+        "src/graph/printers/*.cpp",
+        "src/graph/backends/NEON/*.cpp",
+        "**/*.h",
+    ]),
+    hdrs = glob(["arm_compute/graph/**/*.h"]),
+    visibility = ["//visibility:public"],
+    deps = ["arm_compute_core"],
+)

third_party/compute_library/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2017-2021 Arm Limited
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

third_party/compute_library/compute_library.patch

@@ -0,0 +1,8 @@
+diff --git a/arm_compute_version.embed b/arm_compute_version.embed
+new file mode 100644
+index 000000000..c986ad52a
+--- /dev/null
++++ b/arm_compute_version.embed
+@@ -0,0 +1,1 @@
++"arm_compute_version=v21.02 Build options: {} Git hash=b'N/A'"
+\ No newline at end of file

third_party/mkl/build_defs.bzl

@@ -32,6 +32,7 @@ def if_mkl(if_true, if_false = []):
       may need it. It may be deleted in future with refactoring.
     """
     return select({
+        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": if_true,
         "//tensorflow:linux_x86_64": if_true,
         "//tensorflow:windows": if_true,
         "//conditions:default": if_false,
@@ -100,7 +101,7 @@ def mkl_deps():
       inclusion in the deps attribute of rules.
     """
     return select({
-        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": ["@mkl_dnn_v1//:mkl_dnn_aarch64"],
+        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": ["@mkl_dnn_acl_compatible//:mkl_dnn_acl"],
         "//tensorflow:linux_x86_64": ["@mkl_dnn_v1//:mkl_dnn"],
         "//tensorflow:windows": ["@mkl_dnn_v1//:mkl_dnn"],
         "//conditions:default": [],

third_party/mkl_dnn/BUILD

@@ -27,6 +27,14 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+config_setting(
+    name = "build_with_mkl_aarch64",
+    define_values = {
+        "build_with_mkl_aarch64": "true",
+    },
+    visibility = ["//visibility:public"],
+)
+
 bzl_library(
     name = "build_defs_bzl",
     srcs = ["build_defs.bzl"],

third_party/mkl_dnn/build_defs.bzl

@@ -29,3 +29,9 @@ def if_mkldnn_openmp(if_true, if_false = []):
         "@org_tensorflow//third_party/mkl_dnn:build_with_mkldnn_openmp": if_true,
         "//conditions:default": if_false,
     })
+
+def if_mkldnn_aarch64_acl(if_true, if_false = []):
+    return select({
+        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": if_true,
+        "//conditions:default": if_false,
+    })

third_party/mkl_dnn/mkldnn_acl.BUILD

@@ -0,0 +1,70 @@
+exports_files(["LICENSE"])
+
+load(
+    "@org_tensorflow//third_party:common.bzl",
+    "template_rule",
+)
+
+_DNNL_RUNTIME_OMP = {
+    "#cmakedefine DNNL_CPU_THREADING_RUNTIME DNNL_RUNTIME_${DNNL_CPU_THREADING_RUNTIME}": "#define DNNL_CPU_THREADING_RUNTIME DNNL_RUNTIME_OMP",
+    "#cmakedefine DNNL_CPU_RUNTIME DNNL_RUNTIME_${DNNL_CPU_RUNTIME}": "#define DNNL_CPU_RUNTIME DNNL_RUNTIME_OMP",
+    "#cmakedefine DNNL_GPU_RUNTIME DNNL_RUNTIME_${DNNL_GPU_RUNTIME}": "#define DNNL_GPU_RUNTIME DNNL_RUNTIME_NONE",
+    "#cmakedefine DNNL_WITH_SYCL": "/* #undef DNNL_WITH_SYCL */",
+    "#cmakedefine DNNL_WITH_LEVEL_ZERO": "/* #undef DNNL_WITH_LEVEL_ZERO */",
+    "#cmakedefine DNNL_SYCL_CUDA": "/* #undef DNNL_SYCL_CUDA */",
+}
+
+template_rule(
+    name = "dnnl_config_h",
+    src = "include/oneapi/dnnl/dnnl_config.h.in",
+    out = "include/oneapi/dnnl/dnnl_config.h",
+    substitutions = _DNNL_RUNTIME_OMP,
+)
+
+template_rule(
+    name = "dnnl_version_h",
+    src = "include/oneapi/dnnl/dnnl_version.h.in",
+    out = "include/oneapi/dnnl/dnnl_version.h",
+    substitutions = {
+        "@DNNL_VERSION_MAJOR@": "2",
+        "@DNNL_VERSION_MINOR@": "1",
+        "@DNNL_VERSION_PATCH@": "0",
+        "@DNNL_VERSION_HASH@": "fbdfeea2642fec05387ed37d565cf904042f507e",
+    },
+)
+
+cc_library(
+    name = "mkl_dnn_acl",
+    srcs = glob(
+        [
+            "src/common/*.cpp",
+            "src/common/*.hpp",
+            "src/cpu/**/*.cpp",
+            "src/cpu/**/*.hpp",
+        ],
+        exclude = ["src/cpu/x64/**/*"],
+    ) + [
+        ":dnnl_config_h",
+        ":dnnl_version_h",
+    ],
+    hdrs = glob(["include/*"]),
+    copts = [
+        "-fexceptions",
+        "-UUSE_MKL",
+        "-UUSE_CBLAS",
+    ],
+    defines = ["DNNL_AARCH64_USE_ACL=1"],
+    includes = [
+        "include",
+        "src",
+        "src/common",
+        "src/cpu",
+        "src/cpu/gemm",
+    ],
+    linkopts = ["-lgomp"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@compute_library//:arm_compute_graph",
+        "@compute_library//:arm_compute_runtime",
+    ],
+)

third_party/mkl_dnn/mkldnn_v1.BUILD

@@ -141,36 +141,3 @@ cc_library(
         [],
     ),
 )
-
-cc_library(
-    name = "mkl_dnn_aarch64",
-    srcs = glob([
-        "src/common/*.cpp",
-        "src/common/*.hpp",
-        "src/cpu/*.cpp",
-        "src/cpu/*.hpp",
-        "src/cpu/rnn/*.cpp",
-        "src/cpu/rnn/*.hpp",
-        "src/cpu/matmul/*.cpp",
-        "src/cpu/matmul/*.hpp",
-        "src/cpu/gemm/**/*",
-    ]) + [
-        ":dnnl_config_h",
-        ":dnnl_version_h",
-    ],
-    hdrs = glob(["include/*"]),
-    copts = [
-        "-fexceptions",
-        "-UUSE_MKL",
-        "-UUSE_CBLAS",
-    ],
-    includes = [
-        "include",
-        "src",
-        "src/common",
-        "src/cpu",
-        "src/cpu/gemm",
-    ],
-    linkopts = ["-lgomp"],
-    visibility = ["//visibility:public"],
-)