Roll forward with fix for Windows & PPC.

PiperOrigin-RevId: 308975869
Change-Id: I22abc77396e8d2a9c8f6a0907bf9d1d22169cc93

tensorflow/opensource_only.files

@@ -263,6 +263,7 @@ tensorflow/third_party/toolchains/remote_config/containers.bzl
 tensorflow/third_party/toolchains/remote_config/rbe_config.bzl
 tensorflow/third_party/wrapt.BUILD
 tensorflow/third_party/zlib.BUILD
+tensorflow/tools/build_info/BUILD
 tensorflow/tools/ci_build/release/common.sh
 tensorflow/tools/ci_build/release/common_win.bat
 tensorflow/tools/ci_build/release/macos/cpu_libtensorflow/build.sh
@@ -393,6 +394,7 @@ tensorflow/tools/def_file_filter/BUILD
 tensorflow/tools/def_file_filter/BUILD.tpl
 tensorflow/tools/def_file_filter/def_file_filter.py.tpl
 tensorflow/tools/def_file_filter/def_file_filter_configure.bzl
+tensorflow/tools/git/BUILD
 tensorflow/tools/lib_package/BUILD
 tensorflow/tools/lib_package/LibTensorFlowTest.java
 tensorflow/tools/lib_package/README.md

tensorflow/python/BUILD

@@ -4,7 +4,7 @@
 #  ":platform" - Low-level and platform-specific Python code.
 
 load("//tensorflow:tensorflow.bzl", "py_strict_library")
-load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cuda_library", "tf_gen_op_wrapper_py", "tf_py_build_info_genrule", "tf_py_test")
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cuda_library", "tf_gen_op_wrapper_py", "tf_py_test")
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_python_pybind_extension")
@@ -26,6 +26,9 @@ load("//tensorflow:tensorflow.bzl", "tf_external_workspace_visible")
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_pybind_cc_library_wrapper")
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_py_build_info_genrule")
 load("//tensorflow/core/platform:build_config.bzl", "pyx_library", "tf_additional_all_protos", "tf_additional_lib_deps", "tf_proto_library", "tf_proto_library_py", "tf_protos_grappler")  # @unused
 load("//tensorflow/core/platform:build_config_root.bzl", "if_static", "tf_additional_plugin_deps", "tf_additional_xla_deps_py")
 load("//tensorflow/python:build_defs.bzl", "tf_gen_op_wrapper_private_py")
@@ -231,6 +234,8 @@ py_library(
     ],
 )
 
