[ROCm] remove triton-rocm commit pin and merge pins with triton.txt (#133438)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/133438
Approved by: https://github.com/jithunnair-amd, https://github.com/malfet

Co-authored-by: Jithun Nair <37884920+jithunnair-amd@users.noreply.github.com>

.ci/docker/centos-rocm/Dockerfile

@@ -108,10 +108,10 @@ ENV CMAKE_C_COMPILER cc
 ENV CMAKE_CXX_COMPILER c++
 COPY ./common/install_triton.sh install_triton.sh
 COPY ./common/common_utils.sh common_utils.sh
-COPY ci_commit_pins/triton-rocm.txt triton-rocm.txt
+COPY ci_commit_pins/triton.txt triton.txt
 COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
-RUN rm install_triton.sh common_utils.sh triton-rocm.txt triton_version.txt
+RUN rm install_triton.sh common_utils.sh triton.txt triton_version.txt
 
 # Install AOTriton (Early fail)
 COPY ./aotriton_version.txt aotriton_version.txt

.ci/docker/ci_commit_pins/triton-rocm.txt

@@ -1 +0,0 @@
-21eae954efa5bf584da70324b640288c3ee7aede

.ci/docker/ci_commit_pins/triton.txt

@@ -1 +1 @@
-dedb7bdf339a3546896d4820366ca562c586bfa0
+757b6a61e7df814ba806f498f8bb3160f84b120c

.ci/docker/common/install_triton.sh

@@ -12,10 +12,7 @@ conda_reinstall() {
   as_jenkins conda install -q -n py_$ANACONDA_PYTHON_VERSION -y --force-reinstall $*
 }
 
-if [ -n "${ROCM_VERSION}" ]; then
+if [ -n "${XPU_VERSION}" ]; then
-  TRITON_REPO="https://github.com/openai/triton"
-  TRITON_TEXT_FILE="triton-rocm"
-elif [ -n "${XPU_VERSION}" ]; then
   TRITON_REPO="https://github.com/intel/intel-xpu-backend-for-triton"
   TRITON_TEXT_FILE="triton-xpu"
 else

.ci/docker/ubuntu-rocm/Dockerfile

@@ -100,10 +100,10 @@ ARG TRITON
 # try to reach out to S3, which docker build runners don't have access
 COPY ./common/install_triton.sh install_triton.sh
 COPY ./common/common_utils.sh common_utils.sh
-COPY ci_commit_pins/triton-rocm.txt triton-rocm.txt
+COPY ci_commit_pins/triton.txt triton.txt
 COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
-RUN rm install_triton.sh common_utils.sh triton-rocm.txt triton_version.txt
+RUN rm install_triton.sh common_utils.sh triton.txt triton_version.txt
 
 # Install AOTriton
 COPY ./aotriton_version.txt aotriton_version.txt

.circleci/scripts/binary_populate_env.sh

@@ -90,7 +90,7 @@ fi
 if [[ "$PACKAGE_TYPE" =~ .*wheel.* && -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*rocm.* && $(uname) == "Linux" ]]; then
     TRITON_REQUIREMENT="pytorch-triton-rocm==${TRITON_VERSION}; ${TRITON_CONSTRAINT}"
     if [[ -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*dev.* ]]; then
-        TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton-rocm.txt)
+        TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton.txt)
         TRITON_REQUIREMENT="pytorch-triton-rocm==${TRITON_VERSION}+${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
     fi
     if [[ -z "${PYTORCH_EXTRA_INSTALL_REQUIREMENTS:-}" ]]; then

.github/scripts/build_triton_wheel.py

@@ -15,9 +15,7 @@ REPO_DIR = SCRIPT_DIR.parent.parent
 
 def read_triton_pin(device: str = "cuda") -> str:
     triton_file = "triton.txt"
-    if device == "rocm":
+    if device == "xpu":
-        triton_file = "triton-rocm.txt"
-    elif device == "xpu":
         triton_file = "triton-xpu.txt"
     with open(REPO_DIR / ".ci" / "docker" / "ci_commit_pins" / triton_file) as f:
         return f.read().strip()
@@ -50,25 +48,6 @@ def patch_init_py(
         f.write(orig)
 
 
-# TODO: remove patch_setup_py() once we have a proper fix for https://github.com/triton-lang/triton/issues/4527
-def patch_setup_py(path: Path) -> None:
-    with open(path) as f:
-        orig = f.read()
-    try:
-        orig = check_and_replace(
-            orig,
-            "https://tritonlang.blob.core.windows.net/llvm-builds/",
-            "https://oaitriton.blob.core.windows.net/public/llvm-builds/",
-        )
-        with open(path, "w") as f:
-            f.write(orig)
-    except RuntimeError as e:
-        print(
-            f"Applying patch_setup_py() for llvm-build package failed: {e}.",
-            "If you are trying to build a newer version of Triton, you can ignore this.",
-        )
-
-
 def build_triton(
     *,
     version: str,
@@ -110,9 +89,6 @@ def build_triton(
         else:
             check_call(["git", "checkout", commit_hash], cwd=triton_basedir)
 
-        # TODO: remove this and patch_setup_py() once we have a proper fix for https://github.com/triton-lang/triton/issues/4527
-        patch_setup_py(triton_pythondir / "setup.py")
-
         if build_conda:
             with open(triton_basedir / "meta.yaml", "w") as meta:
                 print(

.github/workflows/build-triton-wheel.yml

@@ -13,7 +13,6 @@ on:
       - .github/scripts/build_triton_wheel.py
       - .github/ci_commit_pins/triton.txt
       - .ci/docker/ci_commit_pins/triton.txt
-      - .ci/docker/ci_commit_pins/triton-rocm.txt
       - .ci/docker/ci_commit_pins/triton-xpu.txt
   pull_request:
     paths:
@@ -21,7 +20,6 @@ on:
       - .github/scripts/build_triton_wheel.py
       - .github/ci_commit_pins/triton.txt
       - .ci/docker/ci_commit_pins/triton.txt
-      - .ci/docker/ci_commit_pins/triton-rocm.txt
       - .ci/docker/ci_commit_pins/triton-xpu.txt
 
 concurrency:

CODEOWNERS

@@ -57,7 +57,6 @@ nn/qat/ @jerryzh168
 # Docker
 /.ci/docker/ @jeffdaily
 /.ci/docker/ci_commit_pins/triton.txt @desertfire @Chillee @eellison @shunting314 @bertmaher @jeffdaily @jataylo @jithunnair-amd @pruthvistony
-/.ci/docker/ci_commit_pins/triton-rocm.txt @jeffdaily @jataylo @jithunnair-amd @pruthvistony
 /.ci/docker/ci_commit_pins/triton-xpu.txt @EikanWang @gujinghui
 
 # Github Actions