{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "view-in-github"
},
"source": [
"![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Mpn1ti5Urdsv"
},
"source": [
"# Ep7: GPU and Hardware Acceleration, Part 1"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qXysoqn-vZuF"
},
"source": [
"## Install packages \n",
"\n",
"For this course, we will use some ongoing development in tvm, which is an open-source machine learning compilation framework. We provide the following command to install a packaged version for mlc course. This particular notebook depends on a CUDA11 environment."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Xe3vClsD9jlq",
"outputId": "b9c660d4-8bea-4921-c8b5-0ac61536ac21"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
"Looking in links: https://mlc.ai/wheels\n",
"Collecting mlc-ai-nightly-cu110\n",
" Downloading https://github.com/mlc-ai/utils/releases/download/v0.9.dev0/mlc_ai_nightly_cu110-0.9.dev1956%2Bge3f218d71-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (237.1 MB)\n",
"\u001b[K |████████████████████████████████| 237.1 MB 11 kB/s \n",
"\u001b[?25hRequirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (1.7.3)\n",
"Requirement already satisfied: Pygments in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (2.6.1)\n",
"Requirement already satisfied: attrs in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (22.1.0)\n",
"Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (4.4.2)\n",
"Requirement already satisfied: cloudpickle in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (1.3.0)\n",
"Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (1.21.6)\n",
"Requirement already satisfied: tornado in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (5.1.1)\n",
"Requirement already satisfied: psutil in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly-cu110) (5.4.8)\n",
"Collecting synr==0.6.0\n",
" Downloading synr-0.6.0-py3-none-any.whl (18 kB)\n",
"Installing collected packages: synr, mlc-ai-nightly-cu110\n",
"Successfully installed mlc-ai-nightly-cu110-0.9.dev1956+ge3f218d71 synr-0.6.0\n"
]
}
],
"source": [
"!python3 -m pip install mlc-ai-nightly-cu110 -f https://mlc.ai/wheels"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "IH1GR0DjvUiM",
"outputId": "aaab791a-7ba9-4f67-bd47-bded876bb89e"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Mon Aug 29 20:06:46 2022 \n",
"+-----------------------------------------------------------------------------+\n",
"| NVIDIA-SMI 515.65.01 Driver Version: 515.65.01 CUDA Version: 11.7 |\n",
"|-------------------------------+----------------------+----------------------+\n",
"| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n",
"| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n",
"| | | MIG M. |\n",
"|===============================+======================+======================|\n",
"| 0 NVIDIA GeForce ... On | 00000000:03:00.0 On | N/A |\n",
"| 16% 38C P8 14W / 250W | 3MiB / 11264MiB | 0% Default |\n",
"| | | N/A |\n",
"+-------------------------------+----------------------+----------------------+\n",
"| 1 NVIDIA GeForce ... On | 00000000:04:00.0 Off | N/A |\n",
"| 16% 32C P8 20W / 250W | 3MiB / 11264MiB | 0% Default |\n",
"| | | N/A |\n",
"+-------------------------------+----------------------+----------------------+\n",
" \n",
"+-----------------------------------------------------------------------------+\n",
"| Processes: |\n",
"| GPU GI CI PID Type Process name GPU Memory |\n",
"| ID ID Usage |\n",
"|=============================================================================|\n",
"| No running processes found |\n",
"+-----------------------------------------------------------------------------+\n"
]
}
],
"source": [
"!nvidia-smi"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i-14C4skxIrJ"
},
"source": [
"## Prelude\n",
"\n",
"In the past chapter, we discussed MLC flows in CPU environments. This chapter will discuss how to bring some of the optimizations onto GPU. We are going to use CUDA terminology. However, the same set of concepts applies to other kinds of GPUs as well.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BBIuE2jc1DaU"
},
"source": [
"## Preparations\n",
"\n",
"To begin with, let us import the necessary dependencies.\n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"id": "BVp0fHyRkYj6"
},
"outputs": [],
"source": [
"import tvm\n",
"from tvm.ir.module import IRModule\n",
"from tvm.script import tir as T, relax as R\n",
"from tvm import relax\n",
"import numpy as np\n",
"\n",
"# This is needed for deferring annotation parsing in TVMScript\n",
"from __future__ import annotations "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gbu5EAE00kBs"
},
"source": [
"## GPU Architecture\n",
"\n",
"Let us begin by reviewing what a GPU architecture looks like. A typical GPU contains a collection of stream multi-processors, and each multi-processor has many cores. A GPU device is massively parallel and allows us to execute many tasks concurrently."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1K801uQgU0I3"
},
"source": [
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yt0tVBCtX9Wi"
},
"source": [
"To program a GPU, we need to create a set of thread blocks, with each thread mapping to the cores and the thread block map to the stream multiprocessors."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ixU4dDosX2mt"
},
"source": [
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rCSVwxo2fmP-"
},
"source": [
"Let us start GPU programming using a vector add example. The following TensorIR program takes two vectors, A and B, performs element-wise add, and stores the result in C."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"id": "D2vIuJ2i0mo3"
},
"outputs": [],
"source": [
"@tvm.script.ir_module\n",
"class MyModuleVecAdd:\n",
" @T.prim_func\n",
" def main(A: T.Buffer[(1024,), \"float32\"], \n",
" B: T.Buffer[(1024,), \"float32\"], \n",
" C: T.Buffer[(1024,), \"float32\"]) -> None:\n",
" T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
" for i in T.grid(1024):\n",
" with T.block(\"C\"):\n",
" vi = T.axis.remap(\"S\", [i])\n",
" C[vi] = A[vi] + B[vi]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KgknjqDSVBzJ"
},
"source": [
"We first split loop `i` into two loops."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "lBko3swU1C4I",
"outputId": "d25f72f9-d4d9-4c32-820c-e456f2ace26c"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], C: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0, i_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m128\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A[vi], B[vi])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C[vi])\n",
" C[vi] \u001b[38;5;129;01m=\u001b[39;00m A[vi] \u001b[38;5;129;01m+\u001b[39;00m B[vi]\n",
" \n",
"\n"
]
}
],
"source": [
"sch = tvm.tir.Schedule(MyModuleVecAdd)\n",
"block_C = sch.get_block(\"C\")\n",
"i, = sch.get_loops(block=block_C)\n",
"i0, i1 = sch.split(i, [None, 128])\n",
"sch.mod.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9DAXMG9m3nqF"
},
"source": [
"### GPU Thread Blocks"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "XA2FvECSYWVP"
},
"source": [
"Then we bind the iterators to the GPU thread blocks. Each thread is parameterized by two indices -- `threadIdx.x` and `blockIdx.x`. In practice, we can have multiple dimensional thread indices, but we keep them simple as one dimension."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qApYOsF9hFn9"
},
"source": [
""
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "4giP-eG-1bJX",
"outputId": "1673704a-5906-4ab6-8047-f44cde6ca956"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], C: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m8\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m128\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A[vi], B[vi])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C[vi])\n",
" C[vi] \u001b[38;5;129;01m=\u001b[39;00m A[vi] \u001b[38;5;129;01m+\u001b[39;00m B[vi]\n",
" \n",
"\n"
]
}
],
"source": [
"sch.bind(i0, \"blockIdx.x\")\n",
"sch.bind(i1, \"threadIdx.x\")\n",
"sch.mod.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i9lEEl4-omtG"
},
"source": [
"### Build and Run the TensorIR Function on GPU\n",
"\n",
"We can build and test out the resulting function on the GPU."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "3hL1tGI-3DWf",
"outputId": "ba0f5406-7ea8-4b6a-843f-bbf626ee4226"
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[20:07:11] /workspace/tvm/src/target/target_kind.cc:163: Warning: Unable to detect CUDA version, default to \"-arch=sm_20\" instead\n"
]
},
{
"ename": "TVMError",
"evalue": "Traceback (most recent call last):\n 4: TVMFuncCall\n 3: _ZN3tvm7runtime13PackedFuncObj9ExtractorINS0_16PackedFuncSubObjIZNS0_15TypedPackedFuncIFNS0_6ModuleERKNS0_3MapINS_6TargetENS_8IRModuleEvvEES7_EE17AssignTypedLambdaINS_L10__mk_TVM17MUlSB_S7_E_EEEvT_SsEUlRKNS0_7TVMArgsEPNS0_11TVMRetValueEE_EEE4CallEPKS1_SH_SL_\n 2: tvm::TIRToRuntime(tvm::runtime::Map const&, tvm::Target const&)\n 1: tvm::codegen::Build(tvm::IRModule, tvm::Target)\n 0: _ZN3tvm7runtime6deta\n File \"/workspace/tvm/src/target/codegen.cc\", line 58\nTVMError: \n---------------------------------------------------------------\nAn error occurred during the execution of TVM.\nFor more information, please see: https://tvm.apache.org/docs/errors.html\n---------------------------------------------------------------\n Check failed: (bf != nullptr) is false: target.build.cuda is not enabled",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mTVMError\u001b[0m Traceback (most recent call last)",
"\u001b[1;32m/media/pc/data/4tb/lxw/books/mlcbook/doc/study/7_GPU_and_Specialized_Hardware.ipynb Cell 22\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m rt_mod \u001b[39m=\u001b[39m tvm\u001b[39m.\u001b[39;49mbuild(sch\u001b[39m.\u001b[39;49mmod, target\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mcuda\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m 3\u001b[0m A_np \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mrandom\u001b[39m.\u001b[39muniform(size\u001b[39m=\u001b[39m(\u001b[39m1024\u001b[39m,))\u001b[39m.\u001b[39mastype(\u001b[39m\"\u001b[39m\u001b[39mfloat32\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m 4\u001b[0m B_np \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mrandom\u001b[39m.\u001b[39muniform(size\u001b[39m=\u001b[39m(\u001b[39m1024\u001b[39m,))\u001b[39m.\u001b[39mastype(\u001b[39m\"\u001b[39m\u001b[39mfloat32\u001b[39m\u001b[39m\"\u001b[39m)\n",
"File \u001b[0;32m/media/pc/data/4tb/lxw/libs/anaconda3/envs/mlc/lib/python3.10/site-packages/tvm/driver/build_module.py:282\u001b[0m, in \u001b[0;36mbuild\u001b[0;34m(inputs, args, target, target_host, runtime, name, binds)\u001b[0m\n\u001b[1;32m 278\u001b[0m target_host \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mllvm\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mif\u001b[39;00m tvm\u001b[39m.\u001b[39mruntime\u001b[39m.\u001b[39menabled(\u001b[39m\"\u001b[39m\u001b[39mllvm\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mstackvm\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 280\u001b[0m annotated_mods, target_host \u001b[39m=\u001b[39m Target\u001b[39m.\u001b[39mcanon_target_map_and_host(annotated_mods, target_host)\n\u001b[0;32m--> 282\u001b[0m rt_mod_host \u001b[39m=\u001b[39m _driver_ffi\u001b[39m.\u001b[39;49mtir_to_runtime(annotated_mods, target_host)\n\u001b[1;32m 284\u001b[0m annotated_mods, target_host \u001b[39m=\u001b[39m Target\u001b[39m.\u001b[39mcanon_target_map_and_host(annotated_mods, target_host)\n\u001b[1;32m 286\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39misinstance\u001b[39m(target_host, Target):\n",
"File \u001b[0;32m/media/pc/data/4tb/lxw/libs/anaconda3/envs/mlc/lib/python3.10/site-packages/tvm/_ffi/_ctypes/packed_func.py:237\u001b[0m, in \u001b[0;36mPackedFuncBase.__call__\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m 225\u001b[0m ret_tcode \u001b[39m=\u001b[39m ctypes\u001b[39m.\u001b[39mc_int()\n\u001b[1;32m 226\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m 227\u001b[0m _LIB\u001b[39m.\u001b[39mTVMFuncCall(\n\u001b[1;32m 228\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mhandle,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 235\u001b[0m \u001b[39m!=\u001b[39m \u001b[39m0\u001b[39m\n\u001b[1;32m 236\u001b[0m ):\n\u001b[0;32m--> 237\u001b[0m \u001b[39mraise\u001b[39;00m get_last_ffi_error()\n\u001b[1;32m 238\u001b[0m _ \u001b[39m=\u001b[39m temp_args\n\u001b[1;32m 239\u001b[0m _ \u001b[39m=\u001b[39m args\n",
"\u001b[0;31mTVMError\u001b[0m: Traceback (most recent call last):\n 4: TVMFuncCall\n 3: _ZN3tvm7runtime13PackedFuncObj9ExtractorINS0_16PackedFuncSubObjIZNS0_15TypedPackedFuncIFNS0_6ModuleERKNS0_3MapINS_6TargetENS_8IRModuleEvvEES7_EE17AssignTypedLambdaINS_L10__mk_TVM17MUlSB_S7_E_EEEvT_SsEUlRKNS0_7TVMArgsEPNS0_11TVMRetValueEE_EEE4CallEPKS1_SH_SL_\n 2: tvm::TIRToRuntime(tvm::runtime::Map const&, tvm::Target const&)\n 1: tvm::codegen::Build(tvm::IRModule, tvm::Target)\n 0: _ZN3tvm7runtime6deta\n File \"/workspace/tvm/src/target/codegen.cc\", line 58\nTVMError: \n---------------------------------------------------------------\nAn error occurred during the execution of TVM.\nFor more information, please see: https://tvm.apache.org/docs/errors.html\n---------------------------------------------------------------\n Check failed: (bf != nullptr) is false: target.build.cuda is not enabled"
]
}
],
"source": [
"rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
"\n",
"A_np = np.random.uniform(size=(1024,)).astype(\"float32\")\n",
"B_np = np.random.uniform(size=(1024,)).astype(\"float32\")\n",
"A_nd = tvm.nd.array(A_np, tvm.cuda(0))\n",
"B_nd = tvm.nd.array(B_np, tvm.cuda(0))\n",
"C_nd = tvm.nd.array(np.zeros((1024,), dtype=\"float32\"), tvm.cuda(0))\n",
"\n",
"rt_mod[\"main\"](A_nd, B_nd, C_nd)\n",
"print(A_nd)\n",
"print(B_nd)\n",
"print(C_nd)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"id": "NJjgpGtx3b4m"
},
"outputs": [],
"source": [
"np.testing.assert_allclose(C_nd.numpy(), A_np + B_np)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IspVTqVR3tBM"
},
"source": [
"## Window Sum Example"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "u8_oVnfQj6CX"
},
"source": [
"Now, let us move forward to another example -- window sum. This program can be viewed as a basic version of \"convolution\" with a predefined weight `[1,1,1]`. We are taking sliding over the input and add three neighboring values together."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Sbpk4tTHYfCx"
},
"source": [
""
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"id": "j_TJgnwBBNyq"
},
"outputs": [],
"source": [
"@tvm.script.ir_module\n",
"class MyModuleWindowSum:\n",
" @T.prim_func\n",
" def main(A: T.Buffer[(1027,), \"float32\"], \n",
" B: T.Buffer[(1024,), \"float32\"]) -> None:\n",
" T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
" for i in T.grid(1024):\n",
" with T.block(\"C\"):\n",
" vi = T.axis.remap(\"S\", [i])\n",
" B[vi] = A[vi] + A[vi + 1] + A[vi + 2]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MNiC6Aczj9PG"
},
"source": [
"First, we can bind the loop to GPU threads."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "sInzhG_cCUA-",
"outputId": "0812ad60-f8ed-4f33-818e-5df4cdbf0202"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1027\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m8\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m128\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A[vi : vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m3\u001b[39m])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(B[vi])\n",
" B[vi] \u001b[38;5;129;01m=\u001b[39;00m A[vi] \u001b[38;5;129;01m+\u001b[39;00m A[vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m1\u001b[39m] \u001b[38;5;129;01m+\u001b[39;00m A[vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m2\u001b[39m]\n",
" \n",
"\n"
]
}
],
"source": [
"sch = tvm.tir.Schedule(MyModuleWindowSum)\n",
"nthread = 128\n",
"block_C = sch.get_block(\"C\")\n",
"i, = sch.get_loops(block=block_C)\n",
"i0, i1 = sch.split(i, [None, nthread])\n",
"sch.bind(i0, \"blockIdx.x\")\n",
"sch.bind(i1, \"threadIdx.x\")\n",
"sch.mod.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dR5UPy_ml-9F"
},
"source": [
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2YBIc5hWl2sW"
},
"source": [
"Importantly, in this case, there are reuse opportunities. Remember that each GPU thread block contains shared memory that all threads can access within the block. We use `cache_read` to add an intermediate stage that caches segments (in green below) onto the shared memory. After the caching is finished, the threads can then read from the shared memory."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JotZy40CYiei"
},
"source": [
""
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "r0W6HSycCbxP",
"outputId": "cb2abd3c-0cc6-4f77-df34-ff853e812092"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1027\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" A_shared \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00malloc_buffer([\u001b[38;5;28m1027\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, scope\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshared\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m8\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m128\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m130\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mA_shared\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" v0 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1027\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A[v0])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(A_shared[v0])\n",
" A_shared[v0] \u001b[38;5;129;01m=\u001b[39;00m A[v0]\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A_shared[vi : vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m3\u001b[39m])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(B[vi])\n",
" B[vi] \u001b[38;5;129;01m=\u001b[39;00m A_shared[vi] \u001b[38;5;129;01m+\u001b[39;00m A_shared[vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m1\u001b[39m] \u001b[38;5;129;01m+\u001b[39;00m A_shared[vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m2\u001b[39m]\n",
" \n",
"\n"
]
}
],
"source": [
"A_shared = sch.cache_read(block_C, read_buffer_index=0, storage_scope=\"shared\")\n",
"sch.compute_at(A_shared, i1)\n",
"sch.mod.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JrguFDF9nksn"
},
"source": [
"Because the memory is shared across threads, we need to resplit the loop and bind the inner iterator of the fetching process onto the thread indices. This technique is called **cooperative fetching**, where multiple threads work together to bring the data onto the shared memory. The following reading process can be different. "
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "G_d0Hd-pGsIs",
"outputId": "38e59c53-e33c-4507-cdee-ba13f2b3ce09"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1027\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" A_shared \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00malloc_buffer([\u001b[38;5;28m1027\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, scope\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshared\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m8\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m128\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m2\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m128\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mA_shared\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" v0 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1027\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m (ax0_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_1))\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwhere(ax0_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_1 \u001b[38;5;129;01m<\u001b[39;00m \u001b[38;5;28m130\u001b[39m)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A[v0])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(A_shared[v0])\n",
" A_shared[v0] \u001b[38;5;129;01m=\u001b[39;00m A[v0]\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m128\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A_shared[vi : vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m3\u001b[39m])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(B[vi])\n",
" B[vi] \u001b[38;5;129;01m=\u001b[39;00m A_shared[vi] \u001b[38;5;129;01m+\u001b[39;00m A_shared[vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m1\u001b[39m] \u001b[38;5;129;01m+\u001b[39;00m A_shared[vi \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m2\u001b[39m]\n",
" \n",
"\n"
]
}
],
"source": [
"ax = sch.get_loops(A_shared)[-1]\n",
"ax0, ax1 = sch.split(ax, [None, nthread])\n",
"sch.bind(ax1, \"threadIdx.x\")\n",
"sch.mod.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "g_df2KCDovQc"
},
"source": [
"We can inspect the corresponding low-level code (in CUDA). The generated code contains two parts: \n",
"- A host part that calls into the GPU driver\n",
"- A cuda kernel that runs the corresponding computation.\n",
"\n",
"We can print out the cuda kernel using the following code. We still need both the host and kernel code to run the program, so it is only a quick way to inspect what the final code generation result.\n",
"\n",
"Notably, the build process automatically compacts the shared memory stage to use a minimum region used within the thread block."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "2ngiuLMSHWM5",
"outputId": "258b6c4e-daae-4de6-dd94-536a4af56418"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"#ifdef _WIN32\n",
" using uint = unsigned int;\n",
" using uchar = unsigned char;\n",
" using ushort = unsigned short;\n",
" using int64_t = long long;\n",
" using uint64_t = unsigned long long;\n",
"#else\n",
" #define uint unsigned int\n",
" #define uchar unsigned char\n",
" #define ushort unsigned short\n",
" #define int64_t long long\n",
" #define uint64_t unsigned long long\n",
"#endif\n",
"extern \"C\" __global__ void __launch_bounds__(128) main_kernel0(float* __restrict__ A, float* __restrict__ B) {\n",
" __shared__ float A_shared[130];\n",
" for (int ax0_0 = 0; ax0_0 < 2; ++ax0_0) {\n",
" if (((ax0_0 * 64) + (((int)threadIdx.x) >> 1)) < 65) {\n",
" A_shared[((ax0_0 * 128) + ((int)threadIdx.x))] = A[(((((int)blockIdx.x) * 128) + (ax0_0 * 128)) + ((int)threadIdx.x))];\n",
" }\n",
" }\n",
" __syncthreads();\n",
" B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = ((A_shared[((int)threadIdx.x)] + A_shared[(((int)threadIdx.x) + 1)]) + A_shared[(((int)threadIdx.x) + 2)]);\n",
"}\n",
"\n",
"\n"
]
}
],
"source": [
"rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
"print(rt_mod.imported_modules[0].get_source())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uW087bCZbwnX"
},
"source": [
"### Build Code for Other GPU Platforms\n",
"\n",
"A MLC process usually support targeting multiple kinds of hardware platforms, we can generate opencl code(which is another kind of GPU programming model) by changing the target parameter."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "T-l21FKVb1ra",
"outputId": "c61c5f4b-de82-4587-b6fa-bafebfa56627"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"// Function: main_kernel0\n",
"#include \n",
"using namespace metal;\n",
"\n",
"union __TVMArgUnion {\n",
" int v_int[2];\n",
"};\n",
"\n",
"kernel void main_kernel0( device float* A [[ buffer(0) ]],\n",
" device float* B [[ buffer(1) ]],\n",
" uint blockIdx [[threadgroup_position_in_grid]],\n",
" uint threadIdx [[thread_position_in_threadgroup]]\n",
") {\n",
" threadgroup float A_shared[130];\n",
" for (int ax0_0 = 0; ax0_0 < 2; ++ax0_0) {\n",
" if (((ax0_0 * 64) + (((int)threadIdx) >> 1)) < 65) {\n",
" A_shared[((ax0_0 * 128) + ((int)threadIdx))] = A[(((((int)blockIdx) * 128) + (ax0_0 * 128)) + ((int)threadIdx))];\n",
" }\n",
" }\n",
" threadgroup_barrier(mem_flags::mem_threadgroup);\n",
" B[((((int)blockIdx) * 128) + ((int)threadIdx))] = ((A_shared[((int)threadIdx)] + A_shared[(((int)threadIdx) + 1)]) + A_shared[(((int)threadIdx) + 2)]);\n",
"}\n",
"\n",
"\n"
]
}
],
"source": [
"rt_mod = tvm.build(sch.mod, target=\"metal\")\n",
"print(rt_mod.imported_modules[0].get_source())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8yH4IMSMvF9o"
},
"source": [
"## Matrix Multiplication\n",
"\n",
"Let us now get to something slightly more complicated and try out optimizing matrix multiplication on GPU. We will go over two common techniques for GPU performance optimization.\n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"id": "XwRUCQBtxNh3"
},
"outputs": [],
"source": [
"@tvm.script.ir_module\n",
"class MyModuleMatmul:\n",
" @T.prim_func\n",
" def main(A: T.Buffer[(1024, 1024), \"float32\"], \n",
" B: T.Buffer[(1024, 1024), \"float32\"], \n",
" C: T.Buffer[(1024, 1024), \"float32\"]) -> None:\n",
" T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
" for i, j, k in T.grid(1024, 1024, 1024):\n",
" with T.block(\"C\"):\n",
" vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
" with T.init():\n",
" C[vi, vj] = 0.0\n",
" C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OWx5swIIqrJJ"
},
"source": [
"### Local Blocking\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wdECOWAIZ64k"
},
"source": [
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "R2gq6CqTrOdQ"
},
"source": [
"To increase overall memory reuse. We can tile the loops. In particular, we introduce local tiles such that we only need to load stripe of data from A and B once, then use them to perform a `V * V` matrix multiplication result.\n",
"\n",
"This local tiling helps to reduce the memory pressure, as each element in the stripe is reused `V` times."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "k8TSH22_KGEc",
"outputId": "02965ff7-3884-41e0-9e4b-6d7c25742670"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[(\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[(\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], C: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[(\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" C_local \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00malloc_buffer([\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, scope\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlocal\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m16\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.y\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m j_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m16\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m8\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.y\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m j_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m8\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_2_init, j_2_init \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC_init\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_2_init)\n",
" vj \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_2_init)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads()\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C_local[vi, vj])\n",
" C_local[vi, vj] \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mfloat32(\u001b[38;5;28m0\u001b[39m)\n",
" \u001b[38;5;28;01mfor\u001b[39;00m k_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m256\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m k_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00munroll(\u001b[38;5;28m4\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_2, j_2 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC_update\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_2)\n",
" vj \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_2)\n",
" vk \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mreduce(\u001b[38;5;28m1024\u001b[39m, k_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m4\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m k_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(C_local[vi, vj], A[vi, vk], B[vk, vj])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C_local[vi, vj])\n",
" C_local[vi, vj] \u001b[38;5;129;01m=\u001b[39;00m C_local[vi, vj] \u001b[38;5;129;01m+\u001b[39;00m A[vi, vk] \u001b[38;5;129;01m*\u001b[39;00m B[vk, vj]\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0, ax1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC_local\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" v0 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0)\n",
" v1 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(C_local[v0, v1])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C[v0, v1])\n",
" C[v0, v1] \u001b[38;5;129;01m=\u001b[39;00m C_local[v0, v1]\n",
" \n",
"\n"
]
}
],
"source": [
"def blocking(sch, \n",
" tile_local_y, \n",
" tile_local_x, \n",
" tile_block_y, \n",
" tile_block_x,\n",
" tile_k):\n",
" block_C = sch.get_block(\"C\")\n",
" C_local = sch.cache_write(block_C, 0, \"local\")\n",
"\n",
" i, j, k = sch.get_loops(block=block_C)\n",
"\n",
" i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])\n",
" j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])\n",
" k0, k1 = sch.split(loop=k, factors=[None, tile_k])\n",
" sch.unroll(k1)\n",
" sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)\n",
" sch.reverse_compute_at(C_local, j1)\n",
"\n",
" sch.bind(i0, \"blockIdx.y\")\n",
" sch.bind(j0, \"blockIdx.x\")\n",
"\n",
" sch.bind(i1, \"threadIdx.y\")\n",
" sch.bind(j1, \"threadIdx.x\")\n",
" sch.decompose_reduction(block_C, k0)\n",
"\n",
" return sch\n",
"\n",
"sch = tvm.tir.Schedule(MyModuleMatmul)\n",
"sch = blocking(sch, 8, 8, 8, 8, 4)\n",
"sch.mod.show()"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Qu-JY7JUI4us",
"outputId": "64d8a8f6-0c74-4c76-a345-825a741550ed"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"GEMM-Blocking: 2378.997450 GFLOPS\n"
]
}
],
"source": [
"rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
"dev = tvm.cuda(0)\n",
"A_np = np.random.uniform(size=(1024, 1024)).astype(\"float32\")\n",
"B_np = np.random.uniform(size=(1024, 1024)).astype(\"float32\")\n",
"A_nd = tvm.nd.array(A_np, dev)\n",
"B_nd = tvm.nd.array(B_np, dev)\n",
"C_nd = tvm.nd.array(np.zeros((1024, 1024), dtype=\"float32\"), dev)\n",
"\n",
"num_flop = 2 * 1024 * 1024 * 1024\n",
"evaluator = rt_mod.time_evaluator(\"main\", dev, number=10)\n",
"\n",
"\n",
"print(\"GEMM-Blocking: %f GFLOPS\" % (num_flop / evaluator(A_nd, B_nd, C_nd).mean / 1e9))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LQhlKMMIbgqq"
},
"source": [
"## Shared Memory Blocking\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pK7d8BY8aFUz"
},
"source": [
"\n",
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e8CP3jKIrvs_"
},
"source": [
"Our first attempt did not consider the neighboring threads which sit in the same GPU thread block, and we can load the data they commonly need into a piece of shared memory.\n",
"\n",
"The following transformation does that."
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "yCiSsNtjxTsq",
"outputId": "9df52e4f-d45b-4ce2-a1df-042e1def162d"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n",
"\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n",
" \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n",
" \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(A: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[(\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], B: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[(\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], C: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[(\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
" \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n",
" T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n",
" \u001b[38;5;30;03m# body\u001b[39;00m\n",
" \u001b[38;5;30;03m# with T.block(\"root\")\u001b[39;00m\n",
" C_local \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00malloc_buffer([\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, scope\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlocal\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
" A_shared \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00malloc_buffer([\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, scope\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshared\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
" B_shared \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00malloc_buffer([\u001b[38;5;28m1024\u001b[39m, \u001b[38;5;28m1024\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, scope\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshared\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m16\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.y\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m j_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m16\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mblockIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_1_j_1_fused \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m64\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m i_2_init, j_2_init \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC_init\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1_j_1_fused \u001b[38;5;129;01m/\u001b[39;00m\u001b[38;5;129;01m/\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_2_init)\n",
" vj \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1_j_1_fused \u001b[38;5;129;01m%\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_2_init)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads()\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C_local[vi, vj])\n",
" C_local[vi, vj] \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mfloat32(\u001b[38;5;28m0\u001b[39m)\n",
" \u001b[38;5;28;01mfor\u001b[39;00m k_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m128\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_ax1_fused_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m2\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_ax1_fused_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m64\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_ax1_fused_2 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mvectorized(\u001b[38;5;28m4\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mA_shared\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" v0 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m (ax0_ax1_fused_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m256\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m4\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_2) \u001b[38;5;129;01m/\u001b[39;00m\u001b[38;5;129;01m/\u001b[39;00m \u001b[38;5;28m8\u001b[39m)\n",
" v1 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, k_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m (ax0_ax1_fused_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m256\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m4\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_2) \u001b[38;5;129;01m%\u001b[39;00m \u001b[38;5;28m8\u001b[39m)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(A[v0, v1])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(A_shared[v0, v1])\n",
" A_shared[v0, v1] \u001b[38;5;129;01m=\u001b[39;00m A[v0, v1]\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_ax1_fused_0 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m2\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_ax1_fused_1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mthread_binding(\u001b[38;5;28m64\u001b[39m, thread\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mthreadIdx.x\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0_ax1_fused_2 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mvectorized(\u001b[38;5;28m4\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mB_shared\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" v0 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, k_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m (ax0_ax1_fused_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m256\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m4\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_2) \u001b[38;5;129;01m/\u001b[39;00m\u001b[38;5;129;01m/\u001b[39;00m \u001b[38;5;28m64\u001b[39m)\n",
" v1 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m (ax0_ax1_fused_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m256\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_1 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m4\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0_ax1_fused_2) \u001b[38;5;129;01m%\u001b[39;00m \u001b[38;5;28m64\u001b[39m)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(B[v0, v1])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(B_shared[v0, v1])\n",
" B_shared[v0, v1] \u001b[38;5;129;01m=\u001b[39;00m B[v0, v1]\n",
" \u001b[38;5;28;01mfor\u001b[39;00m k_1, i_2, j_2 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC_update\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" vi \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1_j_1_fused \u001b[38;5;129;01m/\u001b[39;00m\u001b[38;5;129;01m/\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_2)\n",
" vj \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1_j_1_fused \u001b[38;5;129;01m%\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m j_2)\n",
" vk \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mreduce(\u001b[38;5;28m1024\u001b[39m, k_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m k_1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(C_local[vi, vj], A_shared[vi, vk], B_shared[vk, vj])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C_local[vi, vj])\n",
" C_local[vi, vj] \u001b[38;5;129;01m=\u001b[39;00m C_local[vi, vj] \u001b[38;5;129;01m+\u001b[39;00m A_shared[vi, vk] \u001b[38;5;129;01m*\u001b[39;00m B_shared[vk, vj]\n",
" \u001b[38;5;28;01mfor\u001b[39;00m ax0, ax1 \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mgrid(\u001b[38;5;28m8\u001b[39m, \u001b[38;5;28m8\u001b[39m):\n",
" \u001b[38;5;28;01mwith\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mblock(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mC_local\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
" v0 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, i_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1_j_1_fused \u001b[38;5;129;01m/\u001b[39;00m\u001b[38;5;129;01m/\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax0)\n",
" v1 \u001b[38;5;129;01m=\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00maxis\u001b[38;5;129;01m.\u001b[39;00mspatial(\u001b[38;5;28m1024\u001b[39m, j_0 \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m64\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m i_1_j_1_fused \u001b[38;5;129;01m%\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m \u001b[38;5;129;01m+\u001b[39;00m ax1)\n",
" T\u001b[38;5;129;01m.\u001b[39;00mreads(C_local[v0, v1])\n",
" T\u001b[38;5;129;01m.\u001b[39;00mwrites(C[v0, v1])\n",
" C[v0, v1] \u001b[38;5;129;01m=\u001b[39;00m C_local[v0, v1]\n",
" \n",
"\n"
]
}
],
"source": [
"def cache_read_and_coop_fetch(sch, block, nthread, read_idx, read_loc):\n",
" read_cache = sch.cache_read(block=block, read_buffer_index=read_idx, storage_scope=\"shared\")\n",
" sch.compute_at(block=read_cache, loop=read_loc)\n",
" # vectorized cooperative fetch\n",
" inner0, inner1 = sch.get_loops(block=read_cache)[-2:]\n",
" inner = sch.fuse(inner0, inner1)\n",
" _, tx, vec = sch.split(loop=inner, factors=[None, nthread, 4])\n",
" sch.vectorize(vec)\n",
" sch.bind(tx, \"threadIdx.x\")\n",
"\n",
"\n",
"def blocking_with_shared(\n",
" sch, \n",
" tile_local_y, \n",
" tile_local_x, \n",
" tile_block_y, \n",
" tile_block_x,\n",
" tile_k):\n",
" block_C = sch.get_block(\"C\")\n",
" C_local = sch.cache_write(block_C, 0, \"local\")\n",
"\n",
" i, j, k = sch.get_loops(block=block_C)\n",
"\n",
" i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])\n",
" j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])\n",
" k0, k1 = sch.split(loop=k, factors=[None, tile_k])\n",
"\n",
" sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)\n",
" sch.reverse_compute_at(C_local, j1)\n",
"\n",
" sch.bind(i0, \"blockIdx.y\")\n",
" sch.bind(j0, \"blockIdx.x\")\n",
"\n",
" tx = sch.fuse(i1, j1)\n",
" sch.bind(tx, \"threadIdx.x\")\n",
" nthread = tile_block_y * tile_block_x\n",
" cache_read_and_coop_fetch(sch, block_C, nthread, 0, k0)\n",
" cache_read_and_coop_fetch(sch, block_C, nthread, 1, k0) \n",
" sch.decompose_reduction(block_C, k0)\n",
"\n",
" return sch\n",
"\n",
"sch = tvm.tir.Schedule(MyModuleMatmul)\n",
"sch = blocking_with_shared(sch, 8, 8, 8, 8, 8)\n",
"sch.mod.show()\n"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "emNWFAwexXWH",
"outputId": "10a3760d-9b43-4d83-b339-53e022bd1894"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"GEMM-Blocking: 5672.038389 GFLOPS\n"
]
}
],
"source": [
"rt_mod = tvm.build(sch.mod, target=\"cuda\")\n",
"dev = tvm.cuda(0)\n",
"evaluator = rt_mod.time_evaluator(\"main\", dev, number=10)\n",
"\n",
"print(\"GEMM-Blocking: %f GFLOPS\" % (num_flop / evaluator(A_nd, B_nd, C_nd).mean / 1e9))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2aC9IAGGb7xw"
},
"source": [
"## Leveraging Automatic Program Optimization"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uhrY8D8UsE3O"
},
"source": [
"So far, we have been manually writing transformations to optimize the TensorIR program on GPU. We can leverage the automatic program optimization framework to tune the same program. The following code does that, we only set a small number here, and it can take a few min to finish. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"background_save": true,
"base_uri": "https://localhost:8080/"
},
"id": "ANaU_GsMcAP1",
"outputId": "b0f1ae6a-7f2b-4b63-b268-2750fefbb883"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2022-08-06 13:02:40.921 INFO Logging directory: ./tune_tmp/logs\n",
"2022-08-06 13:02:40.928 INFO Logging directory: ./tune_tmp/logs\n",
"2022-08-06 13:02:40.930 INFO Working directory: ./tune_tmp\n",
"2022-08-06 13:02:40.933 INFO Creating JSONDatabase. Workload at: ./tune_tmp/database_workload.json. Tuning records at: ./tune_tmp/database_tuning_record.json\n",
"2022-08-06 13:02:40.935 INFO LocalBuilder: max_workers = 1\n",
"2022-08-06 13:02:41.536 INFO LocalRunner: max_workers = 1\n",
"2022-08-06 13:02:42.113 INFO Initializing Task #0: \"main\"\n",
"2022-08-06 13:02:42.130 INFO \n",
" ID | Name | FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Terminated \n",
"---------------------------------------------------------------------------------------------------------------\n",
" 0 | main | 2147483648 | 1 | N/A | N/A | N/A | 0 | \n",
"---------------------------------------------------------------------------------------------------------------\n",
"Total trials: 0\n",
"Total latency (us): 0\n",
"\n",
"2022-08-06 13:02:42.133 INFO Scheduler picks Task #0: \"main\"\n",
"2022-08-06 13:03:59.506 INFO Sending 64 sample(s) to builder\n"
]
}
],
"source": [
"from tvm import meta_schedule as ms\n",
"\n",
"sch_tuned = ms.tune_tir(\n",
" mod=MyModuleMatmul,\n",
" target=\"nvidia/tesla-p100\",\n",
" config=ms.TuneConfig(\n",
" max_trials_global=64,\n",
" num_trials_per_iter=64,\n",
" ),\n",
" work_dir=\"./tune_tmp\",\n",
" task_name=\"main\"\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MD_C5bkrer6c"
},
"outputs": [],
"source": [
"sch_tuned.mod.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "03Gk8fhRenjU"
},
"outputs": [],
"source": [
"rt_mod = tvm.build(sch_tuned.mod, target=\"nvidia/tesla-p100\")\n",
"dev = tvm.cuda(0)\n",
"evaluator = rt_mod.time_evaluator(\"main\", dev, number=10)\n",
"\n",
"print(\"MetaSchedule: %f GFLOPS\" % (num_flop / evaluator(A_nd, B_nd, C_nd).mean / 1e9))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pZGWq5EjJ0BA"
},
"source": [
"## Summary\n",
"\n",
"This chapter studies another axis of MLC -- how we can transform our program for hardware acceleration. The MLC process helps us to bridge the input models toward different GPU programming models and environments. We will visit more hardware specialization topics in the incoming chapter as well.\n",
"\n",
"- A typical GPU contains two-level hierarchy. Each thread is indexed by(in cuda terminology) `threadIdx.x` and `blockIdx.x`(there can be multiple dimension indices as well, but they can be fused to one.\n",
"- Shared memory helps cache data commonly used across the threads within the same block.\n",
"- Encourage memory reuse during GPU optimization.\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"include_colab_link": true,
"name": "7 GPU and Specialized Hardware.ipynb",
"provenance": []
},
"gpuClass": "standard",
"kernelspec": {
"display_name": "Python 3.10.4 ('mlc': conda)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.4"
},
"vscode": {
"interpreter": {
"hash": "d8a760899c905ec5a15e0d212432af25d7f0b614c7ae634224dffa77837bb03c"
}
}
},
"nbformat": 4,
"nbformat_minor": 0
}
|