+# TODO(gunan): Investigate making this action hermetic so we do not need
+# to run it locally.
 tf_py_build_info_genrule(
     name = "py_build_info_gen",
     out = "platform/build_info.py",

tensorflow/tensorflow.bzl

@@ -2518,40 +2518,98 @@ def tf_genrule_cmd_append_to_srcs(to_append):
     return ("cat $(SRCS) > $(@) && " + "echo >> $(@) && " + "echo " + to_append +
             " >> $(@)")
 
-def tf_version_info_genrule(name, out):
+def _local_exec_transition_impl(settings, attr):
-    native.genrule(
+    return {
-        name = name,
+        # Force all targets in the subgraph to build on the local machine.
-        srcs = [
+        "//command_line_option:modify_execution_info": ".*=+no-remote-exec",
-            clean_dep("@local_config_git//:gen/spec.json"),
+    }
-            clean_dep("@local_config_git//:gen/head"),
+
-            clean_dep("@local_config_git//:gen/branch_ref"),
+# A transition that forces all targets in the subgraph to be built locally.
-        ],
+_local_exec_transition = transition(
-        outs = [out],
+    implementation = _local_exec_transition_impl,
-        cmd =
+    inputs = [],
-            "$(location //tensorflow/tools/git:gen_git_source) --generate $(SRCS) \"$@\" --git_tag_override=$${GIT_TAG_OVERRIDE:-}",
+    outputs = [
-        local = 1,
+        "//command_line_option:modify_execution_info",
-        exec_tools = [clean_dep("//tensorflow/tools/git:gen_git_source")],
+    ],
+)
+
+def _local_genrule_impl(ctx):
+    ctx.actions.run_shell(
+        outputs = [ctx.outputs.out],
+        inputs = [f for t in ctx.attr.srcs for f in t.files.to_list()],
+        tools = [ctx.executable.exec_tool],
+        arguments = [f.path for t in ctx.attr.srcs for f in t.files.to_list()] +
+                    [ctx.outputs.out.path],
+        command = "%s %s" % (ctx.executable.exec_tool.path, ctx.attr.arguments),
+        execution_requirements = {"no-remote-exec": ""},
+        use_default_shell_env = True,
     )
 
-def tf_py_build_info_genrule(name, out, **kwargs):
+# A genrule that executes locally and forces the tool it runs to be built locally.
-    native.genrule(
+# For python, we want to build all py_binary rules locally that we also want
+# to execute locally, as the remote image might use a different python version.
+# TODO(klimek): Currently we still need to annotate the py_binary rules to use
+# the local platform when building. When we know how to change the platform
+# (https://github.com/bazelbuild/bazel/issues/11035) use this to not require
+# annotating the py_binary rules.
+_local_genrule_internal = rule(
+    implementation = _local_genrule_impl,
+    attrs = {
+        "out": attr.output(),
+        "exec_tool": attr.label(
+            executable = True,
+            cfg = _local_exec_transition,
+            allow_files = True,
+        ),
+        "arguments": attr.string(),
+        "srcs": attr.label_list(
+            allow_files = True,
+        ),
+        "_whitelist_function_transition": attr.label(default = "@bazel_tools//tools/whitelists/function_transition_whitelist"),
+    },
+)
+
+# Wrap the rule in a macro so we can pass in exec_compatible_with.
+def _local_genrule(**kwargs):
+    _local_genrule_internal(
+        exec_compatible_with = [
+            "@local_execution_config_platform//:platform_constraint",
+        ],
+        **kwargs
+    )
+
+def tf_version_info_genrule(name, out):
+    # TODO(gunan): Investigate making this action hermetic so we do not need
+    # to run it locally.
+    _local_genrule(
         name = name,
-        outs = [out],
+        out = out,
-        cmd =
+        exec_tool = "//tensorflow/tools/git:gen_git_source",
-            "$(location //tensorflow/tools/build_info:gen_build_info) --raw_generate \"$@\" " +
+        srcs = [
+            "@local_config_git//:gen/spec.json",
+            "@local_config_git//:gen/head",
+            "@local_config_git//:gen/branch_ref",
+        ],
+        arguments = "--generate \"$@\" --git_tag_override=${GIT_TAG_OVERRIDE:-}",
+    )
+
+def tf_py_build_info_genrule(name, out):
+    _local_genrule(
+        name = name,
+        out = out,
+        exec_tool = "//tensorflow/tools/build_info:gen_build_info",
+        arguments =
+            "--raw_generate \"$@\" " +
             " --is_config_cuda " + if_cuda("True", "False") +
             " --is_config_rocm " + if_rocm("True", "False") +
             " --key_value " +
-            if_cuda(" cuda_version_number=$${TF_CUDA_VERSION:-} cudnn_version_number=$${TF_CUDNN_VERSION:-} ", "") +
+            if_cuda(" cuda_version_number=${TF_CUDA_VERSION:-} cudnn_version_number=${TF_CUDNN_VERSION:-} ", "") +
             if_windows(" msvcp_dll_names=msvcp140.dll,msvcp140_1.dll ", "") +
             if_windows_cuda(" ".join([
                 "nvcuda_dll_name=nvcuda.dll",
-                "cudart_dll_name=cudart64_$$(echo $${TF_CUDA_VERSION:-} | sed \"s/\\.//\").dll",
+                "cudart_dll_name=cudart64_$(echo $${TF_CUDA_VERSION:-} | sed \"s/\\.//\").dll",
-                "cudnn_dll_name=cudnn64_$${TF_CUDNN_VERSION:-}.dll",
+                "cudnn_dll_name=cudnn64_${TF_CUDNN_VERSION:-}.dll",
             ]), ""),
-        local = 1,
-        exec_tools = [clean_dep("//tensorflow/tools/build_info:gen_build_info")],
-        **kwargs
     )
 
 def cc_library_with_android_deps(

tensorflow/tools/build_info/BUILD

@@ -9,8 +9,10 @@ package(
 py_binary(
     name = "gen_build_info",
     srcs = ["gen_build_info.py"],
+    exec_compatible_with = ["@local_execution_config_platform//:platform_constraint"],
     python_version = "PY3",
     srcs_version = "PY2AND3",
+    tags = ["no-remote-exec"],
     deps = [
         "@six_archive//:six",
     ],

tensorflow/tools/git/BUILD

@@ -11,6 +11,8 @@ package(
 py_binary(
     name = "gen_git_source",
     srcs = ["gen_git_source.py"],
+    exec_compatible_with = ["@local_execution_config_platform//:platform_constraint"],
     python_version = "PY3",
     srcs_version = "PY2AND3",
+    tags = ["no-remote-exec"],
 )

tensorflow/workspace.bzl

@@ -77,10 +77,18 @@ def tf_workspace(path_prefix = "", tf_repo_name = ""):
     tf_repositories(path_prefix, tf_repo_name)
     tf_bind()
 
+# Toolchains & platforms required by Tensorflow to build.
+def tf_toolchains():
+    native.register_execution_platforms("@local_execution_config_platform//:platform")
+    native.register_toolchains("@local_execution_config_python//:py_toolchain")
+
 # Define all external repositories required by TensorFlow
 def tf_repositories(path_prefix = "", tf_repo_name = ""):
     """All external dependencies for TF builds."""
 
+    # Initialize toolchains and platforms.
+    tf_toolchains()
+
     # Loads all external repos to configure RBE builds.
     initialize_rbe_configs()
 

third_party/py/BUILD.tpl

@@ -28,6 +28,8 @@ toolchain(
     name = "py_toolchain",
     toolchain = ":py_runtime_pair",
     toolchain_type = "@bazel_tools//tools/python:toolchain_type",
+    target_compatible_with = [%{PLATFORM_CONSTRAINT}],
+    exec_compatible_with = [%{PLATFORM_CONSTRAINT}],
 )
 
 # To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib

third_party/py/python_configure.bzl

@@ -240,11 +240,15 @@ def _create_local_python_repository(repository_ctx):
         "numpy_include",
     )
 
+    platform_constraint = ""
+    if repository_ctx.attr.platform_constraint:
+        platform_constraint = "\"%s\"" % repository_ctx.attr.platform_constraint
     repository_ctx.template("BUILD", build_tpl, {
         "%{PYTHON_BIN_PATH}": python_bin,
         "%{PYTHON_INCLUDE_GENRULE}": python_include_rule,
         "%{PYTHON_IMPORT_LIB_GENRULE}": python_import_lib_genrule,
         "%{NUMPY_INCLUDE_GENRULE}": numpy_include_rule,
+        "%{PLATFORM_CONSTRAINT}": platform_constraint,
     })
 
 def _create_remote_python_repository(repository_ctx, remote_config_repo):
@@ -268,18 +272,31 @@ _ENVIRONS = [
     PYTHON_LIB_PATH,
 ]
 
+local_python_configure = repository_rule(
+    implementation = _create_local_python_repository,
+    environ = _ENVIRONS,
+    attrs = {
+        "environ": attr.string_dict(),
+        "platform_constraint": attr.string(),
+    },
+)
+
 remote_python_configure = repository_rule(
     implementation = _create_local_python_repository,
     environ = _ENVIRONS,
     remotable = True,
     attrs = {
         "environ": attr.string_dict(),
+        "platform_constraint": attr.string(),
     },
 )
 
 python_configure = repository_rule(
     implementation = _python_autoconf_impl,
     environ = _ENVIRONS + [TF_PYTHON_CONFIG_REPO],
+    attrs = {
+        "platform_constraint": attr.string(),
+    },
 )
 """Detects and configures the local Python.
 

third_party/remote_config/BUILD.tpl

@@ -1,8 +1,26 @@
+# Each platform creates a constraint @<platform>//:platform_constraint that
+# is listed in its constraint_values; rule that want to select a specific
+# platform to run on can put @<platform>//:platform_constraing into their
+# exec_compatible_with attribute.
+# Toolchains can similarly be marked with target_compatible_with or
+# exec_compatible_with to bind them to this platform.
+constraint_setting(
+    name = "platform_setting"
+)
+
+constraint_value(
+    name = "platform_constraint",
+    constraint_setting = ":platform_setting",
+    visibility = ["//visibility:public"],
+)
+
 platform(
     name = "platform",
+    visibility = ["//visibility:public"],
     constraint_values = [
-        "@bazel_tools//platforms:x86_64",
+        "@bazel_tools//platforms:%{cpu}",
         "@bazel_tools//platforms:%{platform}",
+        ":platform_constraint",
     ],
     exec_properties = %{exec_properties},
 )

third_party/remote_config/remote_platform_configure.bzl

@@ -2,6 +2,21 @@
 
 def _remote_platform_configure_impl(repository_ctx):
     platform = repository_ctx.attr.platform
+    if platform == "local":
+        os = repository_ctx.os.name.lower()
+        if os.startswith("windows"):
+            platform = "windows"
+        elif os.startswith("mac os"):
+            platform = "osx"
+        else:
+            platform = "linux"
+
+    cpu = "x86_64"
+    if "MACHTYPE" in repository_ctx.os.environ:
+        machine_type = repository_ctx.os.environ["MACHTYPE"]
+        if machine_type.startswith("ppc"):
+            cpu = "ppc"
+
     exec_properties = repository_ctx.attr.platform_exec_properties
 
     serialized_exec_properties = "{"
@@ -15,6 +30,7 @@ def _remote_platform_configure_impl(repository_ctx):
         {
             "%{platform}": platform,
             "%{exec_properties}": serialized_exec_properties,
+            "%{cpu}": cpu,
         },
     )
 
@@ -22,6 +38,6 @@ remote_platform_configure = repository_rule(
     implementation = _remote_platform_configure_impl,
     attrs = {
         "platform_exec_properties": attr.string_dict(mandatory = True),
-        "platform": attr.string(default = "linux", values = ["linux", "windows"]),
+        "platform": attr.string(default = "linux", values = ["linux", "windows", "local"]),
     },
 )

third_party/toolchains/remote_config/configs.bzl

@@ -1,8 +1,12 @@
 """Configurations of RBE builds used with remote config."""
 
-load("//third_party/toolchains/remote_config:rbe_config.bzl", "tensorflow_rbe_config", "tensorflow_rbe_win_config")
+load("//third_party/toolchains/remote_config:rbe_config.bzl", "tensorflow_local_config", "tensorflow_rbe_config", "tensorflow_rbe_win_config")
 
 def initialize_rbe_configs():
+    tensorflow_local_config(
+        name = "local_execution",
+    )
+
     tensorflow_rbe_config(
         name = "ubuntu16.04-manylinux2010-py3",
         os = "ubuntu16.04-manylinux2010",

third_party/toolchains/remote_config/rbe_config.bzl

@@ -1,6 +1,6 @@
 """Macro that creates external repositories for remote config."""
 
-load("//third_party/py:python_configure.bzl", "remote_python_configure")
+load("//third_party/py:python_configure.bzl", "local_python_configure", "remote_python_configure")
 load("//third_party/gpus:cuda_configure.bzl", "remote_cuda_configure")
 load("//third_party/nccl:nccl_configure.bzl", "remote_nccl_configure")
 load("//third_party/gpus:rocm_configure.bzl", "remote_rocm_configure")
@@ -113,6 +113,7 @@ def _tensorflow_rbe_config(name, compiler, python_version, os, rocm_version = No
             name = "%s_config_python" % name,
             environ = env,
             exec_properties = exec_properties,
+            platform_constraint = "@%s_config_platform//:platform_constraint" % name,
         )
 
         remote_rocm_configure(
@@ -127,10 +128,17 @@ def _tensorflow_rbe_config(name, compiler, python_version, os, rocm_version = No
             "Pool": "default",
         }
 
+        remote_platform_configure(
+            name = "%s_config_platform" % name,
+            platform = "linux",
+            platform_exec_properties = exec_properties,
+        )
+
         remote_python_configure(
             name = "%s_config_python" % name,
             environ = env,
             exec_properties = exec_properties,
+            platform_constraint = "@%s_config_platform//:platform_constraint" % name,
         )
     else:
         fail("Neither cuda_version, rocm_version nor python_version specified.")
@@ -156,7 +164,20 @@ def _tensorflow_rbe_win_config(name, python_bin_path, container_name = "windows-
         name = "%s_config_python" % name,
         environ = env,
         exec_properties = exec_properties,
+        platform_constraint = "@%s_config_platform//:platform_constraint" % name,
+    )
+
+def _tensorflow_local_config(name):
+    remote_platform_configure(
+        name = "%s_config_platform" % name,
+        platform = "local",
+        platform_exec_properties = {},
+    )
+    local_python_configure(
+        name = "%s_config_python" % name,
+        platform_constraint = "@%s_config_platform//:platform_constraint" % name,
     )
 
 tensorflow_rbe_config = _tensorflow_rbe_config
 tensorflow_rbe_win_config = _tensorflow_rbe_win_config
+tensorflow_local_config = _tensorflow_local_config