Skip to content

Commit

Permalink
Skeleton code for nvcc compilation planner
Browse files Browse the repository at this point in the history
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
  • Loading branch information
zhuhan0 authored and facebook-github-bot committed Feb 22, 2025
1 parent d0089ab commit ddb9a4a
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 9 deletions.
30 changes: 21 additions & 9 deletions prelude/cxx/compile.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ load(
"HeaderExtension",
"HeadersDepFiles",
)
load("@prelude//cxx:cuda.bzl", "cuda_compile")
load("@prelude//cxx:cxx_toolchain_types.bzl", "CxxToolchainInfo")
load("@prelude//cxx:cxx_utility.bzl", "cxx_attrs_get_allow_cache_upload")
load(
Expand Down Expand Up @@ -385,15 +386,26 @@ def _compile_single_cxx(
)
cmd.add(cmd_args(external_debug_info.as_output(), format = "--fbcc-create-external-debug-info={}"))

ctx.actions.run(
cmd,
category = src_compile_cmd.cxx_compile_cmd.category,
identifier = identifier,
dep_files = action_dep_files,
allow_cache_upload = src_compile_cmd.cxx_compile_cmd.allow_cache_upload,
allow_dep_file_cache_upload = False,
**error_handler_args
)
if src_compile_cmd.src.extension == ".cu":
cuda_compile(
ctx,
cmd,
src_compile_cmd,
identifier,
action_dep_files,
allow_dep_file_cache_upload = False,
error_handler_args = error_handler_args,
)
else:
ctx.actions.run(
cmd,
category = src_compile_cmd.cxx_compile_cmd.category,
identifier = identifier,
dep_files = action_dep_files,
allow_cache_upload = src_compile_cmd.cxx_compile_cmd.allow_cache_upload,
allow_dep_file_cache_upload = False,
**error_handler_args
)

# If we're building with split debugging, where the debug info is in the
# original object, then add the object as external debug info
Expand Down
34 changes: 34 additions & 0 deletions prelude/cxx/cuda.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
# of this source tree.

load("@prelude//cxx:compile_types.bzl", "CxxSrcCompileCommand")

CudaCompileStyle = enum(
# Use NVCC as the compiler driver and compile a CUDA file in a single Buck
# action.
Expand All @@ -13,3 +15,35 @@ CudaCompileStyle = enum(
# sub-command.
"dist",
)

def cuda_compile(
ctx: AnalysisContext,
cmd: cmd_args,
src_compile_cmd: CxxSrcCompileCommand,
identifier: str,
action_dep_files: dict[str, ArtifactTag],
allow_dep_file_cache_upload: bool,
error_handler_args: dict[str, [typing.Callable, None]]):
if ctx.attrs.cuda_compile_style == CudaCompileStyle("mono").value:
ctx.actions.run(
cmd,
category = src_compile_cmd.cxx_compile_cmd.category,
identifier = identifier,
dep_files = action_dep_files,
allow_cache_upload = src_compile_cmd.cxx_compile_cmd.allow_cache_upload,
allow_dep_file_cache_upload = allow_dep_file_cache_upload,
**error_handler_args
)
elif ctx.attrs.cuda_compile_style == CudaCompileStyle("dist").value:
cmd.add("-_NVCC_DRYRUN_")
ctx.actions.run(
cmd,
category = "cuda_compile_prepare",
identifier = identifier,
dep_files = action_dep_files,
allow_cache_upload = True,
allow_dep_file_cache_upload = allow_dep_file_cache_upload,
**error_handler_args
)
else:
fail("Unsupported CUDA compile style: {}".format(ctx.attrs.cuda_compile_style))

0 comments on commit ddb9a4a

Please sign in to comment.