search

intel / intel-xpu-backend-for-triton

OpenAI Triton backend for Intel® GPUs
MIT License
144 stars 44 forks source link

readme

Build and test Triton wheels

Intel® XPU Backend for Triton*

This is the development repository of Intel® XPU Backend for Triton*, a new Triton backend for Intel GPUs. Intel® XPU Backend for Triton* is a out of tree backend module for Triton used to provide best-in-class performance and productivity on any Intel GPUs for PyTorch and standalone usage.

Compatibility

Note that Intel® XPU Backend for Triton* is not compatible with Intel® Extension for PyTorch* and Intel® oneAPI Base Toolkit*.

Quick Installation

Prerequisites

  1. Latest Rolling Release or Long Term Support Release of GPU driver
  2. Latest release of Intel® Deep Learning Essentials

Install PyTorch and Triton from nightly wheels

Currently, Intel® XPU Backend for Triton* requires a special version of PyTorch and both can be installed from nightly wheels. Navigate to the nightly wheels workflow, select the most recent successful run on the top of the page and download an artifact for the corresponding Python version. Extract the archive and in the extracted directory execute:

pip install torch-*.whl triton-*.whl

Before using Intel® XPU Backend for Triton* you need to initialize the toolchain. The default location is /opt/intel/oneapi (if installed as a root user) or ~/intel/oneapi (if installed as a regular user).

# replace /opt/intel/oneapi with the actual location of Intel® Deep Learning Essentials
source /opt/intel/oneapi/setvars.sh

Install from source

Prerequisites

  1. Latest Rolling Release or Long Term Support Release of GPU driver
  2. Latest release of Intel® Deep Learning Essentials

Compile PyTorch and Triton from source

Currently, Intel® XPU Backend for Triton* requires a special version of PyTorch and both need to be compiled at the same time.

Before compiling PyTorch and Intel® XPU Backend for Triton* you need to initialize the toolchain. The default location is /opt/intel/oneapi (if installed as a root user) or ~/intel/oneapi (if installed as a regular user).

# replace /opt/intel/oneapi with the actual location of Intel® Deep Learning Essentials
source /opt/intel/oneapi/setvars.sh

Clone this repository:

git clone https://github.com/intel/intel-xpu-backend-for-triton.git
cd intel-xpu-backend-for-triton

To avoid potential conflicts with installed packages it is recommended to create and activate a new Python virtual environment:

python -m venv .venv --prompt triton
source .venv/bin/activate

Compile and install PyTorch:

scripts/install-pytorch.sh --source

Compile and install Intel® XPU Backend for Triton*:

scripts/compile-triton.sh

Building with a custom LLVM

Triton uses LLVM to generate code for GPUs and CPUs. Normally, the Triton build downloads a prebuilt LLVM, but you can also build LLVM from source and use that.

LLVM does not have a stable API, so the Triton build will not work at an arbitrary LLVM version.

  1. Find the version of LLVM that Triton builds against. Check cmake/llvm-hash.txt to see the current version.

  2. Checkout LLVM at this revision to the directory llvm, which must be in the same directory as intel-xpu-backend-for-triton:

  3. In the directory intel-xpu-backend-for-triton, build Triton with custom LLVM:

    ./scripts/compile-triton.sh --llvm --triton

Tips for building

Running tests

There currently isn't a turnkey way to run all the Triton tests, but you can follow the following recipe.

scripts/test-triton.sh

Tips for hacking

For detailed instructions on how to debug Triton's frontend, please refer to this tutorial. The following includes additional tips for hacking on Triton's backend.

Helpful environment variables

Usage Guide

Code Modifications

Intel® XPU Backend for Triton* requires a special version of PyTorch that can be built from sources or installed from nightly wheels.

  1. Add import torch for xpu support.
  2. Put the tensor and models to XPU by calling to('xpu').

This repository contains modified tutorials that must be used with Intel® XPU Backend for Triton*.

The following examples show modifications for the user code.

Example 1 : Triton Kernel

This example is a modified version of Vector Add triton kernel. Please refer to Vector Add for detailed comments and illustration about the code semantics.

Comparing to the original code, the following code modifies:

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(
    x_ptr,
    y_ptr,
    output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

def add(x: torch.Tensor, y: torch.Tensor):
    # Put the tensor to xpu
    output = torch.empty_like(x).xpu()
    assert x.is_xpu and y.is_xpu and output.is_xpu
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)

    return output

# For manual_seed, needs to use API for XPU
torch.xpu.manual_seed(0)
size = 512
# For tensors, needs to be put on XPU
x = torch.rand(size, device='xpu')
y = torch.rand(size, device='xpu')
output_torch = x + y
output_triton = add(x, y)
print(output_torch)
print(output_triton)
print(
    f'The maximum difference between torch and triton is '
    f'{torch.max(torch.abs(output_torch - output_triton))}'
)

Example 2 : End-to-End Model

Triton is transparent for end-to-end models. One could easily use torch.compile with inductor as backend by default. It will automatically generates triton kernel and gets benefit from it.

import torch
from torch._dynamo.testing import rand_strided

from torch.nn import *
class simpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # tensors inside model should be on xpu
        self.y = rand_strided((32, 8), (8, 1), device='xpu:0', dtype=torch.float32)

    def forward(self, x):
        z = x + self.y
        return z

# tensors passed to the model should be on xpu
x = rand_strided((32, 8), (8, 1), device='xpu:0', dtype=torch.float32)
xpu_model = simpleModel()
# Call torch.compile for optimization
optimized_mod = torch.compile(xpu_model)

graph_result = optimized_mod(x)

Performance Analysis Guide

There are several ways of doing performance analysis. We recommend using torch.profiler for end-to-end performance analysis and using Intel® VTune™ Profiler for more detailed kernel analysis. Note that the user needs to explicitly set TRITON_XPU_PROFILE=1 when the user needs to enable kernel profiling.

export TRITON_XPU_PROFILE=1

Contributing

Community contributions are more than welcome, whether it be to fix bugs or to add new features at github. For more detailed instructions, please visit our contributor's guide.

License

MIT License. As found in LICENSE file.

Security

See Intel's Security Center for information on how to report a potential security issue or vulnerability.

See also: Security Policy.