File tree Expand file tree Collapse file tree 1 file changed +20
-0
lines changed Expand file tree Collapse file tree 1 file changed +20
-0
lines changed Original file line number Diff line number Diff line change @@ -30,6 +30,7 @@ def run(
30
30
use_glog = False ,
31
31
verbose = False ,
32
32
cmake_opts = "" ,
33
+ target_cuda = None ,
33
34
):
34
35
build_system = None
35
36
@@ -65,6 +66,15 @@ def run(
65
66
]
66
67
if cmake_opts :
67
68
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
+ ]
68
78
subprocess .check_call (
69
79
cmake_args , shell = False , cwd = setup_dir , env = os .environ
70
80
)
@@ -131,6 +141,15 @@ def run(
131
141
default = "" ,
132
142
type = str ,
133
143
)
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
+ )
134
153
args = parser .parse_args ()
135
154
136
155
args_to_validate = [
@@ -186,4 +205,5 @@ def run(
186
205
use_glog = args .glog ,
187
206
verbose = args .verbose ,
188
207
cmake_opts = args .cmake_opts ,
208
+ target_cuda = args .target_cuda ,
189
209
)
You can’t perform that action at this time.
0 commit comments