{"id":278,"date":"2021-06-20T17:43:38","date_gmt":"2021-06-20T09:43:38","guid":{"rendered":"http:\/\/8.154.33.202\/?p=278"},"modified":"2021-06-20T17:43:38","modified_gmt":"2021-06-20T09:43:38","slug":"numba-%e6%93%8d%e4%bd%9cpytorch-gpu-tensor","status":"publish","type":"post","link":"http:\/\/8.154.33.202\/?p=278","title":{"rendered":"numba \u64cd\u4f5cpytorch GPU tensor"},"content":{"rendered":"\n

numba\u53ef\u4ee5\u5bf9numpy\u7684array\u4f7f\u7528gpu\u8fdb\u884c\u5e76\u884c\u8ba1\u7b97\uff0c\u4ece\u800c\u8fdb\u884c\u52a0\u901f\u3002<\/p>\n\n\n\n

import numpy as np\nfrom numba import cuda\n\n\u81ea\u5b9a\u4e49numba cuda\u51fd\u6570\n@cuda.jit\ndef mat_add(A, B, C):\n    j, i = cuda.grid(2)\n    C[i, j] = A[i, j] + B[i, j]\n\na = np.arange(12).reshape(3, 4)\nb = np.arange(12, 24).reshape(3, 4)\nc = cuda.device_array_like(a_in)\n\na_dev = cuda.to_device(a)\nb_dev = cuda.to_device(b)\n\nblocks = (1, 1)\nthreads_per_block = (4, 3)\n\nmat_add[blocks, threads_per_block](a_dev, b_dev, c)\nprint(c.copy_to_host())<\/pre>\n\n\n\n
\u5176\u5b83<\/p>\n\n\n\n
stride_j, stride_i = cuda.gridsize(2)\n\n\u4f7f\u7528\u5171\u4eab\u5185\u5b58\ntmp = cuda.shared.array((32, 32), numba.types.int32)\n\u5bf9\u5171\u4eab\u5185\u5b58\u7684\u8bfb\u5199\u8981\u6ce8\u610f\u8fdb\u884c\u540c\u6b65\ncuda.syncthreads()<\/pre>\n\n\n\n
pytorch\u73b0\u5728\u63d0\u4f9b\u7ed9numba\u4e00\u4e2a\u63a5\u53e3\uff0c\u53ef\u4ee5\u4f7f\u7528numba\u76f4\u63a5\u5bf9pytorch\u7684gpu tensor\u8fdb\u884c\u64cd\u4f5c\uff0c\u4ece\u800c\u53ef\u4ee5\u6bd4\u8f83\u65b9\u4fbf\u5730\u5b9e\u73b0\u81ea\u5df1\u60f3\u8981\u7684\u4e00\u4e9b\u8fd0\u7b97\u3002<\/p>\n\n\n\n
import torch\nimport numba.cuda\n\na = torch.arange(12).reshape(3,4).cuda()\na_nb = numba.cuda.as_cuda_array(a)\n# \u8fd9\u91cca_nb\u548ca\u7684\u6570\u636e\u662f\u540c\u4e00\u7247\u5185\u5b58\u7a7a\u95f4\u7684\uff0c\u56e0\u6b64a_nb\u5bf9\u6570\u636e\u4fee\u6539\uff0ca\u7684\u6570\u636e\u4e5f\u4f1a\u540c\u6837\u4fee\u6539\uff0c\u53cd\u4e4b\u4ea6\u7136\u3002\n# \u5f97\u5230a_nb\u540e\uff0c\u5c31\u53ef\u4ee5\u7528\u81ea\u5b9a\u4e49\u7684\u5e76\u884c\u8fd0\u7b97\u51fd\u6570\u8fdb\u884c\u64cd\u4f5c\u4e86<\/pre>\n","protected":false},"excerpt":{"rendered":"
numba\u53ef\u4ee5\u5bf9numpy\u7684array\u4f7f\u7528gpu\u8fdb\u884c\u5e76\u884c\u8ba1\u7b97\uff0c\u4ece\u800c\u8fdb\u884c\u52a0\u901f\u3002 import numpy as  … <\/p>\n
\u7ee7\u7eed\u9605\u8bfb\u201cnumba \u64cd\u4f5cpytorch GPU tensor\u201d<\/span><\/a><\/p>\n","protected":false},"author":1,"featured_media":0,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[2],"tags":[],"class_list":["post-278","post","type-post","status-publish","format-standard","hentry","category-python"],"_links":{"self":[{"href":"http:\/\/8.154.33.202\/index.php?rest_route=\/wp\/v2\/posts\/278"}],"collection":[{"href":"http:\/\/8.154.33.202\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"http:\/\/8.154.33.202\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"http:\/\/8.154.33.202\/index.php?rest_route=\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"http:\/\/8.154.33.202\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=278"}],"version-history":[{"count":1,"href":"http:\/\/8.154.33.202\/index.php?rest_route=\/wp\/v2\/posts\/278\/revisions"}],"predecessor-version":[{"id":279,"href":"http:\/\/8.154.33.202\/index.php?rest_route=\/wp\/v2\/posts\/278\/revisions\/279"}],"wp:attachment":[{"href":"http:\/\/8.154.33.202\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=278"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"http:\/\/8.154.33.202\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=278"},{"taxonomy":"post_tag","embeddable":true,"href":"http:\/\/8.154.33.202\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=278"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}