ROCm · i-chaochen · Oct 5, 2023 · Oct 5, 2023 · Oct 5, 2023 · Oct 5, 2023
diff --git a/third_party/gpus/cuda/cuda_config.h.tpl b/third_party/gpus/cuda/cuda_config.h.tpl
@@ -25,7 +25,6 @@ limitations under the License.
 #define TF_CUFFT_VERSION "%{cufft_version}"
 #define TF_CUSPARSE_VERSION "%{cusparse_version}"
 #define TF_CUDNN_VERSION "%{cudnn_version}"
-#define TF_NCCL_VERSION "%{nccl_version}"
 
 #define TF_CUDA_TOOLKIT_PATH "%{cuda_toolkit_path}"
 

diff --git a/third_party/nccl/archive.BUILD b/third_party/nccl/archive.BUILD
@@ -5,6 +5,8 @@ licenses(["notice"])
 
 exports_files(["LICENSE.txt"])
 
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load(
     "@local_config_cuda//cuda:build_defs.bzl",
     "cuda_library",
@@ -15,6 +17,44 @@ load(
     "gen_device_srcs",
 )
 
+NCCL_MAJOR = 2
+
+NCCL_MINOR = 16
+
+NCCL_PATCH = 5
+
+NCCL_VERSION = NCCL_MAJOR * 10000 + NCCL_MINOR * 100 + NCCL_PATCH  # e.g., 21605
+
+expand_template(
+    name = "nccl_header_version",
+    out = "src/nccl.h",
+    substitutions = {
+        "${nccl:Major}": str(NCCL_MAJOR),
+        "${nccl:Minor}": str(NCCL_MINOR),
+        "${nccl:Patch}": str(NCCL_PATCH),
+        "${nccl:Suffix}": "\"\"",
+        "${nccl:Version}": str(NCCL_VERSION),
+    },
+    template = "src/nccl.h.in",
+)
+
+# This additional header allows us to determine the configured NCCL version
+# without including the rest of NCCL.
+write_file(
+    name = "nccl_config_header",
+    out = "nccl_config.h",
+    content = [
+        "#define TF_NCCL_VERSION \"{}\"".format(NCCL_MAJOR),
+    ],
+)
+
+cc_library(
+    name = "nccl_config",
+    hdrs = ["nccl_config.h"],
+    include_prefix = "third_party/nccl",
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "src_hdrs",
     hdrs = [

diff --git a/third_party/nccl/archive.patch b/third_party/nccl/archive.patch
@@ -30,31 +30,6 @@ diff --git a/src/collectives/device/sendrecv.cu b/src/collectives/device/sendrec
 similarity index 100%
 rename from src/collectives/device/sendrecv.cu
 rename to src/collectives/device/sendrecv.cu.cc
-diff --git a/src/nccl.h.in b/src/nccl.h
-similarity index 98%
-rename from src/nccl.h.in
-rename to src/nccl.h
-index 985274e..7ebb1e1 100644
---- a/src/nccl.h.in
-+++ b/src/nccl.h
-@@ -13,12 +13,12 @@
- #include <cuda_bf16.h>
- #endif
-
--#define NCCL_MAJOR ${nccl:Major}
--#define NCCL_MINOR ${nccl:Minor}
--#define NCCL_PATCH ${nccl:Patch}
--#define NCCL_SUFFIX "${nccl:Suffix}"
-+#define NCCL_MAJOR 2
-+#define NCCL_MINOR 16
-+#define NCCL_PATCH 5
-+#define NCCL_SUFFIX ""
-
--#define NCCL_VERSION_CODE ${nccl:Version}
-+#define NCCL_VERSION_CODE 2165
- #define NCCL_VERSION(X,Y,Z) (((X) <= 2 && (Y) <= 8) ? (X) * 1000 + (Y) * 100 + (Z) : (X) * 10000 + (Y) * 100 + (Z))
-
- #ifdef __cplusplus
 diff --git a/src/include/nvtx.h b/src/include/nvtx.h
 index 2aeb932..cdc67d2 100644
 --- a/src/include/nvtx.h

diff --git a/third_party/nccl/nccl_configure.bzl b/third_party/nccl/nccl_configure.bzl
@@ -50,6 +50,13 @@ cc_library(
   name = "nccl",
   visibility = ["//visibility:public"],
 )
+
+cc_library(
+  name = "nccl_config",
+  hdrs = ["nccl_config.h"],
+  include_prefix = "third_party/nccl",
+  visibility = ["//visibility:public"],
+)
 """
 
 _NCCL_ARCHIVE_BUILD_CONTENT = """
@@ -64,6 +71,12 @@ alias(
   actual = "@nccl_archive//:nccl",
   visibility = ["//visibility:public"],
 )
+
+alias(
+  name = "nccl_config",
+  actual = "@nccl_archive//:nccl_config",
+  visibility = ["//visibility:public"],
+)
 """
 
 _NCCL_ARCHIVE_STUB_BUILD_CONTENT = """
@@ -84,6 +97,12 @@ alias(
   actual = "@nccl_archive//:nccl_headers",
   visibility = ["//visibility:public"],
 )
+
+alias(
+  name = "nccl_config",
+  actual = "@nccl_archive//:nccl_config",
+  visibility = ["//visibility:public"],
+)
 """
 
 def _label(file):
@@ -142,6 +161,7 @@ def _nccl_autoconf_impl(repository_ctx):
         get_cpu_value(repository_ctx) not in ("Linux", "FreeBSD")):
         # Add a dummy build file to make bazel query happy.
         repository_ctx.file("BUILD", _NCCL_DUMMY_BUILD_CONTENT)
+        repository_ctx.file("nccl_config.h", "#define TF_NCCL_VERSION \"\"")
     elif get_host_environ(repository_ctx, "TF_NCCL_CONFIG_REPO") != None:
         _create_remote_nccl_repository(repository_ctx, get_host_environ(repository_ctx, "TF_NCCL_CONFIG_REPO"))
     else:

diff --git a/third_party/nccl/system.BUILD.tpl b/third_party/nccl/system.BUILD.tpl
@@ -1,5 +1,6 @@
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load(
-    "@tsl//tsl/platform/default:cuda_build_defs.bzl",
+    "@xla//tensorflow/platform/default:cuda_build_defs.bzl",
     "cuda_rpath_flags"
 )
 
@@ -20,6 +21,16 @@ cc_library(
     linkopts = cuda_rpath_flags("nvidia/nccl/lib"),
 )
 
+cc_library(
+    name = "nccl_headers",
+    hdrs = ["nccl.h"],
+    include_prefix = "third_party/nccl",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@local_config_cuda//cuda:cuda_headers",
+    ],
+)
+
 genrule(
     name = "nccl-files",
     outs = [
@@ -32,3 +43,20 @@ cp "%{nccl_library_dir}/libnccl.so.%{nccl_version}" \
   "$(@D)/libnccl.so.%{nccl_version}"
 """,
 )
+
+# This additional header allows us to determine the configured NCCL version
+# without including the rest of NCCL.
+write_file(
+    name = "nccl_config_header",
+    out = "nccl_config.h",
+    content = [
+        "#define TF_NCCL_VERSION \"%{nccl_version}\""
+    ]
+)
+
+cc_library(
+    name = "nccl_config",
+    hdrs = ["nccl_config.h"],
+    include_prefix = "third_party/nccl",
+    visibility = ["//visibility:public"],
+)
diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
@@ -181,18 +181,6 @@ diff --ruN a/stablehlo/CMakeLists.txt b/stablehlo/CMakeLists.txt
 
  #-------------------------------------------------------------------------------
  # Directory setup
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll b/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
---- stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
-+++ stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
-@@ -141,7 +141,7 @@
- Pattern =>
-   replace op<stablehlo.multiply>(input0 : Value<_: Tosa_Tensor>,
-                             input1 : Value<_: Tosa_Tensor>)
--     with op<tosa.mul>(input0, input1) {shift = attr<"0 : i32">};
-+     with op<tosa.mul>(input0, input1) {shift = attr<"0 : i8">};
- Pattern =>
-   replace op<stablehlo.or>(input0 : Value<_: Tosa_Tensor>,
-                       input1 : Value<_: Tosa_Tensor>)
 diff --ruN a/stablehlo/stablehlo/dialect/Base.cpp b/stablehlo/stablehlo/dialect/Base.cpp
 --- stablehlo/stablehlo/dialect/Base.cpp
 +++ stablehlo/stablehlo/dialect/Base.cpp

diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
     # LINT.IfChange
-    STABLEHLO_COMMIT = "fe5fa2690e33241b51cefdb27a1fff0fd0626047"
-    STABLEHLO_SHA256 = "e4246b176a911ecc8fc9f42622ba1a3997f88633d47dd90ad641ee3d973b6792"
+    STABLEHLO_COMMIT = "46a2506405a0c1ade5327b9e476a32781221b075"
+    STABLEHLO_SHA256 = "9232d6131a17a264ed9a6c79116fdb53787aadce5c14a6c68399c8617eaddd58"
     # LINT.ThenChange(Google-internal path)
 
     tf_http_archive(

diff --git a/third_party/tsl/third_party/gpus/cuda/cuda_config.h.tpl b/third_party/tsl/third_party/gpus/cuda/cuda_config.h.tpl
@@ -25,7 +25,6 @@ limitations under the License.
 #define TF_CUFFT_VERSION "%{cufft_version}"
 #define TF_CUSPARSE_VERSION "%{cusparse_version}"
 #define TF_CUDNN_VERSION "%{cudnn_version}"
-#define TF_NCCL_VERSION "%{nccl_version}"
 
 #define TF_CUDA_TOOLKIT_PATH "%{cuda_toolkit_path}"
 

diff --git a/third_party/tsl/third_party/nccl/archive.BUILD b/third_party/tsl/third_party/nccl/archive.BUILD
@@ -5,6 +5,8 @@ licenses(["notice"])
 
 exports_files(["LICENSE.txt"])
 
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load(
     "@local_config_cuda//cuda:build_defs.bzl",
     "cuda_library",
@@ -15,6 +17,44 @@ load(
     "gen_device_srcs",
 )
 
+NCCL_MAJOR = 2
+
+NCCL_MINOR = 16
+
+NCCL_PATCH = 5
+
+NCCL_VERSION = NCCL_MAJOR * 10000 + NCCL_MINOR * 100 + NCCL_PATCH  # e.g., 21605
+
+expand_template(
+    name = "nccl_header_version",
+    out = "src/nccl.h",
+    substitutions = {
+        "${nccl:Major}": str(NCCL_MAJOR),
+        "${nccl:Minor}": str(NCCL_MINOR),
+        "${nccl:Patch}": str(NCCL_PATCH),
+        "${nccl:Suffix}": "\"\"",
+        "${nccl:Version}": str(NCCL_VERSION),
+    },
+    template = "src/nccl.h.in",
+)
+
+# This additional header allows us to determine the configured NCCL version
+# without including the rest of NCCL.
+write_file(
+    name = "nccl_config_header",
+    out = "nccl_config.h",
+    content = [
+        "#define TF_NCCL_VERSION \"{}\"".format(NCCL_MAJOR),
+    ],
+)
+
+cc_library(
+    name = "nccl_config",
+    hdrs = ["nccl_config.h"],
+    include_prefix = "third_party/nccl",
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "src_hdrs",
     hdrs = [

diff --git a/third_party/tsl/third_party/nccl/archive.patch b/third_party/tsl/third_party/nccl/archive.patch
@@ -30,31 +30,6 @@ diff --git a/src/collectives/device/sendrecv.cu b/src/collectives/device/sendrec
 similarity index 100%
 rename from src/collectives/device/sendrecv.cu
 rename to src/collectives/device/sendrecv.cu.cc
-diff --git a/src/nccl.h.in b/src/nccl.h
-similarity index 98%
-rename from src/nccl.h.in
-rename to src/nccl.h
-index 985274e..7ebb1e1 100644
---- a/src/nccl.h.in
-+++ b/src/nccl.h
-@@ -13,12 +13,12 @@
- #include <cuda_bf16.h>
- #endif
-
--#define NCCL_MAJOR ${nccl:Major}
--#define NCCL_MINOR ${nccl:Minor}
--#define NCCL_PATCH ${nccl:Patch}
--#define NCCL_SUFFIX "${nccl:Suffix}"
-+#define NCCL_MAJOR 2
-+#define NCCL_MINOR 16
-+#define NCCL_PATCH 5
-+#define NCCL_SUFFIX ""
-
--#define NCCL_VERSION_CODE ${nccl:Version}
-+#define NCCL_VERSION_CODE 2165
- #define NCCL_VERSION(X,Y,Z) (((X) <= 2 && (Y) <= 8) ? (X) * 1000 + (Y) * 100 + (Z) : (X) * 10000 + (Y) * 100 + (Z))
-
- #ifdef __cplusplus
 diff --git a/src/include/nvtx.h b/src/include/nvtx.h
 index 2aeb932..cdc67d2 100644
 --- a/src/include/nvtx.h

diff --git a/third_party/tsl/third_party/nccl/nccl_configure.bzl b/third_party/tsl/third_party/nccl/nccl_configure.bzl
@@ -50,6 +50,13 @@ cc_library(
   name = "nccl",
   visibility = ["//visibility:public"],
 )
+
+cc_library(
+  name = "nccl_config",
+  hdrs = ["nccl_config.h"],
+  include_prefix = "third_party/nccl",
+  visibility = ["//visibility:public"],
+)
 """
 
 _NCCL_ARCHIVE_BUILD_CONTENT = """
@@ -64,6 +71,12 @@ alias(
   actual = "@nccl_archive//:nccl",
   visibility = ["//visibility:public"],
 )
+
+alias(
+  name = "nccl_config",
+  actual = "@nccl_archive//:nccl_config",
+  visibility = ["//visibility:public"],
+)
 """
 
 _NCCL_ARCHIVE_STUB_BUILD_CONTENT = """
@@ -84,6 +97,12 @@ alias(
   actual = "@nccl_archive//:nccl_headers",
   visibility = ["//visibility:public"],
 )
+
+alias(
+  name = "nccl_config",
+  actual = "@nccl_archive//:nccl_config",
+  visibility = ["//visibility:public"],
+)
 """
 
 def _label(file):
@@ -142,6 +161,7 @@ def _nccl_autoconf_impl(repository_ctx):
         get_cpu_value(repository_ctx) not in ("Linux", "FreeBSD")):
         # Add a dummy build file to make bazel query happy.
         repository_ctx.file("BUILD", _NCCL_DUMMY_BUILD_CONTENT)
+        repository_ctx.file("nccl_config.h", "#define TF_NCCL_VERSION \"\"")
     elif get_host_environ(repository_ctx, "TF_NCCL_CONFIG_REPO") != None:
         _create_remote_nccl_repository(repository_ctx, get_host_environ(repository_ctx, "TF_NCCL_CONFIG_REPO"))
     else: