Enable windows build with compiled llvm_openmp source

tensorflow/compiler/mlir/lite/BUILD

@@ -10,6 +10,12 @@ load(
     "//third_party/mlir:tblgen.bzl",
     "gentbl",
 )
+load("//tensorflow:tensorflow.bzl", "if_windows")
+load(
+    "@org_tensorflow//third_party/llvm_openmp:openmp.bzl",
+    "windows_llvm_openmp_deps",
+    "windows_llvm_openmp_linkopts",
+)
 
 package(
     default_visibility = [
@@ -634,7 +640,7 @@ tf_native_cc_binary(
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:TableGen",
         "@llvm-project//mlir:TableGen",
-    ],
+    ] + if_windows(windows_llvm_openmp_deps()),
 )
 
 gentbl(

tensorflow/compiler/mlir/lite/quantization/BUILD

@@ -13,6 +13,12 @@ load(
     "//third_party/mlir:tblgen.bzl",
     "gentbl",
 )
+load("//tensorflow:tensorflow.bzl", "if_windows")
+load(
+    "@org_tensorflow//third_party/llvm_openmp:openmp.bzl",
+    "windows_llvm_openmp_deps",
+    "windows_llvm_openmp_linkopts",
+)
 
 package(
     default_visibility = [
@@ -148,7 +154,7 @@ tf_native_cc_binary(
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:TableGen",
         "@llvm-project//mlir:TableGen",
-    ],
+    ] + if_windows(windows_llvm_openmp_deps()),
 )
 
 cc_library(

tensorflow/compiler/mlir/xla/BUILD

@@ -2,6 +2,12 @@ load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load("//third_party/mlir:tblgen.bzl", "gentbl")
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test", "tf_native_cc_binary")
 load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
+load("//tensorflow:tensorflow.bzl", "if_windows")
+load(
+    "@org_tensorflow//third_party/llvm_openmp:openmp.bzl",
+    "windows_llvm_openmp_deps",
+    "windows_llvm_openmp_linkopts",
+)
 
 package(
     default_visibility = [":friends"],
@@ -481,7 +487,7 @@ tf_native_cc_binary(
         "@llvm-project//llvm:TableGen",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TableGen",
-    ],
+    ] + if_windows(windows_llvm_openmp_deps()),
 )
 
 gentbl(

tensorflow/core/common_runtime/BUILD

@@ -40,6 +40,12 @@ load(
     "//third_party/mkl:build_defs.bzl",
     "if_mkl",
 )
+load("//tensorflow:tensorflow.bzl", "if_windows")
+load(
+    "@org_tensorflow//third_party/llvm_openmp:openmp.bzl",
+    "windows_llvm_openmp_deps",
+    "windows_llvm_openmp_linkopts",
+)
 
 default_package_visibility = [
     "//tensorflow:internal",
@@ -90,7 +96,8 @@ cc_library(
         ":core_cpu",
         "//tensorflow/core/common_runtime/gpu:gpu_runtime",
         "//tensorflow/core/common_runtime/pluggable_device:pluggable_device_runtime",
-    ] + if_libtpu(["//tensorflow/core/tpu:tpu_runtime"]),
+    ] + if_libtpu(["//tensorflow/core/tpu:tpu_runtime"])
+    + if_windows(windows_llvm_openmp_deps()),
 )
 
 filegroup(
@@ -1213,7 +1220,11 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
-    ],
+    ] + if_windows(windows_llvm_openmp_deps()),
+    linkopts = select({
+        "//tensorflow:windows": [windows_llvm_openmp_linkopts()],
+        "//conditions:default": [],
+        }),
 )
 
 cc_library(
@@ -1543,7 +1554,11 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
-    ] + if_mkl([":mkl_cpu_allocator"]),
+    ] + if_mkl([":mkl_cpu_allocator"]) + if_windows(windows_llvm_openmp_deps()),
+    linkopts = select({
+        "//tensorflow:windows": [windows_llvm_openmp_linkopts()],
+        "//conditions:default": [],
+        }),
 )
 
 cc_library(

tensorflow/tensorflow.bzl

@@ -41,6 +41,11 @@ load(
     "if_mkldnn_aarch64_acl",
     "if_mkldnn_openmp",
 )
+load(
+    "//third_party/llvm_openmp:openmp.bzl",
+    "windows_llvm_openmp_deps",
+    "windows_llvm_openmp_linkopts",
+)
 load("@bazel_skylib//lib:new_sets.bzl", "sets")
 load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 
@@ -397,7 +402,7 @@ def tf_openmp_copts():
         # "//third_party/mkl:build_with_mkl_windows_openmp": ["/openmp"],
         # copybara:uncomment_end_and_comment_begin
         "@org_tensorflow//third_party/mkl:build_with_mkl_lnx_openmp": ["-fopenmp"],
-        "@org_tensorflow//third_party/mkl:build_with_mkl_windows_openmp": ["/openmp"],
+        "@org_tensorflow//third_party/mkl:build_with_mkl_windows_openmp": ["/openmp:llvm"],
         # copybara:comment_end
         "//conditions:default": [],
     })
