Skip to content

Commit 6c33b40

Browse files
Extend build_locally.py with --target-cuda to simplify passing CUDA target
1 parent 716d3ff commit 6c33b40

File tree

1 file changed

+20
-0
lines changed

1 file changed

+20
-0
lines changed

scripts/build_locally.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ def run(
3030
use_glog=False,
3131
verbose=False,
3232
cmake_opts="",
33+
target_cuda=None,
3334
):
3435
build_system = None
3536

@@ -65,6 +66,15 @@ def run(
6566
]
6667
if cmake_opts:
6768
cmake_args += cmake_opts.split()
69+
if target_cuda is not None:
70+
if not target_cuda.strip():
71+
raise ValueError(
72+
"--target-cuda can not be an empty string. "
73+
"Use --target-cuda=<arch> or --target-cuda"
74+
)
75+
cmake_args += [
76+
f"-DDPCTL_TARGET_CUDA={target_cuda}",
77+
]
6878
subprocess.check_call(
6979
cmake_args, shell=False, cwd=setup_dir, env=os.environ
7080
)
@@ -131,6 +141,15 @@ def run(
131141
default="",
132142
type=str,
133143
)
144+
driver.add_argument(
145+
"--target-cuda",
146+
nargs="?",
147+
const="ON",
148+
help="Enable CUDA target for build; "
149+
"optionally specify architecture (e.g., --target-cuda=sm_80)",
150+
default=None,
151+
type=str,
152+
)
134153
args = parser.parse_args()
135154

136155
args_to_validate = [
@@ -186,4 +205,5 @@ def run(
186205
use_glog=args.glog,
187206
verbose=args.verbose,
188207
cmake_opts=args.cmake_opts,
208+
target_cuda=args.target_cuda,
189209
)

0 commit comments

Comments
 (0)