Skip to content

Commit ddb9a4a

Browse files
zhuhan0facebook-github-bot
authored andcommitted
Skeleton code for nvcc compilation planner
Summary: Add a `-_NVCC_DRYRUN_` option to wrap_nvcc.py. If provided, wrap_nvcc.py will generate several files related to how the CUDA file will be compiled (to be implemented): * A plain dump of the `nvcc -dryrun` output. * A list of environment variables for the compilation sub-commands to run with. * A dependency graph of the compilation sub-commands. This diff adds the skeleton API for this `nvcc_compilation_plan`. This currently runs in addition to the regular monolithic CUDA compilation, so that we still produce a valid C++ output object when `_NVCC_DRYRUN_` is passed. Once the implementation is complete, I'll stop running the mono CUDA compilation in this branch. Reviewed By: get9 Differential Revision: D69757705 fbshipit-source-id: 5f5c2dc2074c1d9dc6debc20545aaa6d7dcef464
1 parent d0089ab commit ddb9a4a

File tree

2 files changed

+55
-9
lines changed

2 files changed

+55
-9
lines changed

prelude/cxx/compile.bzl

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ load(
2020
"HeaderExtension",
2121
"HeadersDepFiles",
2222
)
23+
load("@prelude//cxx:cuda.bzl", "cuda_compile")
2324
load("@prelude//cxx:cxx_toolchain_types.bzl", "CxxToolchainInfo")
2425
load("@prelude//cxx:cxx_utility.bzl", "cxx_attrs_get_allow_cache_upload")
2526
load(
@@ -385,15 +386,26 @@ def _compile_single_cxx(
385386
)
386387
cmd.add(cmd_args(external_debug_info.as_output(), format = "--fbcc-create-external-debug-info={}"))
387388

388-
ctx.actions.run(
389-
cmd,
390-
category = src_compile_cmd.cxx_compile_cmd.category,
391-
identifier = identifier,
392-
dep_files = action_dep_files,
393-
allow_cache_upload = src_compile_cmd.cxx_compile_cmd.allow_cache_upload,
394-
allow_dep_file_cache_upload = False,
395-
**error_handler_args
396-
)
389+
if src_compile_cmd.src.extension == ".cu":
390+
cuda_compile(
391+
ctx,
392+
cmd,
393+
src_compile_cmd,
394+
identifier,
395+
action_dep_files,
396+
allow_dep_file_cache_upload = False,
397+
error_handler_args = error_handler_args,
398+
)
399+
else:
400+
ctx.actions.run(
401+
cmd,
402+
category = src_compile_cmd.cxx_compile_cmd.category,
403+
identifier = identifier,
404+
dep_files = action_dep_files,
405+
allow_cache_upload = src_compile_cmd.cxx_compile_cmd.allow_cache_upload,
406+
allow_dep_file_cache_upload = False,
407+
**error_handler_args
408+
)
397409

398410
# If we're building with split debugging, where the debug info is in the
399411
# original object, then add the object as external debug info

prelude/cxx/cuda.bzl

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
66
# of this source tree.
77

8+
load("@prelude//cxx:compile_types.bzl", "CxxSrcCompileCommand")
9+
810
CudaCompileStyle = enum(
911
# Use NVCC as the compiler driver and compile a CUDA file in a single Buck
1012
# action.
@@ -13,3 +15,35 @@ CudaCompileStyle = enum(
1315
# sub-command.
1416
"dist",
1517
)
18+
19+
def cuda_compile(
20+
ctx: AnalysisContext,
21+
cmd: cmd_args,
22+
src_compile_cmd: CxxSrcCompileCommand,
23+
identifier: str,
24+
action_dep_files: dict[str, ArtifactTag],
25+
allow_dep_file_cache_upload: bool,
26+
error_handler_args: dict[str, [typing.Callable, None]]):
27+
if ctx.attrs.cuda_compile_style == CudaCompileStyle("mono").value:
28+
ctx.actions.run(
29+
cmd,
30+
category = src_compile_cmd.cxx_compile_cmd.category,
31+
identifier = identifier,
32+
dep_files = action_dep_files,
33+
allow_cache_upload = src_compile_cmd.cxx_compile_cmd.allow_cache_upload,
34+
allow_dep_file_cache_upload = allow_dep_file_cache_upload,
35+
**error_handler_args
36+
)
37+
elif ctx.attrs.cuda_compile_style == CudaCompileStyle("dist").value:
38+
cmd.add("-_NVCC_DRYRUN_")
39+
ctx.actions.run(
40+
cmd,
41+
category = "cuda_compile_prepare",
42+
identifier = identifier,
43+
dep_files = action_dep_files,
44+
allow_cache_upload = True,
45+
allow_dep_file_cache_upload = allow_dep_file_cache_upload,
46+
**error_handler_args
47+
)
48+
else:
49+
fail("Unsupported CUDA compile style: {}".format(ctx.attrs.cuda_compile_style))

0 commit comments

Comments
 (0)