mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
a38b572de3
Summary: This PR for the ROCm target does the following: * enable some unit tests on ROCm * fix a missing static_cast that breaks BatchNorm call on ROCm * fix BatchNorm to work on ROCm w/ ROCm warp sizes etc * improve the pyhipify script by introducing kernel scope to some transpilations and other improvements * fix a linking issue on ROCm * for more unit test sets: mark currently broken tests broken (to be fixed) * enable THINLTO (phase one) to parallelize linking * address the first failing of the elementwise kernel by removing non-working ROCm specialization Pull Request resolved: https://github.com/pytorch/pytorch/pull/10266 Differential Revision: D9184178 Pulled By: ezyang fbshipit-source-id: 03bcd1fe4ca4dd3241f09634dbd42b6a4c350297
16 lines
409 B
Python
16 lines
409 B
Python
import os
|
|
from .env import check_env_flag, check_negative_env_flag
|
|
|
|
# Get ROCm Home Path
|
|
ROCM_HOME = os.getenv("ROCM_HOME", "/opt/rocm")
|
|
ROCM_VERSION = ""
|
|
USE_ROCM = False
|
|
|
|
# Check if ROCm disabled.
|
|
if check_negative_env_flag("USE_ROCM"):
|
|
USE_ROCM = False
|
|
else:
|
|
# If ROCM home exists or we explicitly enable ROCm
|
|
if os.path.exists(ROCM_HOME) or check_env_flag('USE_ROCM'):
|
|
USE_ROCM = True
|