[rocm6.1] Update Triton dependency logic (#47)

* Use hashed triton version

* skip triton dependency for PyTorch < 2.0

manywheel/build_rocm.sh

@@ -274,14 +274,18 @@ fi
 # Add triton install dependency
 # No triton dependency for now on 3.12 since we don't have binaries for it
 # and torch.compile doesn't work.
-if [[ $(uname) == "Linux" && "$DESIRED_PYTHON" != "3.12" ]]; 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)
+PYTORCH_VERSION=$(cat $PYTORCH_ROOT/version.txt | grep -oP "[0-9]+\.[0-9]+\.[0-9]+")
+# Assuming PYTORCH_VERSION=x.y.z, if x >= 2
+if [ ${PYTORCH_VERSION%%\.*} -ge 2 ]; then
+    if [[ $(uname) == "Linux" && "$DESIRED_PYTHON" != "3.12" ]]; 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}"
-    else
-        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}"
+        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
 fi