Move macros in third_party/mkl_dnn/build_defs.bzl to tsl/mkl/build_defs.bzl

PiperOrigin-RevId: 654171036

tensorflow/BUILD

@@ -1110,7 +1110,6 @@ bzl_library(
         "//tensorflow/core/platform:rules_cc_bzl",
         "//third_party/compute_library:build_defs_bzl",
         "//third_party/llvm_openmp:openmp_bzl",
-        "//third_party/mkl_dnn:build_defs_bzl",
         "@bazel_skylib//lib:new_sets",
         "@bazel_skylib//rules:common_settings",
         "@local_config_cuda//cuda:build_defs_bzl",

tensorflow/opensource_only.files

@@ -275,7 +275,6 @@ tf_staging/third_party/llvm_openmp/expand_cmake_vars:.py
 tf_staging/third_party/llvm_openmp/openmp.bzl:
 tf_staging/third_party/mkl/BUILD:
 tf_staging/third_party/mkl_dnn/LICENSE:
-tf_staging/third_party/mkl_dnn/build_defs.bzl:
 tf_staging/third_party/mkl_dnn/mkldnn_acl.BUILD:
 tf_staging/third_party/mkl_dnn/mkldnn_v1.BUILD:
 tf_staging/third_party/mpi/.gitignore:

tensorflow/tensorflow.bzl

@@ -20,6 +20,9 @@ load(
     "if_enable_mkl",
     "if_mkl",
     "if_mkl_ml",
+    "if_mkldnn_aarch64_acl",
+    "if_mkldnn_aarch64_acl_openmp",
+    "if_mkldnn_openmp",
     "onednn_v3_define",
 )
 
@@ -57,12 +60,6 @@ load(
     "//third_party/llvm_openmp:openmp.bzl",
     "windows_llvm_openmp_linkopts",
 )
-load(
-    "//third_party/mkl_dnn:build_defs.bzl",
-    "if_mkldnn_aarch64_acl",
-    "if_mkldnn_aarch64_acl_openmp",
-    "if_mkldnn_openmp",
-)
 load(
     "@local_config_rocm//rocm:build_defs.bzl",
     "if_rocm",

third_party/mkl_dnn/BUILD

@@ -1,5 +1,3 @@
-load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
-
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
     default_visibility = ["//visibility:public"],
@@ -7,8 +5,3 @@ package(
 )
 
 exports_files(["LICENSE"])
-
-bzl_library(
-    name = "build_defs_bzl",
-    srcs = ["build_defs.bzl"],
-)

third_party/mkl_dnn/build_defs.bzl

@@ -1,34 +0,0 @@
-"""Starlark macros for oneDNN.
-
-if_mkldnn_openmp checks if we are building x86 backend with OpenMP.
-if_mkldnn_aarch64_acl checks if we are building with Arm Compute Library.
-if_mkldnn_aarch64_acl_openmp checks if we are building ACL with OpenMP.
-"""
-
-def if_mkldnn_openmp(if_true, if_false = []):
-    """Returns `if_true` if OpenMP is used with oneDNN.
-
-    Shorthand for select()'ing on whether we're building with
-    oneDNN open source library only with openmp
-
-    Returns a select statement which evaluates to if_true if we're building
-    with oneDNN open source library only with OpenMP. Otherwise, the
-    select statement evaluates to if_false.
-
-    """
-    return select({
-        "@local_xla//xla/tsl/mkl:build_with_mkldnn_openmp": if_true,
-        "//conditions:default": if_false,
-    })
-
-def if_mkldnn_aarch64_acl(if_true, if_false = []):
-    return select({
-        "@local_xla//xla/tsl/mkl:build_with_mkl_aarch64": if_true,
-        "//conditions:default": if_false,
-    })
-
-def if_mkldnn_aarch64_acl_openmp(if_true, if_false = []):
-    return select({
-        "@local_xla//xla/tsl/mkl:build_with_mkl_aarch64_openmp": if_true,
-        "//conditions:default": if_false,
-    })

third_party/mkl_dnn/mkldnn_v1.BUILD

@@ -1,7 +1,6 @@
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
-load("@local_tsl//third_party/mkl_dnn:build_defs.bzl", "if_mkldnn_openmp")
 load("@local_xla//xla/tsl:tsl.bzl", "tf_openmp_copts")
-load("@local_xla//xla/tsl/mkl:build_defs.bzl", "if_mkl", "if_mkl_ml")
+load("@local_xla//xla/tsl/mkl:build_defs.bzl", "if_mkl", "if_mkl_ml", "if_mkldnn_openmp")
 
 exports_files(["LICENSE"])
 

third_party/xla/third_party/tsl/opensource_only.files

@@ -59,7 +59,6 @@ third_party/llvm_openmp/expand_cmake_vars:.py
 third_party/llvm_openmp/openmp.bzl:
 third_party/mkl/BUILD:
 third_party/mkl_dnn/LICENSE:
-third_party/mkl_dnn/build_defs.bzl:
 third_party/mkl_dnn/mkldnn_acl.BUILD:
 third_party/mkl_dnn/mkldnn_v1.BUILD:
 third_party/nccl/BUILD:

third_party/xla/third_party/tsl/third_party/mkl_dnn/BUILD

@@ -1,5 +1,3 @@
-load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
-
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
     default_visibility = ["//visibility:public"],
@@ -10,8 +8,3 @@ exports_files(
     ["LICENSE"],
     visibility = ["//visibility:public"],
 )
-
-bzl_library(
-    name = "build_defs_bzl",
-    srcs = ["build_defs.bzl"],
-)

third_party/xla/third_party/tsl/third_party/mkl_dnn/build_defs.bzl

@@ -1,34 +0,0 @@
-"""Starlark macros for oneDNN.
-
-if_mkldnn_openmp checks if we are building x86 backend with OpenMP.
-if_mkldnn_aarch64_acl checks if we are building with Arm Compute Library.
-if_mkldnn_aarch64_acl_openmp checks if we are building ACL with OpenMP.
-"""
-
-def if_mkldnn_openmp(if_true, if_false = []):
-    """Returns `if_true` if OpenMP is used with oneDNN.
-
-    Shorthand for select()'ing on whether we're building with
-    oneDNN open source library only with openmp
-
-    Returns a select statement which evaluates to if_true if we're building
-    with oneDNN open source library only with OpenMP. Otherwise, the
-    select statement evaluates to if_false.
-
-    """
-    return select({
-        "@local_xla//xla/tsl/mkl:build_with_mkldnn_openmp": if_true,
-        "//conditions:default": if_false,
-    })
-
-def if_mkldnn_aarch64_acl(if_true, if_false = []):
-    return select({
-        "@local_xla//xla/tsl/mkl:build_with_mkl_aarch64": if_true,
-        "//conditions:default": if_false,
-    })
-
-def if_mkldnn_aarch64_acl_openmp(if_true, if_false = []):
-    return select({
-        "@local_xla//xla/tsl/mkl:build_with_mkl_aarch64_openmp": if_true,
-        "//conditions:default": if_false,
-    })

third_party/xla/third_party/tsl/third_party/mkl_dnn/mkldnn_v1.BUILD

@@ -1,7 +1,6 @@
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
-load("@local_tsl//third_party/mkl_dnn:build_defs.bzl", "if_mkldnn_openmp")
 load("@local_xla//xla/tsl:tsl.bzl", "tf_openmp_copts")
-load("@local_xla//xla/tsl/mkl:build_defs.bzl", "if_mkl", "if_mkl_ml")
+load("@local_xla//xla/tsl/mkl:build_defs.bzl", "if_mkl", "if_mkl_ml", "if_mkldnn_openmp")
 
 exports_files(["LICENSE"])
 

third_party/xla/xla/tsl/BUILD

@@ -518,7 +518,6 @@ bzl_library(
         "@local_config_cuda//cuda:build_defs_bzl",
         "@local_config_rocm//rocm:build_defs_bzl",
         "@local_config_tensorrt//:build_defs_bzl",
-        "@local_tsl//third_party/mkl_dnn:build_defs_bzl",
         "@local_tsl//tsl/platform:rules_cc_bzl",
     ],
 )

third_party/xla/xla/tsl/mkl/build_defs.bzl

@@ -1,10 +1,13 @@
-"""Starlark macros for MKL.
+"""Starlark macros for MKL and oneDNN.
 
 if_mkl is a conditional to check if we are building with MKL.
 if_mkl_ml is a conditional to check if we are building with MKL-ML.
 if_mkl_ml_only is a conditional to check for MKL-ML-only (no MKL-DNN) mode.
 if_mkl_lnx_x64 is a conditional to check for MKL
 if_enable_mkl is a conditional to check if building with MKL and MKL is enabled.
+if_mkldnn_openmp checks if we are building x86 backend with OpenMP.
+if_mkldnn_aarch64_acl checks if we are building with Arm Compute Library.
+if_mkldnn_aarch64_acl_openmp checks if we are building ACL with OpenMP.
 
 mkl_repository is a repository rule for creating MKL repository rule that can
 be pointed to either a local folder, or download it from the internet.
@@ -121,6 +124,34 @@ def onednn_v3_define():
         "//conditions:default": [],
     })
 
