Adding layering check toolchain features

cc/gen_crosstool_cppmap.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+# Copyright 2025 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.
+# ==============================================================================
+#
+# This tool creates a crosstool.cppmap Clang header modules map file for the toolchain.
+# 
+# Usage:
+# cc/gen_crosstool_cppmap.sh //cc/impls/linux_x86_64_linux_x86_64_cuda:toolchain > cc/impls/linux_x86_64_linux_x86_64_cuda/crosstool.cppmap
+
+readonly TOOLCHAIN_TARGET=$1
+shift
+
+
+#echo '# This file is auto-generated'
+echo 'module "crosstool" [system] {'
+# This will output more files than we actually need, but for our purpose of using the header module map 
+# for the layering check that doesn't matter.
+bazel cquery --output=starlark '--starlark:expr="\n".join([f.path for f in target.files.to_list() if len(f.owner.workspace_name) == 0])' $TOOLCHAIN_TARGET \
+  | grep -E '\.h$' | while read f; do
+  # These are the files that live in the rules_ml_toolchain repo. We need to prefix them with the path where they
+  # are being mapped to in the user's repo.
+  echo "  textual header \"external/rules_ml_toolchain/$f\""
+done
+bazel cquery --output=starlark '--starlark:expr="\n".join([f.path for f in target.files.to_list() if len(f.owner.workspace_name) > 0])' $TOOLCHAIN_TARGET \
+  | grep -E '/include/|cc/cuda/' | while read f; do
+  echo "  textual header \"$f\""
+done
+echo '}'

cc/impls/linux_x86_64_linux_x86_64/BUILD

@@ -190,4 +190,7 @@ cc_toolchain(
     supports_param_files = 1,
     toolchain_config = ":config",
     toolchain_identifier = "toolchain_linux_x86_64_linux_x86_64_id",
+    # The crosstool.cppmap needs to be updated when the compiler or the sysroot gets updated.
+    # You can use the cc/gen_crosstool_cppmap.sh tool for that.
+    module_map = "crosstool.cppmap",
 )

cc/impls/linux_x86_64_linux_x86_64_cuda/BUILD

@@ -248,4 +248,7 @@ cc_toolchain(
     supports_param_files = 1,
     toolchain_config = ":config",
     toolchain_identifier = "toolchain_linux_x86_64_linux_x86_64_id",
+    # The crosstool.cppmap needs to be updated when the compiler or the sysroot gets updated.
+    # You can use the cc/gen_crosstool_cppmap.sh tool for that.
+    module_map = "crosstool.cppmap",
 )

cc/update_all_crosstool_cppmaps.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright 2025 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.
+# ==============================================================================
+#
+# This tool queries all registered toolchains under //cc/impls and calls
+# cc/gen_crosstool_cppmap.sh for each of them.
+#
+# Note that this requires that all toolchains need to be configurable on the system
+# where this command is getting executed since gen_crosstool_cppmap.sh calls bazel cquery.
+#
+# Usage: cc/update_all_crosstool_cppmaps.sh
+
+
+set -euo pipefail
+readonly SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
+bazel query 'kind(cc_toolchain, //cc/impls/...)' | grep :toolchain | while read target; do 
+  crosstool_cppmap="${SCRIPT_DIR}/..${target/\//}"
+  crosstool_cppmap="${crosstool_cppmap/:toolchain//crosstool.cppmap}"
+  ${SCRIPT_DIR}/gen_crosstool_cppmap.sh $target >${crosstool_cppmap}
+done

third_party/gpus/cuda/hermetic/BUILD.tpl

