feat: Add features for compiling with -arch=all or -arch=all-major

bazel-contrib · May 31, 2024 · 9e1e7d2 · 9e1e7d2
1 parent b4f9093
commit 9e1e7d2
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 0 deletions.
diff --git a/cuda/private/toolchain_configs/nvcc.bzl b/cuda/private/toolchain_configs/nvcc.bzl
@@ -178,6 +178,42 @@ def _impl(ctx):
         ],
     )
 
+    arch_all_feature = feature(
+        name = "arch_all",
+        enabled = False,
+        flag_sets = [
+            flag_set(
+                actions = [
+                    ACTION_NAMES.cuda_compile,
+                    ACTION_NAMES.device_link,
+                ],
+                flag_groups = [
+                    flag_group(
+                        flags = ["-arch=all"],
+                    ),
+                ],
+            ),
+        ],
+    )
+
+    arch_all_major_feature = feature(
+        name = "arch_all_major",
+        enabled = False,
+        flag_sets = [
+            flag_set(
+                actions = [
+                    ACTION_NAMES.cuda_compile,
+                    ACTION_NAMES.device_link,
+                ],
+                flag_groups = [
+                    flag_group(
+                        flags = ["-arch=all-major"],
+                    ),
+                ],
+            ),
+        ],
+    )
+
     pic_feature = feature(
         name = "pic",
         enabled = True,
@@ -441,6 +477,8 @@ def _impl(ctx):
         nvcc_compile_env_feature,
         nvcc_device_link_env_feature,
         arch_native_feature,
+        arch_all_feature,
+        arch_all_major_feature,
         pic_feature,
         host_compiler_feature,
         supports_compiler_device_link_feature,

diff --git a/cuda/private/toolchain_configs/nvcc_msvc.bzl b/cuda/private/toolchain_configs/nvcc_msvc.bzl
@@ -338,6 +338,42 @@ def _impl(ctx):
         ],
     )
 
+    arch_all_feature = feature(
+        name = "arch_all",
+        enabled = False,
+        flag_sets = [
+            flag_set(
+                actions = [
+                    ACTION_NAMES.cuda_compile,
+                    ACTION_NAMES.device_link,
+                ],
+                flag_groups = [
+                    flag_group(
+                        flags = ["-arch=all"],
+                    ),
+                ],
+            ),
+        ],
+    )
+
+    arch_all_major_feature = feature(
+        name = "arch_all_major",
+        enabled = False,
+        flag_sets = [
+            flag_set(
+                actions = [
+                    ACTION_NAMES.cuda_compile,
+                    ACTION_NAMES.device_link,
+                ],
+                flag_groups = [
+                    flag_group(
+                        flags = ["-arch=all-major"],
+                    ),
+                ],
+            ),
+        ],
+    )
+
     dbg_feature = feature(
         name = "dbg",
         flag_sets = [
@@ -526,6 +562,8 @@ def _impl(ctx):
         supports_compiler_device_link_feature,
         use_local_env_feature,
         arch_native_feature,
+        arch_all_feature,
+        arch_all_major_feature,
         default_compile_flags_feature,
         include_paths_feature,
         defines_feature,