1. Build on Windows

2. Fix OverflowError

    When calling `key = random.PRNGKey(0)` OverflowError: Python int too
    large to convert to C long for casting value 4294967295 (0xFFFFFFFF)
    from python int to int32.

3. fix file path in regex of errors_test

4. handle ValueError of os.path.commonpath

Author: cloudhan

.bazelversion

@@ -0,0 +1 @@
+3.1.0

WORKSPACE

@@ -28,10 +28,10 @@ http_archive(
 #    and update the sha256 with the result.
 http_archive(
     name = "org_tensorflow",
-    sha256 = "579a74ad171d8da7b7193ff863f28482c2e6050c4090650b001fb80bbc46bb0f",
-    strip_prefix = "tensorflow-04f25b55e27be95ec340f414c2a1cabe16be5c2a",
+    sha256 = "f59fc70b349373267a928f37d3856984781214c07ca53ef075c48389778d17cf",
+    strip_prefix = "tensorflow-b7060f1c189a3b2d97c8eca13d5a13acc5553403",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/04f25b55e27be95ec340f414c2a1cabe16be5c2a.tar.gz",
+        "https://github.com/tensorflow/tensorflow/archive/b7060f1c189a3b2d97c8eca13d5a13acc5553403.tar.gz",
     ],
 )
 

build/BUILD.bazel

@@ -15,25 +15,27 @@
 # JAX is Autograd and XLA
 
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
+load("@org_tensorflow//tensorflow:tensorflow.bzl", "if_not_windows")
 
 licenses(["notice"])  # Apache 2
 
 package(default_visibility = ["//visibility:public"])
 
