diff --git a/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl b/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
index ffa305c77..8ad22502a 100644
--- a/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
+++ b/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
@@ -24,7 +24,7 @@ def all_assembly_actions():
     ]
 
 def all_compile_actions():
-    return [
+    retur [
         ACTION_NAMES.assemble,
         ACTION_NAMES.c_compile,
         ACTION_NAMES.cpp_compile,
@@ -300,6 +300,45 @@ def _features(cpu, compiler, ctx):
     if cpu in ["local", "darwin"]:
         return [
             feature(name = "no_legacy_features"),
+            feature(
+             name = "use_module_maps",
+             requires = [feature_set(features = ["module_maps"])],
+             flag_sets = [
+                 flag_set(
+                     actions = all_compile_actions(),
+                     flag_groups = [
+                         flag_group(
+                             flags = [
+                                 "-fmodule-name=%{module_name}",
+                                 "-fmodule-map-file=%{module_map_file}",
+                             ],
+                         ),
+                     ],
+                 ),
+             ],
+           ),
+           feature(name = "module_maps", enabled = True),
+           feature(
+               name = "layering_check",
+               implies = ["use_module_maps"],
+               flag_sets = [
+                   flag_set(
+                       actions = all_compile_actions(),
+                       flag_groups = [
+                           flag_group(flags = [
+                               "-fmodules-strict-decluse"
+                               "-Wprivate-header",
+                           ]),
+                           flag_group(
+                               iterate_over = "dependent_module_map_files",
+                               flags = [
+                                   "-fmodule-map-file=%{dependent_module_map_files}",
+                               ],
+                           ),
+                       ],
+                   ),
+               ],
+           ),
             feature(
                 name = "all_compile_flags",
                 enabled = True,
diff --git a/xla/literal.cc b/xla/literal.cc
index 2cef23917..d362f8af9 100644
--- a/xla/literal.cc
+++ b/xla/literal.cc
@@ -47,6 +47,7 @@ limitations under the License.
 #include "tsl/platform/errors.h"
 #include "tsl/platform/float8.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/platform.h"
 #include "tsl/platform/mem.h"
 #include "tsl/util/byte_swap_array.h"