[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:
tvukovic-amd
2025-05-27 19:40:25 +00:00
committed by PyTorch MergeBot
parent fa6ca59079
commit b367e5f6a6

View File

@ -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()