Add therock

.ci/docker/build.sh

@@ -189,6 +189,20 @@ case "$tag" in
       INDUCTOR_BENCHMARKS=yes
     fi
     ;;
+  pytorch-linux-noble-rocm-nightly-py3)
+    ANACONDA_PYTHON_VERSION=3.12
+    GCC_VERSION=11
+    VISION=yes
+    ROCM_VERSION=nightly
+    NINJA_VERSION=1.9.0
+    TRITON=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    PYTORCH_ROCM_ARCH="gfx942"
+    USE_THEROCK_NIGHTLY=1
+    THEROCK_NIGHTLY_INDEX_URL=https://rocm.nightlies.amd.com/v2/gfx94X-dcgpu/
+    ;;
   pytorch-linux-jammy-xpu-n-1-py3)
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
@@ -385,6 +399,8 @@ docker build \
        --build-arg "SKIP_SCCACHE_INSTALL=${SKIP_SCCACHE_INSTALL:-}" \
        --build-arg "SKIP_LLVM_SRC_BUILD_INSTALL=${SKIP_LLVM_SRC_BUILD_INSTALL:-}" \
        --build-arg "INSTALL_MINGW=${INSTALL_MINGW:-}" \
+       --build-arg "USE_THEROCK_NIGHTLY=${USE_THEROCK_NIGHTLY:-}" \
+       --build-arg "THEROCK_NIGHTLY_INDEX_URL=${THEROCK_NIGHTLY_INDEX_URL:-}" \
        -f $(dirname ${DOCKERFILE})/Dockerfile \
        -t "$tmp_tag" \
        "$@" \

.ci/docker/common/install_miopen.sh

@@ -10,6 +10,12 @@ if [[ -z $ROCM_VERSION ]]; then
     exit 1;
 fi
 
+# Skip MIOpen build for TheRock nightly builds
+if [[ "${USE_THEROCK_NIGHTLY:-0}" == "1" ]]; then
+    echo "Skipping MIOpen custom build for TheRock nightly"
+    exit 0
+fi
+
 IS_UBUNTU=0
 ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
 case "$ID" in

.ci/docker/common/install_rocm.sh

