Add pytorch-triton-rocm as an install dependency for ROCm (pytorch#1463)

* Add pytorch-triton-rocm as an install dependency for ROCm * Update build_rocm.sh
atalman · Aug 22, 2023 · eba456f · eba456f
1 parent a5aa27f
commit eba456f
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
@@ -214,6 +214,18 @@ elif [[ $ROCM_INT -ge 50600 ]]; then
     DEPS_AUX_DSTLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_DST/})
 fi
 
+# Add triton install dependency
+if [[ $(uname) == "Linux" ]]; then
+    TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton-rocm.txt)
+    TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
+
+    if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
+        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton-rocm==${TRITON_VERSION}+${TRITON_SHORTHASH}"
+    else
+        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}+${TRITON_SHORTHASH}"
+    fi
+fi
+
 
 echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}"