Skip to content

Commit

Permalink
We use PYTORCH_ROCM_ARCH env variable as the target list first, if it…
Browse files Browse the repository at this point in the history
…'s not set, then we default to rocm_agent_enumerator for the target list
  • Loading branch information
Alexey Kondratiev committed Apr 8, 2024
1 parent 71e9dc2 commit d8985ef
Showing 1 changed file with 11 additions and 4 deletions.
15 changes: 11 additions & 4 deletions cmake/utils.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -139,16 +139,23 @@ macro(override_gpu_arches GPU_ARCHES GPU_LANG GPU_SUPPORTED_ARCHES)
if (${GPU_LANG} STREQUAL "HIP")
#
# `GPU_ARCHES` controls the `--offload-arch` flags.
# `CMAKE_HIP_ARCHITECTURES` is set up by torch and can be controlled
# via the `PYTORCH_ROCM_ARCH` env variable.
#

# If PYTORCH_ROCM_ARCH env variable exsists, then we take it as a list,
# if not, then we use CMAKE_HIP_ARCHITECTURES which was generated by calling
# "rocm_agent_enumerator" in "enable_language(HIP)"
# (in file Modules/CMakeDetermineHIPCompiler.cmake)
#
if(DEFINED ENV{PYTORCH_ROCM_ARCH})
set(HIP_ARCHITECTURES $ENV{PYTORCH_ROCM_ARCH})
else()
set(HIP_ARCHITECTURES ${CMAKE_HIP_ARCHITECTURES})
endif()
#
# Find the intersection of the supported + detected architectures to
# set the module architecture flags.
#
set(${GPU_ARCHES})
foreach (_ARCH ${CMAKE_HIP_ARCHITECTURES})
foreach (_ARCH ${HIP_ARCHITECTURES})
if (_ARCH IN_LIST _GPU_SUPPORTED_ARCHES_LIST)
list(APPEND ${GPU_ARCHES} ${_ARCH})
endif()
Expand Down

0 comments on commit d8985ef

Please sign in to comment.