[ROCm] Use AOTriton as a dynamic library (#1893)

jithunnair-amd · web-flow · commit decc1a16f8a8 · 2024-07-03T13:39:09.000-04:00
* Add AOTriton dependency

* Update aotriton install scripts to download tarball with shared object
diff --git a/common/aotriton_version.txt b/common/aotriton_version.txt
@@ -1,5 +1,5 @@
 0.6b
 manylinux_2_17
-rocm6
+rocm6.1
 04b5df8c8123f90cba3ede7e971e6fbc6040d506
-3db6ecbc915893ff967abd6e1b43bd5f54949868873be60dc802086c3863e648
+77c29fa3f3b614e187d7213d745e989a92708cee2bc6020419ab49019af399d1
diff --git a/common/install_aotriton.sh b/common/install_aotriton.sh
@@ -7,7 +7,7 @@ TARBALL='aotriton.tar.bz2'
 read -d "\n" VER MANYLINUX ROCMBASE PINNED_COMMIT SHA256 < aotriton_version.txt || true
 ARCH=$(uname -m)
 AOTRITON_INSTALL_PREFIX="$1"
-AOTRITON_URL="https://github.com/ROCm/aotriton/releases/download/${VER}/aotriton-${VER}-${MANYLINUX}_${ARCH}-${ROCMBASE}.tar.bz2"
+AOTRITON_URL="https://github.com/ROCm/aotriton/releases/download/${VER}/aotriton-${VER}-${MANYLINUX}_${ARCH}-${ROCMBASE}-shared.tar.bz2"
 
 cd "${AOTRITON_INSTALL_PREFIX}"
 # Must use -L to follow redirects
diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
@@ -184,6 +184,12 @@ do
     OS_SO_FILES[${#OS_SO_FILES[@]}]=$file_name # Append lib to array
 done
 
+# PyTorch-version specific
+# AOTriton dependency only for PyTorch >= 2.4
+if (( $(echo "${PYTORCH_VERSION} 2.4" | awk '{print ($1 >= $2)}') )); then
+    ROCM_SO_FILES+=("libaotriton_v2.so")
+fi
+
 # rocBLAS library files
 if [[ $ROCM_INT -ge 50200 ]]; then
     ROCBLAS_LIB_SRC=$ROCM_HOME/lib/rocblas/library

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ TARBALL='aotriton.tar.bz2'`
`7`	`7`	`read -d "\n" VER MANYLINUX ROCMBASE PINNED_COMMIT SHA256 < aotriton_version.txt \|\| true`
`8`	`8`	`ARCH=$(uname -m)`
`9`	`9`	`AOTRITON_INSTALL_PREFIX="$1"`
`10`		`-AOTRITON_URL="https://github.com/ROCm/aotriton/releases/download/${VER}/aotriton-${VER}-${MANYLINUX}_${ARCH}-${ROCMBASE}.tar.bz2"`
	`10`	`+AOTRITON_URL="https://github.com/ROCm/aotriton/releases/download/${VER}/aotriton-${VER}-${MANYLINUX}_${ARCH}-${ROCMBASE}-shared.tar.bz2"`
`11`	`11`
`12`	`12`	`cd "${AOTRITON_INSTALL_PREFIX}"`
`13`	`13`	`# Must use -L to follow redirects`