diff --git a/third_party/tsl/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl b/third_party/tsl/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl index 848ab280cb636..43c2a9979570d 100755 --- a/third_party/tsl/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl +++ b/third_party/tsl/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl @@ -34,6 +34,7 @@ HIP_RUNTIME_LIBRARY = '%{hip_runtime_library}' ROCR_RUNTIME_PATH = '%{rocr_runtime_path}' ROCR_RUNTIME_LIBRARY = '%{rocr_runtime_library}' VERBOSE = '%{crosstool_verbose}'=='1' +ROCM_AMDGPU_TARGETS = '%{rocm_amdgpu_targets}' def Log(s): print('gpus/crosstool: {0}'.format(s)) @@ -99,6 +100,29 @@ def GetHostCompilerOptions(argv): return opts +def GetHipccOptions(argv): + """Collect the -hipcc_options values from argv. + + Args: + argv: A list of strings, possibly the argv passed to main(). + + Returns: + The string that can be passed directly to hipcc. + """ + + parser = ArgumentParser() + parser.add_argument('--offload-arch', nargs='*', action='append') + # TODO find a better place for this + parser.add_argument('-gline-tables-only', action='store_true') + + args, _ = parser.parse_known_args(argv) + + hipcc_opts = ' -gline-tables-only ' if args.gline_tables_only else '' + if args.offload_arch: + hipcc_opts = hipcc_opts + ' '.join(['--offload-arch=' + a for a in sum(args.offload_arch, [])]) + + return hipcc_opts + def system(cmd): """Invokes cmd with os.system(). @@ -115,7 +139,6 @@ def system(cmd): else: return -os.WTERMSIG(retv) - def InvokeHipcc(argv, log=False): """Call hipcc with arguments assembled from argv. @@ -128,6 +151,7 @@ def InvokeHipcc(argv, log=False): """ host_compiler_options = GetHostCompilerOptions(argv) + hipcc_compiler_options = GetHipccOptions(argv) opt_option = GetOptionValue(argv, 'O') m_options = GetOptionValue(argv, 'm') m_options = ''.join([' -m' + m for m in m_options if m in ['32', '64']]) @@ -166,7 +190,7 @@ def InvokeHipcc(argv, log=False): srcs = ' '.join(src_files) out = ' -o ' + out_file[0] - hipccopts = ' ' + hipccopts = hipcc_compiler_options + ' ' # In hip-clang environment, we need to make sure that hip header is included # before some standard math header like is included in any source. # Otherwise, we get build error. @@ -222,6 +246,7 @@ def main(): if VERBOSE: print('PWD=' + os.getcwd()) if VERBOSE: print('HIPCC_ENV=' + HIPCC_ENV) + if VERBOSE: print('ROCM_AMDGPU_TARGETS= ' + ROCM_AMDGPU_TARGETS) if args.x and args.x[0] == 'rocm': # compilation for GPU objects diff --git a/third_party/tsl/third_party/gpus/rocm_configure.bzl b/third_party/tsl/third_party/gpus/rocm_configure.bzl index e23bc2185f63f..963cfb2a1bb38 100644 --- a/third_party/tsl/third_party/gpus/rocm_configure.bzl +++ b/third_party/tsl/third_party/gpus/rocm_configure.bzl @@ -798,6 +798,9 @@ def _create_local_rocm_repository(repository_ctx): "%{hip_runtime_library}": "amdhip64", "%{crosstool_verbose}": _crosstool_verbose(repository_ctx), "%{gcc_host_compiler_path}": str(cc), + "%{rocm_amdgpu_targets}": ",".join( + ["\"%s\"" % c for c in rocm_config.amdgpu_targets], + ), }, )