[ROCm] Add support for use of the built in (rocm) clang

.github/workflows/rocm_linux_x86_64_linux_x86_64_build.yaml

@@ -86,3 +86,21 @@ jobs:
             --define=using_rocm=true \
             //cc/tests/cpu:all
 
+      - name: Test ROCm GPU code compilation with ROCm distro clang
+        run: |
+          bazel test --color=yes \
+            --config=rocm \
+            --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908" \
+            --@rules_ml_toolchain//common:use_rocm_clang=True \
+            --define=using_rocm=true \
+            //cc/tests/gpu/rocm:all
+
+      - name: Test host (CPU) code compilation with ROCm distro clang
+        run: |
+          bazel test --color=yes \
+            --config=rocm \
+            --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908" \
+            --@rules_ml_toolchain//common:use_rocm_clang=True \
+            --define=using_rocm=true \
+            //cc/tests/cpu:all
+

cc/impls/linux_x86_64_linux_x86_64_rocm/BUILD

@@ -54,7 +54,10 @@ rocm_hipcc_feature(
     clang_version = hipcc_config().clang_version,
     enabled = True,
     hipcc_path = hipcc_config().hipcc_path,
-    host_compiler = "@llvm_linux_x86_64//:clang",
+    host_compiler = select({
+        "//common:is_rocm_clang_enabled": "@config_rocm_hipcc//rocm:clang",
+        "//conditions:default": "@llvm_linux_x86_64//:clang",
+    }),
     rocm_path = hipcc_config().rocm_path,
     rocm_toolkit = "@config_rocm_hipcc//rocm:toolchain_data",
     version = str(hipcc_config().version_number),
@@ -101,19 +104,31 @@ cc_toolchain_import_feature(
 # buildifier: leave-alone
 cc_toolchain_import(
     name = "imports",
-    deps = [
-        # Always include hermetic sysroot headers - needed for GPU compilation
-        # GPU compilation requires hermetic headers (system headers incompatible with ROCm clang)
-        # Host compilation will use system headers when HOST_SYSROOT is set (via wrapper override)
-        # IMPORTANT: std_incs must come before compiler_incs to ensure correct stdatomic.h resolution
-        # The sysroot C++ stdatomic.h wrapper uses #include_next to find clang's built-in stdatomic.h
-        "@sysroot_linux_x86_64//:std_incs",
-        # Hermetic LLVM headers
-        "@llvm_linux_x86_64//:compiler_incs",
-        "@sysroot_linux_x86_64//:sys_incs",
-        ":std_libs",
-        "@sysroot_linux_x86_64//:sys_libs",
-    ],
+    deps = select({
+        "//common:is_rocm_clang_enabled": [
+            # Use ROCm compiler builtins but sysroot for C++ stdlib
+            # (ROCm's C++ headers are incomplete - missing __config_site)
+            "@config_rocm_hipcc//rocm:compiler_incs",
+            "@sysroot_linux_x86_64//:std_incs",
+            "@sysroot_linux_x86_64//:sys_incs",
+            "@sysroot_linux_x86_64//:std_libs",
+            "@sysroot_linux_x86_64//:sys_libs",
+        ],
+        "//conditions:default": [
+            # Use hermetic LLVM toolchain
+            # Always include hermetic sysroot headers - needed for GPU compilation
+            # GPU compilation requires hermetic headers (system headers incompatible with ROCm clang)
+            # Host compilation will use system headers when HOST_SYSROOT is set (via wrapper override)
+            # IMPORTANT: std_incs must come before compiler_incs to ensure correct stdatomic.h resolution
+            # The sysroot C++ stdatomic.h wrapper uses #include_next to find clang's built-in stdatomic.h
+            "@sysroot_linux_x86_64//:std_incs",
+            # Hermetic LLVM headers
+            "@llvm_linux_x86_64//:compiler_incs",
+            "@sysroot_linux_x86_64//:sys_incs",
+            ":std_libs",
+            "@sysroot_linux_x86_64//:sys_libs",
+        ],
+    }),
     visibility = ["//visibility:public"],
 )
 
@@ -284,13 +299,28 @@ FEATURES_ESSENTIAL = [
 
 cc_toolchain_config(
     name = "config",
-    archiver = "@llvm_linux_x86_64//:ar",
-    c_compiler = "@llvm_linux_x86_64//:clang",
-    cc_compiler = "@llvm_linux_x86_64//:clang++",
+    archiver = select({
+        "//common:is_rocm_clang_enabled": "@config_rocm_hipcc//rocm:ar",
+        "//conditions:default": "@llvm_linux_x86_64//:ar",
+    }),
+    c_compiler = select({
+        "//common:is_rocm_clang_enabled": "@config_rocm_hipcc//rocm:clang",
+        "//conditions:default": "@llvm_linux_x86_64//:clang",
+    }),
+    cc_compiler = select({
+        "//common:is_rocm_clang_enabled": "@config_rocm_hipcc//rocm:clang++",
+        "//conditions:default": "@llvm_linux_x86_64//:clang++",
+    }),
     compiler_features = FEATURES_ESSENTIAL,
     cxx_builtin_include_directories = [],
-    linker = "@llvm_linux_x86_64//:ld",
-    strip_tool = "@llvm_linux_x86_64//:strip",
+    linker = select({
+        "//common:is_rocm_clang_enabled": "@config_rocm_hipcc//rocm:ld",
+        "//conditions:default": "@llvm_linux_x86_64//:ld",
+    }),
+    strip_tool = select({
+        "//common:is_rocm_clang_enabled": "@config_rocm_hipcc//rocm:strip",
+        "//conditions:default": "@llvm_linux_x86_64//:strip",
+    }),
     target_cpu = "x86_64",
     target_system_name = "local",
     tool_paths = ROCM_TOOLS,
@@ -311,18 +341,22 @@ cc_toolchain(
     toolchain_identifier = "toolchain_linux_x86_64_linux_x86_64_rocm_id",
 )
 
-# Aliases to select between libc++ and libstdc++
+# Aliases to select between libc++ and libstdc++, and between hermetic LLVM and ROCm LLVM
+# ROCm clang uses sysroot C++ headers (ROCm's headers are incomplete)
 alias(
     name = "std_incs",
     actual = select({
+        "//common:is_rocm_clang_enabled": "@sysroot_linux_x86_64//:std_incs",
         "//common:is_stdlib_libcxx": "@llvm_linux_x86_64//:std_incs",
         "//conditions:default": "@sysroot_linux_x86_64//:std_incs",
     }),
 )
 
+# ROCm clang uses sysroot libraries (ROCm doesn't ship static C++ libs)
 alias(
     name = "std_libs",
     actual = select({
+        "//common:is_rocm_clang_enabled": "@sysroot_linux_x86_64//:std_libs",
         "//common:is_stdlib_libcxx": "@llvm_linux_x86_64//:std_libs",
         "//conditions:default": "@sysroot_linux_x86_64//:std_libs",
     }),

common/BUILD

@@ -144,6 +144,19 @@ config_setting(
     },
 )
 
+# ROCm-specific flags
+bool_flag(
+    name = "use_rocm_clang",
+    build_setting_default = False,
+)
+
+config_setting(
+    name = "is_rocm_clang_enabled",
+    flag_values = {
+        ":use_rocm_clang": "True",
+    },
+)
+
 #######################################################
 # Enable CUDA support flags
 #

gpu/rocm/BUILD.tpl

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 load("@rules_cc//cc:defs.bzl", "cc_library")
+load("@rules_ml_toolchain//third_party/rules_cc_toolchain/features:cc_toolchain_import.bzl", "cc_toolchain_import")
 
 licenses(["restricted"])  # MPL2, portions GPL v3, LGPL v3, BSD-like
 
@@ -74,3 +75,80 @@ config_setting(
     },
     visibility = ["//visibility:public"],
 )
+
+# ROCm distribution's clang compiler (single file for use in attributes)
+filegroup(
+    name = "clang",
+    srcs = glob(
+        ["%{rocm_root}/llvm/bin/clang"],
+        exclude = ["%{rocm_root}/llvm/bin/clang-*"],
+    ),
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "clang++",
+    srcs = glob(["%{rocm_root}/llvm/bin/clang++"]),
+    visibility = ["//visibility:public"],
+)
+
+# All clang binaries (for packaging)
+filegroup(
+    name = "clang_all",
+    srcs = glob([
+        "%{rocm_root}/llvm/bin/clang",
+        "%{rocm_root}/llvm/bin/clang-*",
+    ]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM linker
+filegroup(
+    name = "ld",
+    srcs = glob(["%{rocm_root}/llvm/bin/ld.lld"]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM archiver
+filegroup(
+    name = "ar",
+    srcs = glob(["%{rocm_root}/llvm/bin/llvm-ar"]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM strip
+filegroup(
+    name = "strip",
+    srcs = glob(["%{rocm_root}/llvm/bin/llvm-strip"]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM compiler includes (raw filegroup)
+filegroup(
+    name = "compiler_incs_files",
+    srcs = glob(["%{rocm_root}/llvm/lib/clang/*/include/**"]),
+)
+
+# Wrapped for cc_toolchain_import
+cc_toolchain_import(
+    name = "compiler_incs",
+    hdrs = [":compiler_incs_files"],
+    includes = glob(["%{rocm_root}/llvm/lib/clang/*/include"], exclude_directories = 0),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm llvm-symbolizer for sanitizer stack trace symbolization
+filegroup(
+    name = "llvm-symbolizer",
+    srcs = glob(["%{rocm_root}/llvm/bin/llvm-symbolizer"]),
+    visibility = ["//visibility:public"],
+)
+
+# Distribution libraries needed by llvm-symbolizer
+filegroup(
+    name = "distro_libs",
+    srcs = glob(["%{rocm_root}/llvm/lib/*.so*"]),
+    visibility = ["//visibility:public"],
+)
+
+

gpu/rocm/rocm_dist.BUILD.tpl

@@ -1,5 +1,7 @@
 # Hermetic ROCm distribution
 
+load("@rules_ml_toolchain//third_party/rules_cc_toolchain/features:cc_toolchain_import.bzl", "cc_toolchain_import")
+
 # Export the distribution directory so it can be symlinked by other repositories
 exports_files(["%{rocm_root}"], visibility = ["//visibility:public"])
 
@@ -62,3 +64,80 @@ filegroup(
     ]),
     visibility = ["//visibility:public"],
 )
+
+# ROCm clang compiler (single file for use in attributes)
+filegroup(
+    name = "clang",
+    srcs = glob(
+        ["%{rocm_root}/llvm/bin/clang"],
+        exclude = ["%{rocm_root}/llvm/bin/clang-*"],
+    ),
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "clang++",
+    srcs = glob(["%{rocm_root}/llvm/bin/clang++"]),
+    visibility = ["//visibility:public"],
+)
+
+# All clang binaries (for packaging)
+filegroup(
+    name = "clang_all",
+    srcs = glob([
+        "%{rocm_root}/llvm/bin/clang",
+        "%{rocm_root}/llvm/bin/clang-*",
+    ]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM linker
+filegroup(
+    name = "ld",
+    srcs = glob(["%{rocm_root}/llvm/bin/ld.lld"]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM archiver
+filegroup(
+    name = "ar",
+    srcs = glob(["%{rocm_root}/llvm/bin/llvm-ar"]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM strip
+filegroup(
+    name = "strip",
+    srcs = glob(["%{rocm_root}/llvm/bin/llvm-strip"]),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm LLVM compiler includes (raw filegroup)
+filegroup(
+    name = "compiler_incs_files",
+    srcs = glob([
+        "%{rocm_root}/llvm/lib/clang/*/include/**",
+    ]),
+)
+
+# Wrapped for cc_toolchain_import
+cc_toolchain_import(
+    name = "compiler_incs",
+    hdrs = [":compiler_incs_files"],
+    includes = glob(["%{rocm_root}/llvm/lib/clang/*/include"], exclude_directories = 0),
+    visibility = ["//visibility:public"],
+)
+
+# ROCm llvm-symbolizer for sanitizer stack trace symbolization
+filegroup(
+    name = "llvm-symbolizer",
+    srcs = glob([
+        "%{rocm_root}/llvm/bin/llvm-symbolizer",
+        "%{rocm_root}/lib/llvm/bin/llvm-symbolizer",
+    ]),
+    visibility = ["//visibility:public"],
+)
+
+# Distribution libraries needed by llvm-symbolizer (already defined above as distro_libs)
+
+