xref: /aosp_15_r20/external/pytorch/scripts/install_triton_wheel.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2# Updates Triton to the pinned version for this copy of PyTorch
3BRANCH=$(git rev-parse --abbrev-ref HEAD)
4DOWNLOAD_PYTORCH_ORG="https://download.pytorch.org/whl"
5
6if [[ -z "${USE_XPU}" ]]; then
7    # Default install from PyTorch source
8
9    TRITON_VERSION="pytorch-triton==$(cat .ci/docker/triton_version.txt)"
10    if [[ "$BRANCH" =~ .*release.* ]]; then
11        pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/test/ $TRITON_VERSION
12    else
13        pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+$(head -c 10 .ci/docker/ci_commit_pins/triton.txt)
14    fi
15else
16    # The Triton xpu logic is as follows:
17    # 1. By default, install pre-built whls.
18    # 2. [Not exposed to user] If the user set `TRITON_XPU_BUILD_FROM_SOURCE=1` flag,
19    #    it will install Triton from the source.
20
21    TRITON_VERSION="pytorch-triton-xpu==$(cat .ci/docker/triton_version.txt)"
22    TRITON_XPU_COMMIT_ID="$(head -c 10 .ci/docker/ci_commit_pins/triton-xpu.txt)"
23    if [[ -z "${TRITON_XPU_BUILD_FROM_SOURCE}" ]]; then
24        pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ ${TRITON_VERSION}+${TRITON_XPU_COMMIT_ID}
25    else
26        TRITON_XPU_REPO="https://github.com/intel/intel-xpu-backend-for-triton"
27
28        # force-reinstall to ensure the pinned version is installed
29        pip install --force-reinstall "git+${TRITON_XPU_REPO}@${TRITON_XPU_COMMIT_ID}#subdirectory=python"
30    fi
31fi
32