@@ -26,6 +26,44 @@ install_ubuntu() {
     apt-get install -y libc++1
     apt-get install -y libc++abi1
 
+    # When USE_THEROCK_NIGHTLY=1, install ROCm from TheRock nightly wheels
+    if [[ "${USE_THEROCK_NIGHTLY:-0}" == "1" ]]; then
+      echo "install_rocm.sh: installing ROCm from TheRock nightly wheels"
+
+      # Clean any previous ROCm installation in the base CI image.
+      if [[ -d /opt/rocm ]]; then
+        echo "Removing existing /opt/rocm from base image"
+        rm -rf /opt/rocm
+      fi
+
+      : "${THEROCK_NIGHTLY_INDEX_URL:=https://rocm.nightlies.amd.com/v2/gfx94X-dcgpu/}"
+      python3 -m pip install \
+        --index-url "${THEROCK_NIGHTLY_INDEX_URL}" \
+        "rocm[libraries,devel]"
+
+      # Use the rocm-sdk CLI helper to populate environment defaults
+      ROCM_HOME="$(rocm-sdk path --root)"
+      ROCM_BIN="$(rocm-sdk path --bin)"
+      ROCM_CMAKE_PREFIX="$(rocm-sdk path --cmake)"
+
+      echo "ROCM_HOME=${ROCM_HOME}"
+      echo "ROCM_BIN=${ROCM_BIN}"
+      echo "ROCM_CMAKE_PREFIX=${ROCM_CMAKE_PREFIX}"
+
+      cat >/etc/profile.d/rocm-sdk.sh <<EOF
+export ROCM_HOME="${ROCM_HOME}"
+export PATH="${ROCM_BIN}:\$PATH"
+export CMAKE_PREFIX_PATH="${ROCM_CMAKE_PREFIX}:\${CMAKE_PREFIX_PATH:-}"
+EOF
+
+      export ROCM_HOME
+      export PATH="${ROCM_BIN}:${PATH}"
+      export CMAKE_PREFIX_PATH="${ROCM_CMAKE_PREFIX}:${CMAKE_PREFIX_PATH:-}"
+
+      echo "install_rocm.sh: TheRock nightly ROCm install complete"
+      exit 0
+    fi
+
     # Make sure rocm packages from repo.radeon.com have highest priority
     cat << EOF > /etc/apt/preferences.d/rocm-pin-600
 Package: *

.ci/docker/common/install_rocm_magma.sh

@@ -5,6 +5,13 @@ set -eou pipefail
 
 function do_install() {
     rocm_version=$1
+
+    # Skip magma install for TheRock nightly builds
+    if [[ "${USE_THEROCK_NIGHTLY:-0}" == "1" ]]; then
+        echo "Skipping MAGMA install for TheRock nightly build"
+        return 0
+    fi
+
     if [[ ${rocm_version} =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
         # chop off any patch version
         rocm_version="${rocm_version%.*}"

.ci/docker/ubuntu-rocm/Dockerfile

@@ -52,6 +52,10 @@ ENV INSTALLED_VISION ${VISION}
 
 # Install rocm
 ARG ROCM_VERSION
+ARG USE_THEROCK_NIGHTLY
+ARG THEROCK_NIGHTLY_INDEX_URL
+ENV USE_THEROCK_NIGHTLY=${USE_THEROCK_NIGHTLY}
+ENV THEROCK_NIGHTLY_INDEX_URL=${THEROCK_NIGHTLY_INDEX_URL}
 RUN mkdir ci_commit_pins
 COPY ./common/common_utils.sh common_utils.sh
 COPY ./ci_commit_pins/rocm-composable-kernel.txt ci_commit_pins/rocm-composable-kernel.txt
@@ -64,6 +68,17 @@ RUN bash ./install_rocm_magma.sh ${ROCM_VERSION}
 RUN rm install_rocm_magma.sh
 ADD ./common/install_miopen.sh install_miopen.sh
 RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
+
+RUN if [ "${USE_THEROCK_NIGHTLY}" = "1" ]; then \
+        ROCM_PATH_DETECTED="$(python3 -m rocm_sdk path --root)"; \
+    else \
+        ROCM_PATH_DETECTED="/opt/rocm"; \
+    fi; \
+    echo "ROCM_PATH=${ROCM_PATH_DETECTED}" >> /etc/environment && \
+    echo "PATH=${ROCM_PATH_DETECTED}/bin:${ROCM_PATH_DETECTED}/hcc/bin:${ROCM_PATH_DETECTED}/hip/bin:${ROCM_PATH_DETECTED}/opencl/bin:${ROCM_PATH_DETECTED}/llvm/bin:\${PATH}" >> /etc/environment && \
+    echo "MAGMA_HOME=${ROCM_PATH_DETECTED}/magma" >> /etc/environment
+
+# Set default ENV for traditional builds - will be overridden by /etc/environment for TheRock
 ENV ROCM_PATH /opt/rocm
 ENV PATH /opt/rocm/bin:$PATH
 ENV PATH /opt/rocm/hcc/bin:$PATH

.github/workflows/docker-builds.yml

@@ -61,6 +61,7 @@ jobs:
           pytorch-linux-jammy-py3.14-clang12,
           pytorch-linux-jammy-rocm-n-py3,
           pytorch-linux-noble-rocm-n-py3,
+          pytorch-linux-noble-rocm-nightly-py3,
           pytorch-linux-jammy-rocm-n-py3-benchmarks,
           pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang12,
           pytorch-linux-jammy-py3.10-gcc11,
@@ -90,6 +91,8 @@ jobs:
     # Docker uploads fail from LF runners, see https://github.com/pytorch/pytorch/pull/137358
     # runs-on: "${{ needs.get-label-type.outputs.label-type }}${{ matrix.runner }}"
     runs-on: "${{ matrix.runner }}"
+    outputs:
+      pytorch-linux-noble-rocm-n-py3: ${{ steps.build-docker-image.outputs.docker-image }}
     steps:
       - name: Clean workspace
         shell: bash

.github/workflows/therock-nightly.yml

@@ -0,0 +1,198 @@
+name: therock-nightly
+
+on:
+  # Run at midnight UTC every day
+  schedule:
+    - cron: "0 0 * * *"
+  # Also allow manual triggering
+  workflow_dispatch:
+
+env:
+  PYTORCH_NIGHTLY_REF: nightly
+  THEROCK_NIGHTLY_INDEX_URL: https://rocm.nightlies.amd.com/v2/gfx94X-dcgpu/
+  ROCM_VERSION: nightly
+  # Base CI image to modify (same as rocm-mi300.yml uses)
+  BASE_CI_IMAGE: pytorch/pytorch-linux-noble-rocm-n-py3:latest
+  # New image name after installing TheRock
+  THEROCK_IMAGE: pytorch-linux-noble-rocm-nightly-py3
+  PYTORCH_ROCM_ARCH: "gfx942"
+  PYTHON_VERSION: "3.12"
+
+jobs:
+  # Job 1: Build TheRock-based Docker image from Dockerfile
+  create-therock-image:
+    runs-on: ubuntu-latest
+    outputs:
+      image-tag: ${{ steps.build-image.outputs.image-tag }}
+
+    steps:
+      - name: Checkout pytorch
+        uses: actions/checkout@v4
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Show runner info
+        run: |
+          uname -a
+          docker version
+          df -h
+
+      - name: Build TheRock-based ROCm CI image
+        id: build-image
+        env:
+          THEROCK_NIGHTLY_INDEX_URL: ${{ env.THEROCK_NIGHTLY_INDEX_URL }}
+        run: |
+          set -euxo pipefail
+
+          # Tag with date for tracking
+          DATE_TAG=$(date +%Y%m%d)
+          IMAGE_TAG="${THEROCK_IMAGE}:nightly-${DATE_TAG}"
+
+          echo "Building TheRock nightly CI image: ${IMAGE_TAG}"
+
+          # Build from Dockerfile with TheRock nightly mode
+          docker build \
+            -f .ci/docker/ubuntu-rocm/Dockerfile \
+            --build-arg UBUNTU_VERSION=24.04 \
+            --build-arg ROCM_VERSION=nightly \
+            --build-arg ANACONDA_PYTHON_VERSION=3.12 \
+            --build-arg PYTORCH_ROCM_ARCH=gfx942 \
+            --build-arg USE_THEROCK_NIGHTLY=1 \
+            --build-arg THEROCK_NIGHTLY_INDEX_URL="${THEROCK_NIGHTLY_INDEX_URL}" \
+            --build-arg BUILD_ENVIRONMENT="therock-nightly-rocm" \
+            -t "${IMAGE_TAG}" \
+            .
+
+          # Also tag as latest
+          docker tag "${IMAGE_TAG}" "${THEROCK_IMAGE}:latest"
+
+          echo "image-tag=${IMAGE_TAG}" >> "$GITHUB_OUTPUT"
+
+          # Save docker image as artifact for next job
+          docker save "${IMAGE_TAG}" | gzip > therock-docker-image.tar.gz
+
+      - name: Upload Docker image as artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: therock-docker-image
+          path: therock-docker-image.tar.gz
+          retention-days: 1
+
+  # Job 2: Build PyTorch nightly using the TheRock image
+  build-pytorch-nightly:
+    needs: create-therock-image
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout pytorch nightly branch
+        uses: actions/checkout@v4
+        with:
+          repository: pytorch/pytorch
+          ref: ${{ env.PYTORCH_NIGHTLY_REF }}
+          submodules: recursive
+
+      - name: Download Docker image artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: therock-docker-image
+
+      - name: Load Docker image
+        run: |
+          echo "Loading TheRock Docker image..."
+          gunzip -c therock-docker-image.tar.gz | docker load
+          docker images
+
+      - name: Build PyTorch nightly using TheRock image
+        env:
+          IMAGE_TAG: ${{ needs.create-therock-image.outputs.image-tag }}
+          PYTORCH_ROCM_ARCH: ${{ env.PYTORCH_ROCM_ARCH }}
+          PYTHON_VERSION: ${{ env.PYTHON_VERSION }}
+        run: |
+          set -euxo pipefail
+
+          # Create output directory on host
+          mkdir -p ${{ github.workspace }}/artifacts
+
+          # Build PyTorch inside the container
+          docker run --rm \
+            -v ${{ github.workspace }}:/workspace \
+            -v ${{ github.workspace }}/artifacts:/artifacts \
+            -w /workspace \
+            -e PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" \
+            -e MAX_JOBS=8 \
+            "${IMAGE_TAG}" \
+            bash -c '
+              set -euxo pipefail
+
+              # Source the rocm-sdk environment
+              if [ -f /etc/profile.d/rocm-sdk.sh ]; then
+                source /etc/profile.d/rocm-sdk.sh
+                echo "Sourced /etc/profile.d/rocm-sdk.sh"
+              fi
+
+              # Verify ROCm is available
+              echo "=== ROCm Environment ==="
+              echo "ROCM_HOME=${ROCM_HOME}"
+              rocm-sdk path --root
+              rocm-sdk --version || true
+              which hipcc && hipcc --version || echo "hipcc not found"
+
+              # Activate conda environment for Python '"${PYTHON_VERSION}"'
+              source /opt/conda/etc/profile.d/conda.sh
+              conda activate py_'"${PYTHON_VERSION}"' || conda activate py_3.12
+
+              echo "=== Python Environment ==="
+              which python
+              python --version
+
+              # Install PyTorch dependencies
+              python -m pip install --upgrade pip
+              python -m pip install -r requirements.txt
+
+              # Initialize submodules if needed
+              git submodule sync
+              git submodule update --init --recursive
+
+              # Configure ROCm build flags
+              export USE_ROCM=1
+              export CMAKE_PREFIX_PATH="${ROCM_HOME}:${CMAKE_PREFIX_PATH:-}"
+              export PATH="${ROCM_HOME}/bin:${PATH}"
+
+              echo "=== Build Configuration ==="
+              echo "USE_ROCM=${USE_ROCM}"
+              echo "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}"
+              echo "MAX_JOBS=${MAX_JOBS}"
+
+              # Build PyTorch wheel
+              echo "=== Starting PyTorch Build ==="
+              python setup.py bdist_wheel
+
+              # Copy artifacts to host-mounted directory
+              cp dist/*.whl /artifacts/
+
+              echo "=== Build Complete ==="
+              ls -lh /artifacts/
+            '
+
+      - name: Upload PyTorch wheel artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: pytorch-therock-nightly-wheels-py${{ env.PYTHON_VERSION }}
+          path: artifacts/*.whl
+          retention-days: 7
+
+      - name: Show build summary
+        run: |
+          echo "## PyTorch TheRock Nightly Build Summary" >> $GITHUB_STEP_SUMMARY
+          echo "" >> $GITHUB_STEP_SUMMARY
+          echo "Built using TheRock nightly wheels from: ${{ env.THEROCK_NIGHTLY_INDEX_URL }}" >> $GITHUB_STEP_SUMMARY
+          echo "" >> $GITHUB_STEP_SUMMARY
+          echo "### Build Configuration" >> $GITHUB_STEP_SUMMARY
+          echo "- Base Image: ${{ env.BASE_CI_IMAGE }}" >> $GITHUB_STEP_SUMMARY
+          echo "- TheRock Image: ${{ env.THEROCK_IMAGE }}" >> $GITHUB_STEP_SUMMARY
+          echo "- Python Version: ${{ env.PYTHON_VERSION }}" >> $GITHUB_STEP_SUMMARY
+          echo "- PyTorch ROCM Arch: ${{ env.PYTORCH_ROCM_ARCH }}" >> $GITHUB_STEP_SUMMARY
+          echo "" >> $GITHUB_STEP_SUMMARY
+          echo "### Artifacts" >> $GITHUB_STEP_SUMMARY
+          ls -lh artifacts/*.whl >> $GITHUB_STEP_SUMMARY || echo "No wheels found"