@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j_0 in tir.grid(128, 32):\n",
" for j_1_init in tir.serial(4):\n",
" with tir.block("Y_init"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 4 + j_1_init)\n",
" tir.reads()\n",
" tir.writes(Y[vi, vj])\n",
" Y[vi, vj] = tir.float32(0)\n",
" for k, j_1 in tir.grid(128, 4):\n",
" with tir.block("Y_update"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 4 + j_1)\n",
" vk = tir.axis.reduce(128, k)\n",
" tir.reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for ax0 in tir.serial(4):\n",
" with tir.block("C"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 4 + ax0)\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j, k in tir.grid(128, 128, 128):\n",
" with tir.block("Y"):\n",
" vi, vj, vk = tir.axis.remap("SSR", [i, j, k])\n",
" tir.reads(A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" with tir.init():\n",
" Y[vi, vj] = tir.float32(0)\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for i, j in tir.grid(128, 128):\n",
" with tir.block("C"):\n",
" vi, vj = tir.axis.remap("SS", [i, j])\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j_0 in tir.grid(128, 32):\n",
" for j_1_init in tir.serial(4):\n",
" with tir.block("Y_init"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 4 + j_1_init)\n",
" tir.reads()\n",
" tir.writes(Y[vi, vj])\n",
" Y[vi, vj] = tir.float32(0)\n",
" for k, j_1 in tir.grid(128, 4):\n",
" with tir.block("Y_update"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 4 + j_1)\n",
" vk = tir.axis.reduce(128, k)\n",
" tir.reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for ax0 in tir.serial(4):\n",
" with tir.block("C"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 4 + ax0)\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j_0 in tir.grid(128, 16):\n",
" for k, j_1 in tir.grid(128, 8):\n",
" with tir.block("Y"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 8 + j_1)\n",
" vk = tir.axis.reduce(128, k)\n",
" tir.reads(A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" with tir.init():\n",
" Y[vi, vj] = tir.float32(0)\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for ax0 in tir.serial(8):\n",
" with tir.block("C"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 8 + ax0)\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j_0 in tir.grid(128, 15):\n",
" for k, j_1 in tir.grid(128, 9):\n",
" with tir.block("Y"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 9 + j_1)\n",
" vk = tir.axis.reduce(128, k)\n",
" tir.where(j_0 * 9 + j_1 < 128)\n",
" tir.reads(A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" with tir.init():\n",
" Y[vi, vj] = tir.float32(0)\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for ax0 in tir.serial(9):\n",
" with tir.block("C"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 9 + ax0)\n",
" tir.where(j_0 * 9 + ax0 < 128)\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j_0 in tir.grid(128, 4):\n",
" for k, j_1 in tir.grid(128, 32):\n",
" with tir.block("Y"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 32 + j_1)\n",
" vk = tir.axis.reduce(128, k)\n",
" tir.reads(A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" with tir.init():\n",
" Y[vi, vj] = tir.float32(0)\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for ax0 in tir.serial(32):\n",
" with tir.block("C"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 32 + ax0)\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i, j_0 in tir.grid(128, 64):\n",
" for k, j_1 in tir.grid(128, 2):\n",
" with tir.block("Y"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 2 + j_1)\n",
" vk = tir.axis.reduce(128, k)\n",
" tir.reads(A[vi, vk], B[vk, vj])\n",
" tir.writes(Y[vi, vj])\n",
" with tir.init():\n",
" Y[vi, vj] = tir.float32(0)\n",
" Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
" for ax0 in tir.serial(2):\n",
" with tir.block("C"):\n",
" vi = tir.axis.spatial(128, i)\n",
" vj = tir.axis.spatial(128, j_0 * 2 + ax0)\n",
" tir.reads(Y[vi, vj])\n",
" tir.writes(C[vi, vj])\n",
" C[vi, vj] = tir.max(Y[vi, vj], tir.float32(0))\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @tir.prim_func\n",
" def mm_relu(A: tir.Buffer[(128, 128), "float32"], B: tir.Buffer[(128, 128), "float32"], C: tir.Buffer[(128, 128), "float32"]) -> None:\n",
" # function attr dict\n",
" tir.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})\n",
" # body\n",
" # with tir.block("root")\n",
" Y = tir.alloc_buffer([128, 128], dtype="float32")\n",
" for i0, i1, i2 in tir.grid(128, 128, 128):\n",
" with tir.block("Y"):\n",
" i, j, k = tir.axis.remap("SSR", [i0, i1, i2])\n",
" tir.reads(A[i, k], B[k, j])\n",
" tir.writes(Y[i, j])\n",
" with tir.init():\n",
" Y[i, j] = tir.float32(0)\n",
" Y[i, j] = Y[i, j] + A[i, k] * B[k, j]\n",
" for i0, i1 in tir.grid(128, 128):\n",
" with tir.block("C"):\n",
" i, j = tir.axis.remap("SS", [i0, i1])\n",
" tir.reads(Y[i, j])\n",
" tir.writes(C[i, j])\n",
" C[i, j] = tir.max(Y[i, j], tir.float32(0))\n",
" \n",
"