【Triton 教程】Libdevice (tl_extra.libdevice) 函式

超神经HyperAI發表於2024-11-28
dev函式

Triton 是一種用於並行程式設計的語言和編譯器。它旨在提供一個基於 Python 的程式設計環境，以高效編寫自定義 DNN 計算核心，並能夠在現代 GPU 硬體上以最大吞吐量執行。

更多 Triton 中文文件可訪問 →https://triton.hyper.ai/

Triton 可以呼叫外部庫中的自定義函式。在這個例子中，我們將使用 libdevice 庫在張量上應用 asin 函式。請參考以下連結獲取關於所有可用 libdevice 函式語義的詳細資訊：

  • CUDA：https://docs.nvidia.com/cuda/libdevice-users-guide/index.html
  • HIP：https://github.com/ROCm/llvm-project/tree/amd-staging/amd/device-libs/ocml/src

在 libdevice.py 中，我們試圖將相同計算但不同資料型別的函式聚合在一起。例如，__nv_asin 和 __nv_asinf 都計算輸入的反正弦的主值，但 __nv_asin 適用於 double 型別，而 __nv_asinf 適用於 float 型別。使用 Triton，您可以簡單地呼叫 tl.math.asin。根據輸入和輸出型別，Triton 會自動選擇正確的底層裝置函式來呼叫。

asin 核心

import torch


import triton
import triton.language as tl
from triton.language.extra import libdevice




@triton.jit
def asin_kernel(
    x_ptr,
    y_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)
    x = libdevice.asin(x)
    tl.store(y_ptr + offsets, x, mask=mask)

使用預設的 libdevice 庫路徑

可以使用 triton/language/math.py 中編碼的預設 libdevice 庫路徑。

torch.manual_seed(0)
size = 98432
x = torch.rand(size, device='cuda')
output_triton = torch.zeros(size, device='cuda')
output_torch = torch.asin(x)
assert x.is_cuda and output_triton.is_cuda
n_elements = output_torch.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
asin_kernel[grid](x, output_triton, n_elements, BLOCK_SIZE=1024)
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))}')

Out:

tensor([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351, 0.8149],
device='cuda:0') tensor([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351,
0.8149], device='cuda:0') The maximum difference between torch and triton is 2.384185791015625e-07

定製 libdevice 庫路徑

可以透過將 libdevice 庫的路徑傳遞給 asin 核心來定製 libdevice 庫的路徑。

output_triton = torch.empty_like(x)
asin_kernel[grid](x, output_triton, n_elements, BLOCK_SIZE=1024)
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))}')

Out:

tensor([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351, 0.8149],
device='cuda:0') tensor([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351,
0.8149], device='cuda:0') The maximum difference between torch and triton is 2.384185791015625e-07

Download Jupyter notebook: 07-extern-functions.ipynb

Download Python source code: 07-extern-functions.py

Download zipped: 07-extern-functions.zip

相關文章