mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[ROCm][Windows] Fix building torch 2.8 wheel with ROCm (added hipblasLt and rocblas directories) (#153144)
Since rocblas.dll and hipblaslt.dll are copied to torch/lib, rocblas and hipblaslt directories are needed to be stored there too (otherwise we have an error after wheel installation while searching for files in rocblas/library and hipblaslt/library which doesn't exist). This PR fixes this issue. Pull Request resolved: https://github.com/pytorch/pytorch/pull/153144 Approved by: https://github.com/jeffdaily Co-authored-by: Jeff Daily <jeff.daily@amd.com>
This commit is contained in:
committed by
PyTorch MergeBot
parent
fa6ca59079
commit
b367e5f6a6
19
setup.py
19
setup.py
@ -748,6 +748,25 @@ class build_ext(setuptools.command.build_ext.build_ext):
|
||||
|
||||
self.copy_file(export_lib, target_lib)
|
||||
|
||||
# In ROCm on Windows case copy rocblas and hipblaslt files into
|
||||
# torch/lib/rocblas/library and torch/lib/hipblaslt/library
|
||||
use_rocm = os.environ.get("USE_ROCM")
|
||||
if use_rocm:
|
||||
rocm_dir_path = os.environ.get("ROCM_DIR")
|
||||
rocm_bin_path = os.path.join(rocm_dir_path, "bin")
|
||||
|
||||
rocblas_dir = os.path.join(rocm_bin_path, "rocblas")
|
||||
target_rocblas_dir = os.path.join(target_dir, "rocblas")
|
||||
os.makedirs(target_rocblas_dir, exist_ok=True)
|
||||
self.copy_tree(rocblas_dir, target_rocblas_dir)
|
||||
|
||||
hipblaslt_dir = os.path.join(rocm_bin_path, "hipblaslt")
|
||||
target_hipblaslt_dir = os.path.join(target_dir, "hipblaslt")
|
||||
os.makedirs(target_hipblaslt_dir, exist_ok=True)
|
||||
self.copy_tree(hipblaslt_dir, target_hipblaslt_dir)
|
||||
else:
|
||||
report("The specified environment variable does not exist.")
|
||||
|
||||
def build_extensions(self):
|
||||
self.create_compile_commands()
|
||||
|
||||
|
Reference in New Issue
Block a user