[XPU] Add xpu support of make triton (#126513)

This PR is to add XPU support for `make triton`.

If a user wishes to use Triton with XPU support, the user needs to install the  [intel-xpu-backend-for-triton](https://github.com/intel/intel-xpu-backend-for-triton).

This PR allows the user to easily install Triton for xpu backend support:

```
# clone the pytorch repo
export USE_XPU=1
make triton
```
The XPU version of triton will always be built from the source. It will cat the commit id from `.ci/docker/ci_commit_pins/triton-xpu.txt`, for example, `b8c64f64c18d8cac598b3adb355c21e7439c21de`.

So the final call would be like:

```
pip install --force-reinstall "git+https://github.com/intel/intel-xpu-backend-for-triton@b8c64f64c18d8cac598b3adb355c21e7439c21de#subdirectory=python"
```
Pull Request resolved: pytorch/pytorch#126513
Approved by: https://github.com/EikanWang, https://github.com/atalman

Author: Stonepia

README.md

@@ -213,6 +213,7 @@ conda install -c pytorch magma-cuda121  # or the magma-cuda* that matches your C
 
 # (optional) If using torch.compile with inductor/triton, install the matching version of triton
 # Run from the pytorch directory after cloning
+# For Intel GPU support, please explicitly `export USE_XPU=1` before running command.
 make triton
 ```
 

scripts/install_triton_wheel.sh

@@ -1,11 +1,23 @@
 #!/bin/bash
 # Updates Triton to the pinned version for this copy of PyTorch
 BRANCH=$(git rev-parse --abbrev-ref HEAD)
-TRITON_VERSION="pytorch-triton==$(cat .ci/docker/triton_version.txt)"
-DOWNLOAD_PYTORCH_ORG="https://download.pytorch.org/whl"
 
-if [[ "$BRANCH" =~ .*release.* ]]; then
-    pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/test/ $TRITON_VERSION
+if [[ -z "${USE_XPU}" ]]; then
+    # Default install from PyTorch source
+
+    TRITON_VERSION="pytorch-triton==$(cat .ci/docker/triton_version.txt)"
+    DOWNLOAD_PYTORCH_ORG="https://download.pytorch.org/whl"
+    if [[ "$BRANCH" =~ .*release.* ]]; then
+        pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/test/ $TRITON_VERSION
+    else
+        pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+$(head -c 10 .ci/docker/ci_commit_pins/triton.txt)
+    fi
 else
-    pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+$(head -c 10 .ci/docker/ci_commit_pins/triton.txt)
+    # Always install Triton for XPU from source
+
+    TRITON_XPU_REPO="https://github.com/intel/intel-xpu-backend-for-triton"
+    TRITON_XPU_COMMIT_ID="$(cat .ci/docker/ci_commit_pins/triton-xpu.txt)"
+
+    # force-reinstall to ensure the pinned version is installed
+    pip install --force-reinstall "git+${TRITON_XPU_REPO}@${TRITON_XPU_COMMIT_ID}#subdirectory=python"
 fi