Skip to content

Commit

Permalink
Refactor PyTorch wheel and libtorch build scripts for ROCm (#1232)
Browse files Browse the repository at this point in the history
* Refactor wheel and libtorch build scripts (ROCm#7)

* Update to so patching for ROCm

Wildcard used in grep to grab the actual numbered so file referenced
in patchelf. This allows the removal of specifying the so number in
DEPS_LIST & DEPS_SONAME

This commit also adds the functionality for trimming so names to
build_libtorch.sh from build_common.sh

* Refactor to remove switch statement in build_rocm.sh

This commit refactors build_rocm.sh and brings in a few major updates:
 - No longer required to specify the full .so name (with number) for ROCm libraries
       - The .so versions are copied and the patching code will fix the links to point to this version
 - No longer required to specify paths for ROCm libraries allowing the removal of the large switch
       - Paths are acquired programmatically with find
 - No longer required to specify both the path and filename for the OS specific libraries
       - Programatically extract file name from the path
 - Automatically extract Tensile/Kernels files for the architectures specified in PYTORCH_ROCM_ARCH
   and any non-arch specific files e.g. TensileLibrary.dat

* rocfft/hipfft link to libhiprtc.so in ROCm5.4 (ROCm#15)

Co-authored-by: Jack Taylor <108682042+jataylo@users.noreply.github.com>
  • Loading branch information
jithunnair-amd and jataylo committed Jan 12, 2023
1 parent d168714 commit 0b51144
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 565 deletions.
4 changes: 2 additions & 2 deletions manywheel/build_common.sh
Expand Up @@ -279,9 +279,9 @@ replace_needed_sofiles() {
find $1 -name '*.so*' | while read sofile; do
origname=$2
patchedname=$3
if [[ "$origname" != "$patchedname" ]]; then
if [[ "$origname" != "$patchedname" ]] || [[ "$DESIRED_CUDA" == *"rocm"* ]]; then
set +e
$PATCHELF_BIN --print-needed $sofile | grep $origname 2>&1 >/dev/null
origname=$($PATCHELF_BIN --print-needed $sofile | grep "$origname*")
ERRCODE=$?
set -e
if [ "$ERRCODE" -eq "0" ]; then
Expand Down
17 changes: 13 additions & 4 deletions manywheel/build_libtorch.sh
Expand Up @@ -228,6 +228,11 @@ fname_with_sha256() {
fi
}

fname_without_so_number() {
LINKNAME=$(echo $1 | sed -e 's/\.so.*/.so/g')
echo "$LINKNAME"
}

make_wheel_record() {
FPATH=$1
if echo $FPATH | grep RECORD >/dev/null 2>&1; then
Expand Down Expand Up @@ -277,8 +282,12 @@ for pkg in /$LIBTORCH_HOUSE_DIR/libtorch*.zip; do
if [[ "$filepath" != "$destpath" ]]; then
cp $filepath $destpath
fi

patchedpath=$(fname_with_sha256 $destpath)

if [[ "$DESIRED_CUDA" == *"rocm"* ]]; then
patchedpath=$(fname_without_so_number $destpath)
else
patchedpath=$(fname_with_sha256 $destpath)
fi
patchedname=$(basename $patchedpath)
if [[ "$destpath" != "$patchedpath" ]]; then
mv $destpath $patchedpath
Expand All @@ -292,9 +301,9 @@ for pkg in /$LIBTORCH_HOUSE_DIR/libtorch*.zip; do
find $PREFIX -name '*.so*' | while read sofile; do
origname=${DEPS_SONAME[i]}
patchedname=${patched[i]}
if [[ "$origname" != "$patchedname" ]]; then
if [[ "$origname" != "$patchedname" ]] || [[ "$DESIRED_CUDA" == *"rocm"* ]]; then
set +e
$PATCHELF_BIN --print-needed $sofile | grep $origname 2>&1 >/dev/null
origname=$($PATCHELF_BIN --print-needed $sofile | grep "$origname*")
ERRCODE=$?
set -e
if [ "$ERRCODE" -eq "0" ]; then
Expand Down

0 comments on commit 0b51144

Please sign in to comment.