8000 Add AOTriton install step for ROCm manylinux images (#1862) (#1867) · pytorch/builder@01a9531 · GitHub
[go: up one dir, main page]
Skip to content

Commit 01a9531

Browse files
atalmanjithunnair-amd
atalman
and
jithunnair-amd
authored
Add AOTriton install step for ROCm manylinux images (#1862) (#1867)
* Add AOTriton install step for ROCm * No common_utils.sh needed Co-authored-by: Jithun Nair <37884920+jithunnair-amd@users.noreply.github.com>
1 parent eb11301 commit 01a9531

File tree

4 files changed

+38
-0
lines changed

4 files changed

+38
-0
lines changed

‎common/aotriton_version.txt

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
0.6b
2+
manylinux_2_17
3+
rocm6
4+
04b5df8c8123f90cba3ede7e971e6fbc6040d506
5+
3db6ecbc915893ff967abd6e1b43bd5f54949868873be60dc802086c3863e648

‎common/install_aotriton.sh

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#!/bin/bash
2+
3+
set -ex
4+
5+
TARBALL='aotriton.tar.bz2'
6+
# This read command alwasy returns with exit code 1
7+
read -d "\n" VER MANYLINUX ROCMBASE PINNED_COMMIT SHA256 < aotriton_version.txt || true
8+
ARCH=$(uname -m)
9+
AOTRITON_INSTALL_PREFIX="$1"
10+
AOTRITON_URL="https://github.com/ROCm/aotriton/releases/download/${VER}/aotriton-${VER}-${MANYLINUX}_${ARCH}-${ROCMBASE}.tar.bz2"
11+
12+
cd "${AOTRITON_INSTALL_PREFIX}"
13+
# Must use -L to follow redirects
14+
curl -L --retry 3 -o "${TARBALL}" "${AOTRITON_URL}"
15+
ACTUAL_SHA256=$(sha256sum "${TARBALL}" | cut -d " " -f 1)
16+
if [ "${SHA256}" != "${ACTUAL_SHA256}" ]; then
17+
echo -n "Error: The SHA256 of downloaded tarball is ${ACTUAL_SHA256},"
18+
echo " which does not match the expected value ${SHA256}."
19+
exit
20+
fi
21+
tar xf "${TARBALL}" && rm -rf "${TARBALL}"

‎libtorch/Dockerfile

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,12 @@ RUN apt-get update -y && \
8181
RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
8282
RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
8383

84+
# Install AOTriton
85+
COPY ./common/aotriton_version.txt aotriton_version.txt
86+
COPY ./common/install_aotriton.sh install_aotriton.sh
87+
RUN bash ./install_aotriton.sh /opt/rocm && rm install_aotriton.sh aotriton_version.txt
88+
ENV AOTRITON_INSTALLED_PREFIX /opt/rocm/aotriton
89+
8490
FROM ${BASE_TARGET} as final
8591
COPY --from=openssl /opt/openssl /opt/openssl
8692
# Install patchelf

‎manywheel/Dockerfile

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,3 +173,9 @@ ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
173173
RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
174174
ADD ./common/install_miopen.sh install_miopen.sh
175175
RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
176+
177+
# Install AOTriton
178+
COPY ./common/aotriton_version.txt aotriton_version.txt
179+
COPY ./common/install_aotriton.sh install_aotriton.sh
180+
RUN bash ./install_aotriton.sh /opt/rocm && rm install_aotriton.sh aotriton_version.txt
181+
ENV AOTRITON_INSTALLED_PREFIX /opt/rocm/aotriton

0 commit comments

Comments
 (0)
0