-sh_binary(
+py_binary(
     name = "install_xla_in_source_tree",
-    srcs = ["install_xla_in_source_tree.sh"],
+    srcs = ["install_xla_in_source_tree.py"],
     data = [
         "@org_tensorflow//tensorflow/compiler/xla/python:xla_client",
-        "@org_tensorflow//tensorflow/compiler/xla/python/tpu_driver/client:py_tpu_client",
         "//jaxlib",
         "//jaxlib:lapack.so",
         "//jaxlib:_pocketfft.so",
         "//jaxlib:pocketfft_flatbuffers_py",
-    ] + if_cuda([
+    ] + if_not_windows([
+        "@org_tensorflow//tensorflow/compiler/xla/python/tpu_driver/client:py_tpu_client",
+    ]) + if_cuda([
         "//jaxlib:cublas_kernels",
         "//jaxlib:cusolver_kernels",
         "//jaxlib:cuda_prng_kernels",
     ]),
-    deps = ["@bazel_tools//tools/bash/runfiles"],
+    deps = ["@bazel_tools//tools/python/runfiles"],
 )

build/build.py

@@ -43,6 +43,10 @@
 # pylint: enable=g-import-not-at-top
 
 
+def is_windows():
+    return sys.platform.startswith("win32")
+
+
 def shell(cmd):
   output = subprocess.check_output(cmd)
   return output.decode("UTF-8").strip()
@@ -52,7 +56,8 @@ def shell(cmd):
 
 def get_python_bin_path(python_bin_path_flag):
   """Returns the path to the Python interpreter to use."""
-  return python_bin_path_flag or sys.executable
+  path = python_bin_path_flag or sys.executable
+  return path.replace(os.sep, "/")
 
 
 def get_python_version(python_bin_path):
@@ -189,7 +194,8 @@ def check_bazel_version(bazel_path, min_version, max_version):
 build --repo_env TF_NEED_CUDA="{tf_need_cuda}"
 build --action_env TF_CUDA_COMPUTE_CAPABILITIES="{cuda_compute_capabilities}"
 build --distinct_host_configuration=false
-build --copt=-Wno-sign-compare
+build:linux --copt=-Wno-sign-compare
+build:macos --copt=-Wno-sign-compare
 build -c opt
 build:opt --copt=-march=native
 build:opt --host_copt=-march=native
@@ -205,9 +211,12 @@ def check_bazel_version(bazel_path, min_version, max_version):
 build --define open_source_build=true
 
 # Disable enabled-by-default TensorFlow features that we don't care about.
-build --define=no_aws_support=true
-build --define=no_gcp_support=true
-build --define=no_hdfs_support=true
+build:linux --define=no_aws_support=true
+build:macos --define=no_aws_support=true
+build:linux --define=no_gcp_support=true
+build:macos --define=no_gcp_support=true
+build:linux --define=no_hdfs_support=true
+build:macos --define=no_hdfs_support=true
 build --define=no_kafka_support=true
 build --define=no_ignite_support=true
 build --define=grpc_no_ares=true
@@ -218,16 +227,49 @@ def check_bazel_version(bazel_path, min_version, max_version):
 build --spawn_strategy=standalone
 build --strategy=Genrule=standalone
 
-build --cxxopt=-std=c++14
-build --host_cxxopt=-std=c++14
+build --enable_platform_specific_config
+
+# Tensorflow uses M_* math constants that only get defined by MSVC headers if
+# _USE_MATH_DEFINES is defined.
+build:windows --copt=/D_USE_MATH_DEFINES
+build:windows --host_copt=/D_USE_MATH_DEFINES
+
+# Make sure to include as little of windows.h as possible
+build:windows --copt=-DWIN32_LEAN_AND_MEAN
+build:windows --host_copt=-DWIN32_LEAN_AND_MEAN
+build:windows --copt=-DNOGDI
+build:windows --host_copt=-DNOGDI
+
+# https://devblogs.microsoft.com/cppblog/announcing-full-support-for-a-c-c-conformant-preprocessor-in-msvc/
+# otherwise, there will be some compiling error due to preprocessing.
+build:windows --copt=/Zc:preprocessor
+
+build:linux --cxxopt=-std=c++14
+build:linux --host_cxxopt=-std=c++14
+
+build:macos --cxxopt=-std=c++14
+build:macos --host_cxxopt=-std=c++14
+
+build:windows --cxxopt=/std:c++14
+build:windows --host_cxxopt=/std:c++14
+
+# Generate PDB files, to generate useful PDBs, in opt compilation_mode
+# --copt /Z7 is needed.
+build:windows --linkopt=/DEBUG
+build:windows --host_linkopt=/DEBUG
+build:windows --linkopt=/OPT:REF
+build:windows --host_linkopt=/OPT:REF
+build:windows --linkopt=/OPT:ICF
+build:windows --host_linkopt=/OPT:ICF
 
 # Suppress all warning messages.
 build:short_logs --output_filter=DONT_MATCH_ANYTHING
 """
 
 
 
-def write_bazelrc(cuda_toolkit_path=None, cudnn_install_path=None, **kwargs):
+def write_bazelrc(cuda_toolkit_path=None, cudnn_install_path=None,
+                  cuda_version=None, cudnn_version=None, **kwargs):
   with open("../.bazelrc", "w") as f:
     f.write(BAZELRC_TEMPLATE.format(**kwargs))
     if cuda_toolkit_path:
@@ -236,7 +278,12 @@ def write_bazelrc(cuda_toolkit_path=None, cudnn_install_path=None, **kwargs):
     if cudnn_install_path:
       f.write("build --action_env CUDNN_INSTALL_PATH=\"{cudnn_install_path}\"\n"
               .format(cudnn_install_path=cudnn_install_path))
-
+    if cuda_version:
+      f.write("build --action_env TF_CUDA_VERSION=\"{cuda_version}\"\n"
+              .format(cuda_version=cuda_version))
+    if cudnn_version:
+      f.write("build --action_env TF_CUDNN_VERSION=\"{cudnn_version}\"\n"
+              .format(cudnn_version=cudnn_version))
 
 BANNER = r"""
      _   _  __  __
@@ -317,6 +364,14 @@ def main():
       "--cudnn_path",
       default=None,
       help="Path to CUDNN libraries.")
+  parser.add_argument(
+      "--cuda_version",
+      default=None,
+      help="CUDA toolkit version, e.g., 11.1")
+  parser.add_argument(
+      "--cudnn_version",
+      default=None,
+      help="CUDNN version, e.g., 8")
   parser.add_argument(
       "--cuda_compute_capabilities",
       default="3.5,5.2,6.0,6.1,7.0",
@@ -331,6 +386,12 @@ def main():
       help="Additional options to pass to bazel.")
   args = parser.parse_args()
 
+  if is_windows() and args.enable_cuda:
+    if args.cuda_version is None:
+      parser.error("--cuda_version is needed for Windows CUDA build.")
+    if args.cudnn_version is None:
+      parser.error("--cudnn_version is needed for Windows CUDA build.")
+
   print(BANNER)
   os.chdir(os.path.dirname(__file__ or args.prog) or '.')
 
@@ -357,12 +418,18 @@ def main():
     if cudnn_install_path:
       print("CUDNN library path: {}".format(cudnn_install_path))
     print("CUDA compute capabilities: {}".format(args.cuda_compute_capabilities))
+    if args.cuda_version:
+      print("CUDA version: {}".format(args.cuda_version))
+    if args.cudnn_version:
+      print("CUDNN version: {}".format(args.cudnn_version))
   write_bazelrc(
       python_bin_path=python_bin_path,
       tf_need_cuda=1 if args.enable_cuda else 0,
       cuda_toolkit_path=cuda_toolkit_path,
       cudnn_install_path=cudnn_install_path,
-      cuda_compute_capabilities=args.cuda_compute_capabilities)
+      cuda_compute_capabilities=args.cuda_compute_capabilities,
+      cuda_version=args.cuda_version,
+      cudnn_version=args.cudnn_version)
 
   print("\nBuilding XLA and installing it in the jaxlib source tree...")
   config_args = args.bazel_options