+def if_mkldnn_openmp(if_true, if_false = []):
+    """Returns `if_true` if OpenMP is used with oneDNN.
+
+    Shorthand for select()'ing on whether we're building with
+    oneDNN open source library only with openmp
+
+    Returns a select statement which evaluates to if_true if we're building
+    with oneDNN open source library only with OpenMP. Otherwise, the
+    select statement evaluates to if_false.
+
+    """
+    return select({
+        "@local_xla//xla/tsl/mkl:build_with_mkldnn_openmp": if_true,
+        "//conditions:default": if_false,
+    })
+
+def if_mkldnn_aarch64_acl(if_true, if_false = []):
+    return select({
+        "@local_xla//xla/tsl/mkl:build_with_mkl_aarch64": if_true,
+        "//conditions:default": if_false,
+    })
+
+def if_mkldnn_aarch64_acl_openmp(if_true, if_false = []):
+    return select({
+        "@local_xla//xla/tsl/mkl:build_with_mkl_aarch64_openmp": if_true,
+        "//conditions:default": if_false,
+    })
+
 def _enable_local_mkl(repository_ctx):
     return _TF_MKL_ROOT in repository_ctx.os.environ
 

third_party/xla/xla/tsl/tsl.bzl

@@ -9,18 +9,15 @@ load(
     "//xla/tsl/mkl:build_defs.bzl",
     "if_enable_mkl",
     "if_mkl",
+    "if_mkldnn_aarch64_acl",
+    "if_mkldnn_aarch64_acl_openmp",
+    "if_mkldnn_openmp",
     "onednn_v3_define",
 )
 load(
     "//third_party/compute_library:build_defs.bzl",
     "if_enable_acl",
 )
-load(
-    "@local_tsl//third_party/mkl_dnn:build_defs.bzl",
-    "if_mkldnn_aarch64_acl",
-    "if_mkldnn_aarch64_acl_openmp",
-    "if_mkldnn_openmp",
-)
 load(
     "@local_config_rocm//rocm:build_defs.bzl",
     "if_rocm",