@@ -762,7 +767,7 @@ def tf_native_cc_binary(
         name = name,
         copts = copts,
         linkopts = select({
-            clean_dep("//tensorflow:windows"): [],
+            clean_dep("//tensorflow:windows"): [windows_llvm_openmp_linkopts()],
             clean_dep("//tensorflow:macos"): [
                 "-lm",
             ],
@@ -1097,7 +1102,7 @@ def tf_cc_test(
             clean_dep("//tensorflow:android"): [
                 "-pie",
             ],
-            clean_dep("//tensorflow:windows"): [],
+            clean_dep("//tensorflow:windows"): [windows_llvm_openmp_linkopts()],
             clean_dep("//tensorflow:macos"): [
                 "-lm",
             ],
@@ -2028,14 +2033,14 @@ def pywrap_tensorflow_macro(
             "-Wl,-w",
             "-Wl,-exported_symbols_list,$(location %s.lds)" % vscriptname,
         ],
-        clean_dep("//tensorflow:windows"): [],
+        clean_dep("//tensorflow:windows"): [windows_llvm_openmp_linkopts()],
         "//conditions:default": [
             "-Wl,--version-script",
             "$(location %s.lds)" % vscriptname,
         ],
     })
     extra_deps += select({
-        clean_dep("//tensorflow:windows"): [],
+        clean_dep("//tensorflow:windows"): windows_llvm_openmp_deps(),
         "//conditions:default": [
             "%s.lds" % vscriptname,
         ],
@@ -2763,7 +2768,8 @@ def tf_python_pybind_extension(
         features = features,
         copts = copts,
         hdrs = hdrs,
-        deps = deps + tf_binary_pybind_deps() + if_mkl_ml(["//third_party/mkl:intel_binary_blob"]),
+        deps = deps + tf_binary_pybind_deps() + if_mkl_ml(["//third_party/mkl:intel_binary_blob"])
+            + if_windows(windows_llvm_openmp_deps()),
         defines = defines,
         visibility = visibility,
         link_in_framework = True,

tensorflow/tools/pip_package/BUILD

@@ -6,6 +6,11 @@ load("//third_party/mkl:build_defs.bzl", "if_enable_mkl", "if_mkl", "if_mkl_ml")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load("@local_config_syslibs//:build_defs.bzl", "if_not_system_lib")
 load("//tensorflow/core/platform:build_config_root.bzl", "tf_additional_license_deps")
+load(
+    "@org_tensorflow//third_party/llvm_openmp:openmp.bzl",
+    "windows_llvm_openmp_deps",
+    "windows_llvm_openmp_linkopts",
+)
 
 package(default_visibility = ["//visibility:private"])
 
@@ -296,7 +301,8 @@ sh_binary(
            select({
                "//tensorflow:dynamic_loaded_kernels": DYNAMIC_LOADED_KERNELS,
                "//conditions:default": [],
-           }) + if_mkl_ml(["//third_party/mkl:intel_binary_blob"]),
+           }) + if_mkl_ml(["//third_party/mkl:intel_binary_blob"])
+           + if_windows(windows_llvm_openmp_deps()),
 )
 
 # A genrule for generating a marker file for the pip package on Windows

tensorflow/tools/pip_package/build_pip_package.sh

@@ -145,6 +145,8 @@ function prepare_src() {
       bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles/org_tensorflow \
       "${XLA_AOT_RUNTIME_SOURCES}/"
     RUNFILES=bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles/org_tensorflow
+    cp bazel-bin/external/llvm_openmp/libiomp5md.dll ${TMPDIR}/tensorflow/python
+    cp bazel-bin/external/llvm_openmp/libiomp5md.dll.if.lib ${TMPDIR}/tensorflow/python
   else
     RUNFILES=bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow
     if [ -d bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/external ]; then

third_party/llvm_openmp/BUILD

@@ -1,25 +1,33 @@
 # Build file for OpenMP library that is part of llvm
 
+load("@org_tensorflow//tensorflow:tensorflow.bzl", "if_windows", "if_linux_x86_64", "if_macos")
+
 load(
     "@org_tensorflow//third_party/llvm:llvm.bzl",
     "cmake_var_string",
     "expand_cmake_vars",
 )
 load(
     "@org_tensorflow//third_party/llvm_openmp:openmp.bzl",
+    "libiomp5_cc_binary",
     "dict_add",
 )
 
 exports_files(["LICENSE.txt"])
 
+kmp_i18n_os_type = select({
+    "@org_tensorflow//tensorflow:windows": "win",
+    "//conditions:default": "lin",
+})
+
 genrule(
     name = "kmp_i18n_id",
     srcs = [
         "runtime/tools/message-converter.pl",
         "runtime/src/i18n/en_US.txt",
     ],
     outs = ["include/kmp_i18n_id.inc"],
-    cmd = "perl $(location runtime/tools/message-converter.pl) --os=lin --prefix=kmp_i18n --enum=$@  $(location runtime/src/i18n/en_US.txt)",
+    cmd = "perl $(location runtime/tools/message-converter.pl) --os=" + kmp_i18n_os_type + " --prefix=kmp_i18n --enum=$@  $(location runtime/src/i18n/en_US.txt)",
 )
 
 genrule(
@@ -29,7 +37,16 @@ genrule(
         "runtime/src/i18n/en_US.txt",
     ],
     outs = ["include/kmp_i18n_default.inc"],
-    cmd = "perl $(location runtime/tools/message-converter.pl) --os=lin --prefix=kmp_i18n --default=$@ $(location runtime/src/i18n/en_US.txt)",
+    cmd = "perl $(location runtime/tools/message-converter.pl) --os=" + kmp_i18n_os_type + " --prefix=kmp_i18n --default=$@ $(location runtime/src/i18n/en_US.txt)",
+)
+
+genrule(
+    name = "generate_def",
+    srcs = [
+        "runtime/tools/generate-def.pl",
+    ],
+    outs = ["omp_dll.def"],
+    cmd = "perl $(location runtime/tools/generate-def.pl) -D arch_32e -D msvc_compat -D norm -o omp_dll.tmp external/llvm_openmp/runtime/src/dllexports; mv omp_dll.tmp $@",
 )
 
 # Bazel doesn't accept .txt as an input, rename the ldscript to .inc to workaround.
@@ -158,48 +175,8 @@ common_includes = [
 # ......
 # ......
 
-cc_binary(
-    name = "libiomp5.so",
-    srcs = cppsources + [
-        #linux specific files
-        "runtime/src/z_Linux_util.cpp",
-        "runtime/src/kmp_gsupport.cpp",
-        "runtime/src/z_Linux_asm.S",
-    ] + srcdeps,
-    copts = ["-Domp_EXPORTS -D_GNU_SOURCE -D_REENTRANT"],
-    includes = common_includes,
-    linkopts = ["-lpthread -ldl -Wl,--version-script=$(location :ldscript)"],
-    linkshared = True,
-    visibility = ["//visibility:public"],
-)
 
-cc_binary(
-    name = "libiomp5md.dll",
-    srcs = cppsources + [
-        #window specific files
-        "runtime/src/z_Windows_NT_util.cpp",
-        "runtime/src/z_Windows_NT-586_util.cpp",
-    ] + srcdeps + [":openmp_asm"],
-    copts = ["/Domp_EXPORTS /D_M_AMD64 /DOMPT_SUPPORT=0 /D_WINDOWS /D_WINNT /D_USRDLL"],
-    includes = common_includes,
-    linkopts = ["/MACHINE:X64"],
-    linkshared = True,
-    visibility = ["//visibility:public"],
-)
+if_linux_x86_64( libiomp5_cc_binary("libiomp5.so", cppsources, srcdeps, common_includes) )
+if_macos( libiomp5_cc_binary("libiomp5.dylib", cppsources, srcdeps, common_includes) )
+if_windows( libiomp5_cc_binary("libiomp5md.dll", cppsources, srcdeps, common_includes) )
 
-# MacOS build has not been tested, however since the MacOS build of openmp
-# uses the same configuration as Linux, the following should work.
-cc_binary(
-    name = "libiomp5.dylib",
-    srcs = cppsources + [
-        #linux/MacOS specific files
-        "runtime/src/z_Linux_util.cpp",
-        "runtime/src/kmp_gsupport.cpp",
-        "runtime/src/z_Linux_asm.S",
-    ] + srcdeps,
-    copts = ["-Domp_EXPORTS -D_GNU_SOURCE -D_REENTRANT"],
-    includes = common_includes,
-    linkopts = ["-lpthread -ldl -Wl,--version-script=$(location :ldscript)"],
-    linkshared = True,
-    visibility = ["//visibility:public"],
-)

third_party/llvm_openmp/openmp.bzl

@@ -3,6 +3,22 @@ TODO(Intel-tf): Delete this and reuse a similar function in third_party/llvm
 after the TF 2.4 branch cut has passed.
 """
 
+load(
+    "//tensorflow/core/platform:rules_cc.bzl",
+    "cc_binary",
+)
+
+#WINDOWS_MSVC_LLVM_OPENMP_LIBPATH = "bazel-out/x64_windows-opt/bin/external/llvm_openmp/libiomp5md.dll.if.lib"
+WINDOWS_MSVC_LLVM_OPENMP_LIBPATH = "$(BINDIR)/external/llvm_openmp/libiomp5md.dll.if.lib"
+WINDOWS_MSVC_LLVM_OPENMP_LINKOPTS = "/NODEFAULTLIB:libomp /DEFAULTLIB:" + WINDOWS_MSVC_LLVM_OPENMP_LIBPATH
+
+def windows_llvm_openmp_deps():
+    return ["@llvm_openmp//:libiomp5md.dll"]
+
+
+def windows_llvm_openmp_linkopts():
+    return WINDOWS_MSVC_LLVM_OPENMP_LINKOPTS
+
 def dict_add(*dictionaries):
     """Returns a new `dict` that has all the entries of the given dictionaries.
 
@@ -19,3 +35,58 @@ def dict_add(*dictionaries):
     for d in dictionaries:
         result.update(d)
     return result
+
+def select_os_specific(L, M, W):
+    return select({
+        "@org_tensorflow//tensorflow:linux_x86_64": L,
+        "@org_tensorflow//tensorflow:macos": M,
+        "@org_tensorflow//tensorflow:windows": W,
+        "//conditions:default": L,
+    })
+
+def libname_os_specific():
+    return "" + select_os_specific(L = "libiomp5.so", M = "libiomp5.dylib", W = "libiomp5md.dll")
+
+def libiomp5_cc_binary(libname, cppsources, srcdeps, common_includes):
+    cc_binary(
+        name = libname,
+        srcs = cppsources + srcdeps +
+            select_os_specific(
+            L = [
+                #linux & macos specific files
+                "runtime/src/z_Linux_util.cpp",
+                "runtime/src/kmp_gsupport.cpp",
+                "runtime/src/z_Linux_asm.S",
+            ],
+            M = [
+                #linux & macos specific files
+                "runtime/src/z_Linux_util.cpp",
+                "runtime/src/kmp_gsupport.cpp",
+                "runtime/src/z_Linux_asm.S",
+            ],
+            W = [
+                #window specific files
+                "runtime/src/z_Windows_NT_util.cpp",
+                "runtime/src/z_Windows_NT-586_util.cpp",
+                ":openmp_asm",
+            ]),
+        copts = select_os_specific(
+            L = ["-Domp_EXPORTS -D_GNU_SOURCE -D_REENTRANT"],
+            M = ["-Domp_EXPORTS -D_GNU_SOURCE -D_REENTRANT"],
+            W = ["/Domp_EXPORTS /D_M_AMD64 /DOMPT_SUPPORT=0 /D_WINDOWS /D_WINNT /D_USRDLL"],
+        ),
+        includes = common_includes,
+        linkopts = select_os_specific(
+            L = ["-lpthread -ldl -Wl,--version-script=$(location :ldscript)"],
+            M = ["-lpthread -ldl -Wl,--version-script=$(location :ldscript)"],
+            W = ["/MACHINE:X64"],
+        ),
+        linkshared = True,
+        additional_linker_inputs = [":generate_def"],
+        win_def_file = ":generate_def",
+        visibility = ["//visibility:public"],
+    )
+
+
+
+