Triton \u662f\u4e00\u79cd\u7528\u4e8e\u5e76\u884c\u7f16\u7a0b\u7684\u8bed\u8a00\u548c\u7f16\u8bd1\u5668\u3002\u5b83\u65e8\u5728\u63d0\u4f9b\u4e00\u4e2a\u57fa\u4e8e Python \u7684\u7f16\u7a0b\u73af\u5883\uff0c\u4ee5\u9ad8\u6548\u7f16\u5199\u81ea\u5b9a\u4e49 DNN \u8ba1\u7b97\u5185\u6838\uff0c\u5e76\u80fd\u591f\u5728\u73b0\u4ee3 GPU \u786c\u4ef6\u4e0a\u4ee5\u6700\u5927\u541e\u5410\u91cf\u8fd0\u884c\u3002<\/p>\n
\u66f4\u591a Triton \u4e2d\u6587\u6587\u6863\u53ef\u8bbf\u95ee \u2192https:\/\/triton.hyper.ai\/<\/a><\/p>\n Triton \u53ef\u4ee5\u8c03\u7528\u5916\u90e8\u5e93\u4e2d\u7684\u81ea\u5b9a\u4e49\u51fd\u6570\u3002\u5728\u8fd9\u4e2a\u4f8b\u5b50\u4e2d\uff0c\u6211\u4eec\u5c06\u4f7f\u7528 libdevice \u5e93\u5728\u5f20\u91cf\u4e0a\u5e94\u7528 asin \u51fd\u6570\u3002\u8bf7\u53c2\u8003\u4ee5\u4e0b\u94fe\u63a5\u83b7\u53d6\u5173\u4e8e\u6240\u6709\u53ef\u7528 libdevice \u51fd\u6570\u8bed\u4e49\u7684\u8be6\u7ec6\u4fe1\u606f\uff1a<\/p>\n \u5728 libdevice.py \u4e2d\uff0c\u6211\u4eec\u8bd5\u56fe\u5c06\u76f8\u540c\u8ba1\u7b97\u4f46\u4e0d\u540c\u6570\u636e\u7c7b\u578b\u7684\u51fd\u6570\u805a\u5408\u5728\u4e00\u8d77\u3002\u4f8b\u5982\uff0c__nv_asin \u548c __nv_asinf \u90fd\u8ba1\u7b97\u8f93\u5165\u7684\u53cd\u6b63\u5f26\u7684\u4e3b\u503c\uff0c\u4f46 __nv_asin \u9002\u7528\u4e8e double \u7c7b\u578b\uff0c\u800c __nv_asinf \u9002\u7528\u4e8e float \u7c7b\u578b\u3002\u4f7f\u7528 Triton\uff0c\u60a8\u53ef\u4ee5\u7b80\u5355\u5730\u8c03\u7528 tl.math.asin\u3002\u6839\u636e\u8f93\u5165\u548c\u8f93\u51fa\u7c7b\u578b\uff0cTriton \u4f1a\u81ea\u52a8\u9009\u62e9\u6b63\u786e\u7684\u5e95\u5c42\u8bbe\u5907\u51fd\u6570\u6765\u8c03\u7528\u3002<\/p>\n \u53ef\u4ee5\u4f7f\u7528 triton\/language\/math.py \u4e2d\u7f16\u7801\u7684\u9ed8\u8ba4 libdevice \u5e93\u8def\u5f84\u3002<\/p>\n Out:<\/p>\n tensor([0.4105, 0.5430, 0.0249, …, 0.0424, 0.5351, 0.8149], \u53ef\u4ee5\u901a\u8fc7\u5c06 libdevice \u5e93\u7684\u8def\u5f84\u4f20\u9012\u7ed9 asin \u5185\u6838\u6765\u5b9a\u5236 libdevice \u5e93\u7684\u8def\u5f84\u3002<\/p>\n Out:<\/p>\n tensor([0.4105, 0.5430, 0.0249, …, 0.0424, 0.5351, 0.8149], Download Jupyter notebook: 07-extern-functions.ipynb<\/a><\/p>\n Download Python source code: 07-extern-functions.py<\/a><\/p>\n\n
asin \u5185\u6838<\/h2>\n
import torch\n\n\nimport triton\nimport triton.language as tl\nfrom triton.language.extra import libdevice\n\n\n\n\n@triton.jit\ndef asin_kernel(\n x_ptr,\n y_ptr,\n n_elements,\n BLOCK_SIZE: tl.constexpr,\n):\n pid = tl.program_id(axis=0)\n block_start = pid * BLOCK_SIZE\n offsets = block_start + tl.arange(0, BLOCK_SIZE)\n mask = offsets < n_elements\n x = tl.load(x_ptr + offsets, mask=mask)\n x = libdevice.asin(x)\n tl.store(y_ptr + offsets, x, mask=mask)<\/code><\/pre>\n\u4f7f\u7528\u9ed8\u8ba4\u7684 libdevice \u5e93\u8def\u5f84<\/h2>\n
torch.manual_seed(0)\nsize = 98432\nx = torch.rand(size, device='cuda')\noutput_triton = torch.zeros(size, device='cuda')\noutput_torch = torch.asin(x)\nassert x.is_cuda and output_triton.is_cuda\nn_elements = output_torch.numel()\ngrid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\nasin_kernel[grid](x, output_triton, n_elements, BLOCK_SIZE=1024)\nprint(output_torch)\nprint(output_triton)\nprint(f'The maximum difference between torch and triton is '\n f'{torch.max(torch.abs(output_torch - output_triton))}')<\/code><\/pre>\n
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<\/p><\/blockquote>\n\u5b9a\u5236 libdevice \u5e93\u8def\u5f84<\/h2>\n
output_triton = torch.empty_like(x)\nasin_kernel[grid](x, output_triton, n_elements, BLOCK_SIZE=1024)\nprint(output_torch)\nprint(output_triton)\nprint(f'The maximum difference between torch and triton is '\n f'{torch.max(torch.abs(output_torch - output_triton))}')<\/code><\/pre>\n
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<\/p><\/blockquote>\n