{"id":10020405,"date":"2024-11-21T04:05:28","date_gmt":"2024-11-20T20:05:28","guid":{"rendered":"https:\/\/itmax.cn\/?p=10020405"},"modified":"2024-11-21T04:05:28","modified_gmt":"2024-11-20T20:05:28","slug":"%e3%80%90tvm-%e6%95%99%e7%a8%8b%e3%80%91%e4%bd%bf%e7%94%a8%e5%85%83%e7%bb%84%e8%be%93%e5%85%a5%ef%bc%88tuple-inputs%ef%bc%89%e8%bf%9b%e8%a1%8c%e8%ae%a1%e7%ae%97%e5%92%8c%e5%bd%92%e7%ba%a6","status":"publish","type":"post","link":"https:\/\/itmax.cn\/?p=10020405","title":{"rendered":"\u3010TVM \u6559\u7a0b\u3011\u4f7f\u7528\u5143\u7ec4\u8f93\u5165\uff08Tuple Inputs\uff09\u8fdb\u884c\u8ba1\u7b97\u548c\u5f52\u7ea6"},"content":{"rendered":"

Apache TVM \u662f\u4e00\u4e2a\u7aef\u5230\u7aef\u7684\u6df1\u5ea6\u5b66\u4e60\u7f16\u8bd1\u6846\u67b6\uff0c\u9002\u7528\u4e8e CPU\u3001GPU \u548c\u5404\u79cd\u673a\u5668\u5b66\u4e60\u52a0\u901f\u82af\u7247\u3002\u66f4\u591a TVM \u4e2d\u6587\u6587\u6863\u53ef\u8bbf\u95ee \u2192 https:\/\/tvm.hyper.ai\/<\/a><\/p>\n

\u4f5c\u8005\uff1aZiheng Jiang<\/a><\/p>\n

\u82e5\u8981\u5728\u5355\u4e2a\u5faa\u73af\u4e2d\u8ba1\u7b97\u5177\u6709\u76f8\u540c shape \u7684\u591a\u4e2a\u8f93\u51fa\uff0c\u6216\u6267\u884c\u591a\u4e2a\u503c\u7684\u5f52\u7ea6\uff0c\u4f8b\u5982 argmax\u3002\u8fd9\u4e9b\u95ee\u9898\u53ef\u4ee5\u901a\u8fc7\u5143\u7ec4\u8f93\u5165\u6765\u89e3\u51b3\u3002<\/p>\n

\u672c\u6559\u7a0b\u4ecb\u7ecd\u4e86 TVM \u4e2d\u5143\u7ec4\u8f93\u5165\u7684\u7528\u6cd5\u3002<\/p>\n

from __future__ import absolute_import, print_function\n\nimport tvm\nfrom tvm import te\nimport numpy as np<\/code><\/pre>\n
\u63cf\u8ff0\u6279\u91cf\u8ba1\u7b97<\/h2>\n
\u5bf9\u4e8e shape \u76f8\u540c\u7684\u7b97\u5b50\uff0c\u82e5\u8981\u5728\u4e0b\u4e00\u4e2a\u8c03\u5ea6\u8fc7\u7a0b\u4e2d\u4e00\u8d77\u8c03\u5ea6\uff0c\u53ef\u4ee5\u5c06\u5b83\u4eec\u653e\u5728\u4e00\u8d77\u4f5c\u4e3a te.compute \u7684\u8f93\u5165\u3002<\/p>\n
n = te.var("n")\nm = te.var("m")\nA0 = te.placeholder((m, n), name="A0")\nA1 = te.placeholder((m, n), name="A1")\nB0, B1 = te.compute((m, n), lambda i, j: (A0[i, j] + 2, A1[i, j] * 3), name="B")\n\n# \u751f\u6210\u7684 IR \u4ee3\u7801\uff1a\ns = te.create_schedule(B0.op)\nprint(tvm.lower(s, [A0, A1, B0, B1], simple_mode=True))<\/code><\/pre>\n
\u8f93\u51fa\u7ed3\u679c\uff1a<\/p>\n
@main = primfn(A0_1: handle, A1_1: handle, B_2: handle, B_3: handle) -> ()\n  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}\n  buffers = {A0: Buffer(A0_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"),\n             A1: Buffer(A1_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"),\n             B: Buffer(B_4: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto"),\n             B_1: Buffer(B_5: Pointer(float32), float32, [(stride_3: int32*m)], [], type="auto")}\n  buffer_map = {A0_1: A0, A1_1: A1, B_2: B, B_3: B_1}\n  preflattened_buffer_map = {A0_1: A0_3: Buffer(A0_2, float32, [m, n: int32], [stride, stride_4: int32], type="auto"), A1_1: A1_3: Buffer(A1_2, float32, [m, n], [stride_1, stride_5: int32], type="auto"), B_2: B_6: Buffer(B_4, float32, [m, n], [stride_2, stride_6: int32], type="auto"), B_3: B_7: Buffer(B_5, float32, [m, n], [stride_3, stride_7: int32], type="auto")} {\n  for (i: int32, 0, m) {\n    for (j: int32, 0, n) {\n      B[((i*stride_2) + (j*stride_6))] = (A0[((i*stride) + (j*stride_4))] + 2f32)\n      B_1[((i*stride_3) + (j*stride_7))] = (A1[((i*stride_1) + (j*stride_5))]*3f32)\n    }\n  }\n}<\/code><\/pre>\n
\u4f7f\u7528\u534f\u540c\u8f93\u5165\uff08Collaborative Inputs\uff09\u63cf\u8ff0\u5f52\u7ea6<\/h2>\n
\u6709\u65f6\u9700\u8981\u591a\u4e2a\u8f93\u5165\u6765\u8868\u8fbe\u5f52\u7ea6\u7b97\u5b50\uff0c\u5e76\u4e14\u8f93\u5165\u4f1a\u534f\u540c\u5de5\u4f5c\uff0c\u4f8b\u5982 argmax\u3002\u5728\u5f52\u7ea6\u8fc7\u7a0b\u4e2d\uff0cargmax \u8981\u6bd4\u8f83\u64cd\u4f5c\u6570\u7684\u503c\uff0c\u8fd8\u9700\u8981\u4fdd\u7559\u64cd\u4f5c\u6570\u7684\u7d22\u5f15\uff0c\u53ef\u7528 te.comm_reducer() \u8868\u793a\uff1a<\/p>\n
# x \u548c y \u662f\u5f52\u7ea6\u7684\u64cd\u4f5c\u6570\uff0c\u5b83\u4eec\u90fd\u662f\u5143\u7ec4\u7684\u7d22\u5f15\u548c\u503c\u3002\ndef fcombine(x, y):\n    lhs = tvm.tir.Select((x[1] >= y[1]), x[0], y[0])\n    rhs = tvm.tir.Select((x[1] >= y[1]), x[1], y[1])\n    return lhs, rhs\n\n# \u8eab\u4efd\u5143\u7d20\u4e5f\u8981\u662f\u4e00\u4e2a\u5143\u7ec4\uff0c\u6240\u4ee5 `fidentity` \u63a5\u6536\u4e24\u79cd\u7c7b\u578b\u4f5c\u4e3a\u8f93\u5165\u3002\ndef fidentity(t0, t1):\n    return tvm.tir.const(-1, t0), tvm.te.min_value(t1)\n\nargmax = te.comm_reducer(fcombine, fidentity, name="argmax")\n\n# \u63cf\u8ff0\u5f52\u7ea6\u8ba1\u7b97\nm = te.var("m")\nn = te.var("n")\nidx = te.placeholder((m, n), name="idx", dtype="int32")\nval = te.placeholder((m, n), name="val", dtype="int32")\nk = te.reduce_axis((0, n), "k")\nT0, T1 = te.compute((m,), lambda i: argmax((idx[i, k], val[i, k]), axis=k), name="T")\n\n# \u751f\u6210\u7684 IR \u4ee3\u7801\uff1a\ns = te.create_schedule(T0.op)\nprint(tvm.lower(s, [idx, val, T0, T1], simple_mode=True))<\/code><\/pre>\n
\u8f93\u51fa\u7ed3\u679c\uff1a<\/p>\n
@main = primfn(idx_1: handle, val_1: handle, T_2: handle, T_3: handle) -> ()\n  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}\n  buffers = {idx: Buffer(idx_2: Pointer(int32), int32, [(stride: int32*m: int32)], [], type="auto"),\n             val: Buffer(val_2: Pointer(int32), int32, [(stride_1: int32*m)], [], type="auto"),\n             T: Buffer(T_4: Pointer(int32), int32, [(stride_2: int32*m)], [], type="auto"),\n             T_1: Buffer(T_5: Pointer(int32), int32, [(stride_3: int32*m)], [], type="auto")}\n  buffer_map = {idx_1: idx, val_1: val, T_2: T, T_3: T_1}\n  preflattened_buffer_map = {idx_1: idx_3: Buffer(idx_2, int32, [m, n: int32], [stride, stride_4: int32], type="auto"), val_1: val_3: Buffer(val_2, int32, [m, n], [stride_1, stride_5: int32], type="auto"), T_2: T_6: Buffer(T_4, int32, [m], [stride_2], type="auto"), T_3: T_7: Buffer(T_5, int32, [m], [stride_3], type="auto")} {\n  for (i: int32, 0, m) {\n    T[(i*stride_2)] = -1\n    T_1[(i*stride_3)] = -2147483648\n    for (k: int32, 0, n) {\n      T[(i*stride_2)] = @tir.if_then_else((val[((i*stride_1) + (k*stride_5))] <= T_1[(i*stride_3)]), T[(i*stride_2)], idx[((i*stride) + (k*stride_4))], dtype=int32)\n      T_1[(i*stride_3)] = @tir.if_then_else((val[((i*stride_1) + (k*stride_5))] <= T_1[(i*stride_3)]), T_1[(i*stride_3)], val[((i*stride_1) + (k*stride_5))], dtype=int32)\n    }\n  }\n}<\/code><\/pre>\n
\u4f7f\u7528\u5143\u7ec4\u8f93\u5165\u8c03\u5ea6\u64cd\u4f5c<\/h2>\n
\u867d\u7136\u4e00\u6b21 batch \u64cd\u4f5c\u4f1a\u6709\u591a\u4e2a\u8f93\u51fa\uff0c\u4f46\u5b83\u4eec\u53ea\u80fd\u4e00\u8d77\u8c03\u5ea6\u3002<\/p>\n
n = te.var("n")\nm = te.var("m")\nA0 = te.placeholder((m, n), name="A0")\nB0, B1 = te.compute((m, n), lambda i, j: (A0[i, j] + 2, A0[i, j] * 3), name="B")\nA1 = te.placeholder((m, n), name="A1")\nC = te.compute((m, n), lambda i, j: A1[i, j] + B0[i, j], name="C")\n\ns = te.create_schedule(C.op)\ns[B0].compute_at(s[C], C.op.axis[0])\n# \u751f\u6210\u7684 IR \u4ee3\u7801\uff1a\nprint(tvm.lower(s, [A0, A1, C], simple_mode=True))<\/code><\/pre>\n
\u8f93\u51fa\u7ed3\u679c\uff1a<\/p>\n
@main = primfn(A0_1: handle, A1_1: handle, C_1: handle) -> ()\n  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}\n  buffers = {A0: Buffer(A0_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"),\n             A1: Buffer(A1_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"),\n             C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto")}\n  buffer_map = {A0_1: A0, A1_1: A1, C_1: C}\n  preflattened_buffer_map = {A0_1: A0_3: Buffer(A0_2, float32, [m, n: int32], [stride, stride_3: int32], type="auto"), A1_1: A1_3: Buffer(A1_2, float32, [m, n], [stride_1, stride_4: int32], type="auto"), C_1: C_3: Buffer(C_2, float32, [m, n], [stride_2, stride_5: int32], type="auto")} {\n  allocate(B.v0: Pointer(global float32), float32, [n]), storage_scope = global;\n  allocate(B.v1: Pointer(global float32), float32, [n]), storage_scope = global;\n  for (i: int32, 0, m) {\n    for (j: int32, 0, n) {\n      B.v0_1: Buffer(B.v0, float32, [n], [])[j] = (A0[((i*stride) + (j*stride_3))] + 2f32)\n      B.v1_1: Buffer(B.v1, float32, [n], [])[j] = (A0[((i*stride) + (j*stride_3))]*3f32)\n    }\n    for (j_1: int32, 0, n) {\n      C[((i*stride_2) + (j_1*stride_5))] = (A1[((i*stride_1) + (j_1*stride_4))] + B.v0_1[j_1])\n    }\n  }\n}<\/code><\/pre>\n
\u603b\u7ed3<\/h2>\n
\u672c\u6559\u7a0b\u4ecb\u7ecd\u5143\u7ec4\u8f93\u5165\u64cd\u4f5c\u7684\u7528\u6cd5\u3002<\/p>\n