@@ -49,17 +49,35 @@ config_setting(
 # All clients including TensorFlow should use these directives.
 cc_library(
     name = "cuda_headers",
-    hdrs = [
-        "cuda/cuda_config.h",
-    ],
+    hdrs = ["cuda/cuda_config.h"],
     include_prefix = "third_party/gpus",
     includes = [
         ".",  # required to include cuda/cuda/cuda_config.h as cuda/config.h
     ],
+    # For the layering check to work we need to re-export all the headers from
+    # the subtargets. We still also need them in the dependencies for setting
+    # up the right include paths with the include prefixes.
+    textual_hdrs = [
+            "@cuda_cudart//:header_list",
+            "@cuda_cublas//:header_list",
+            "@cuda_profiler_api//:header_list",
+            "@cuda_cccl//:header_list",
+            "@cuda_nvrtc//:header_list",
+            "@cuda_nvtx//:header_list",
+            "@cuda_nvcc//:header_list",
+            "@cuda_cusolver//:header_list",
+            "@cuda_cufft//:header_list",
+            "@cuda_cusparse//:header_list",
+            "@cuda_curand//:header_list",
+            "@cuda_cupti//:header_list",
+            "@cuda_nvml//:header_list",
+            "@cuda_nvjitlink//:header_list",
+           ] + (["@cuda_crt://header_list"] if _cudart_version and int(_cudart_version)>=13 else []),
     deps = [":cudart_headers",
             ":cublas_headers",
             ":profiler_api_headers",
             ":cccl_headers",
+            ":nvrtc_headers",
             ":nvtx_headers",
             ":nvcc_headers",
             ":cusolver_headers",

third_party/gpus/cuda/hermetic/cuda_cccl.BUILD.tpl

@@ -5,14 +5,21 @@ load(
 )
 load("@cuda_cudart//:version.bzl", _cudart_version = "VERSION")
 
-cc_library(
-    name = "headers",
-    hdrs = glob([
+filegroup(
+    name = "header_list",
+    srcs = glob([
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/cub/**",
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/cuda/**",
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/nv/**",
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/thrust/**",
     ]),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "")],
     strip_include_prefix = "include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", ""),

third_party/gpus/cuda/hermetic/cuda_crt.BUILD.tpl

@@ -3,13 +3,19 @@ licenses(["restricted"])  # NVIDIA proprietary license
 load("@cuda_cudart//:version.bzl", _cudart_version = "VERSION")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_newer_than")
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = if_cuda_newer_than(
+filegroup(
+    name = "header_list",
+    %{comment}srcs = if_cuda_newer_than(
         %{comment}"13_0",
         %{comment}if_true = glob(["include/crt/**"]),
         %{comment}if_false = [],
     %{comment}),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cublas.BUILD.tpl

@@ -44,14 +44,21 @@ cc_library(
     %{comment}),
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = [
+filegroup(
+    name = "header_list",
+    %{comment}srcs = [
         %{comment}"include/cublas.h",
         %{comment}"include/cublasLt.h",
         %{comment}"include/cublas_api.h",
         %{comment}"include/cublas_v2.h",
     %{comment}],
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cudart.BUILD.tpl

@@ -55,9 +55,9 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = glob([
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob([
         %{comment}"include/builtin_types.h",
         %{comment}"include/channel_descriptor.h",
         %{comment}"include/common_functions.h",
@@ -140,6 +140,12 @@ cc_library(
         %{comment}"include/vector_functions.hpp",
         %{comment}"include/vector_types.h",
     %{comment}]),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+cc_library(
+    name = "headers",
+    hdrs = [ ":header_list" ],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cudnn.BUILD.tpl

@@ -5,25 +5,25 @@ load(
 )
 
 %{multiline_comment}
-cc_import( 
+cc_import(
     name = "cudnn_ops",
     hdrs = [":headers"],
     shared_library = "lib/libcudnn_ops.so.%{libcudnn_ops_version}",
 )
 
-cc_import( 
+cc_import(
     name = "cudnn_cnn",
     hdrs = [":headers"],
     shared_library = "lib/libcudnn_cnn.so.%{libcudnn_cnn_version}",
 )
 
-cc_import( 
+cc_import(
     name = "cudnn_adv",
     hdrs = [":headers"],
     shared_library = "lib/libcudnn_adv.so.%{libcudnn_adv_version}",
 )
 
-cc_import( 
+cc_import(
     name = "cudnn_graph",
     hdrs = [":headers"],
     shared_library = "lib/libcudnn_graph.so.%{libcudnn_graph_version}",
@@ -55,6 +55,7 @@ cc_import(
 %{multiline_comment}
 cc_library(
     name = "cudnn",
+    hdrs = [":header_list"],
     %{comment}deps = [
       %{comment}":cudnn_engines_precompiled",
       %{comment}":cudnn_ops",
@@ -70,11 +71,17 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = glob([
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob([
         %{comment}"include/cudnn*.h",
     %{comment}]),
+)
+
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cudnn",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cudnn8.BUILD.tpl

@@ -63,11 +63,17 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = glob([
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob([
         %{comment}"include/cudnn*.h",
     %{comment}]),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cudnn",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cufft.BUILD.tpl

@@ -26,12 +26,18 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = glob([
-        %{comment}"include/cudalibxt.h", 
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob([
+        %{comment}"include/cudalibxt.h",
         %{comment}"include/cufft*.h"
     %{comment}]),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cupti.BUILD.tpl

@@ -27,9 +27,9 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = glob([
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob([
         %{comment}"include/Openacc/**",
         %{comment}"include/Openmp/**",
         %{comment}"include/cuda_stdint.h",
@@ -72,6 +72,12 @@ cc_library(
         %{comment}["include/cupti_pmsampling.h",
         %{comment}"include/cupti_profiler_host.h"],
     %{comment}),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/extras/CUPTI/include",
     includes = ["include/"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_curand.BUILD.tpl

@@ -26,9 +26,15 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob(["include/curand*.h"]),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
 cc_library(
     name = "headers",
-    %{comment}hdrs = glob(["include/curand*.h"]),
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include"],
     strip_include_prefix = "include",

third_party/gpus/cuda/hermetic/cuda_cusolver.BUILD.tpl

@@ -32,11 +32,18 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
-cc_library(
-    name = "headers",
-    %{comment}hdrs = glob([
+
+filegroup(
+    name = "header_list",
+    %{comment}srcs = glob([
         %{comment}"include/cusolver*.h",
     %{comment}]),
+    visibility = ["@local_config_cuda//cuda:__pkg__"],
+)
+
+cc_library(
+    name = "headers",
+    hdrs = [":header_list"],
     include_prefix = "third_party/gpus/cuda/include",
     includes = ["include"],
     strip_include_prefix = "include",