Tritonは並列プログラミングのための言語とコンパイラです。カスタムDNN計算カーネルを効率的に記述し、最新のGPUハードウェア上で最大スループットで実行できるようにするためのPythonベースのプログラミング環境を提供するように設計されています。

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))}')

外：

テンソル([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351, 0.8149],
デバイス='cuda:0') テンソル([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351,
0.8149]、デバイス='cuda:0') トーチとトリトンの最大差は2.384185791015625e-07です

libdeviceライブラリパスをカスタマイズする

libdevice ライブラリへのパスは、asin カーネルにパスを渡すことによってカスタマイズできます。

 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))}')

外：

テンソル([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351, 0.8149],
デバイス='cuda:0') テンソル([0.4105, 0.5430, 0.0249, ..., 0.0424, 0.5351,
0.8149]、デバイス='cuda:0') トーチとトリトンの最大差は2.384185791015625e-07です

Jupyterノートブックをダウンロード: 07-extern-functions.ipynb

Pythonソースコードをダウンロード: 07-extern-functions.py

圧縮ファイルをダウンロード: 07-extern-functions.zip