build/install_xla_in_source_tree.py

@@ -0,0 +1,108 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+
+import os
+import sys
+import shutil
+import argparse
+
+from bazel_tools.tools.python.runfiles import runfiles
+
+parser = argparse.ArgumentParser()
+parser.add_argument("target")
+args = parser.parse_args()
+
+r = runfiles.Create()
+
+
+def _is_windows():
+  return sys.platform.startswith("win32")
+
+
+def _copy_so(src_file, dst_dir):
+  src_filename = os.path.basename(src_file)
+  if _is_windows() and src_filename.endswith(".so"):
+    dst_filename = src_filename[:-3] + ".pyd"
+  else:
+    dst_filename = src_filename
+  dst_file = os.path.join(dst_dir, dst_filename)
+  shutil.copyfile(src_file, dst_file)
+
+
+def _copy_normal(src_file, dst_dir):
+  src_filename = os.path.basename(src_file)
+  dst_file = os.path.join(dst_dir, src_filename)
+  shutil.copyfile(src_file, dst_file)
+
+
+def copy(src_file, dst_dir=os.path.join(args.target, "jaxlib")):
+  if src_file.endswith(".so"):
+    _copy_so(src_file, dst_dir)
+  else:
+    _copy_normal(src_file, dst_dir)
+
+
+def patch_copy_xla_client_py(dst_dir=os.path.join(args.target, "jaxlib")):
+  with open(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/xla_client.py")) as f:
+    src = f.read()
+    src = src.replace("from tensorflow.compiler.xla.python import xla_extension as _xla",
+                      "from . import xla_extension as _xla")
+    src = src.replace("from tensorflow.compiler.xla.python.xla_extension import ops",
+                      "from .xla_extension import ops")
+    with open(os.path.join(dst_dir, "xla_client.py"), "w") as f:
+      f.write(src)
+
+
+def patch_copy_tpu_client_py(dst_dir=os.path.join(args.target, "jaxlib")):
+  with open(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client.py")) as f:
+    src = f.read()
+    src = src.replace("from tensorflow.compiler.xla.python import xla_extension as _xla",
+                      "from . import xla_extension as _xla")
+    src = src.replace("from tensorflow.compiler.xla.python import xla_client",
+                      "from . import xla_client")
+    src = src.replace(
+        "from tensorflow.compiler.xla.python.tpu_driver.client import tpu_client_extension as _tpu_client",
+        "from . import tpu_client_extension as _tpu_client")
+    with open(os.path.join(dst_dir, "tpu_client.py"), "w") as f:
+      f.write(src)
+
+
+copy(r.Rlocation("__main__/jaxlib/lapack.so"))
+copy(r.Rlocation("__main__/jaxlib/_pocketfft.so"))
+copy(r.Rlocation("__main__/jaxlib/pocketfft_flatbuffers_py_generated.py"))
+copy(r.Rlocation("__main__/jaxlib/pocketfft.py"))
+if r.Rlocation("__main__/jaxlib/cusolver_kernels.so") is not None:
+  copy(r.Rlocation("__main__/jaxlib/cusolver_kernels.so"))
+  copy(r.Rlocation("__main__/jaxlib/cublas_kernels.so"))
+  copy(r.Rlocation("__main__/jaxlib/cusolver_kernels.so"))
+  copy(r.Rlocation("__main__/jaxlib/cuda_prng_kernels.so"))
+if r.Rlocation("__main__/jaxlib/cusolver_kernels.pyd") is not None:
+  copy(r.Rlocation("__main__/jaxlib/cusolver_kernels.pyd"))
+  copy(r.Rlocation("__main__/jaxlib/cublas_kernels.pyd"))
+  copy(r.Rlocation("__main__/jaxlib/cusolver_kernels.pyd"))
+  copy(r.Rlocation("__main__/jaxlib/cuda_prng_kernels.pyd"))
+copy(r.Rlocation("__main__/jaxlib/version.py"))
+copy(r.Rlocation("__main__/jaxlib/cusolver.py"))
+copy(r.Rlocation("__main__/jaxlib/cuda_prng.py"))
+
+if _is_windows():
+  copy(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/xla_extension.pyd"))
+else:
+  copy(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/xla_extension.so"))
+patch_copy_xla_client_py()
+
+if not _is_windows():
+  copy(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client_extension.so"))
+  patch_copy_tpu_client_py()