TVMOp Tutorial.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Installation and Basic of TVM\n",
    "\n",
    "### [About TVM](https://tvm.ai/about)\n",
    "\n",
    "TVM is an open deep learning compiler stack for CPUs, GPUs, and specialized accelerators. It aims to close the gap between the productivity-focused deep learning frameworks, and the performance- or efficiency-oriented hardware backends. TVM provides the following main features:\n",
    "\n",
    "- Compilation of deep learning models in Keras, MXNet, PyTorch, Tensorflow, CoreML, DarkNet into minimum deployable modules on diverse hardware backends.\n",
    "- Infrastructure to automatic generate and optimize tensor operators on more backend with better performance.\n",
    "\n",
    "![image.png](https://tvm.ai/images/main/tvm-stack.png)\n",
    "\n",
    "### [Background](https://github.com/apache/incubator-mxnet/issues/15465)\n",
    "\n",
    "Currently in MXNet we implement operator kernels in C++. Developers need to specify the detail logic of each computation, which slows down the development process. Given the fact that we’re moving forward to be numpy-compatible, a large amount of operators are to be implemented. Moreover, we also have various of backend to support, including CPUs and GPUs from AMD, ARM, Intel, Nvidia, etc. It requires great effort to implement efficient kernels for each of these hardwares, so as writing test cases for each of the operator+backend combination."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Installation\n",
    "\n",
    "Official Guide: https://docs.tvm.ai/install/from_source.html\n",
    "\n",
    "1. **Install python environment**\n",
    "2. **Download source code**\n",
    "```\n",
    "git clone --recursive https://github.com/dmlc/tvm\n",
    "```\n",
    "3. **Install backend**\n",
    " - Create build directory\n",
    "```\n",
    "cd tvm\n",
    "mkdir build\n",
    "cp cmake/config.cmake build\n",
    "cd build\n",
    "```\n",
    " - Download llvm.\n",
    "```\n",
    "wget http://releases.llvm.org/8.0.0/clang+llvm-8.0.0-x86_64-linux-gnu-ubuntu-16.04.tar.xz\n",
    "tar -xf clang+llvm-8.0.0-x86_64-linux-gnu-ubuntu-16.04.tar.xz\n",
    "mv clang+llvm-8.0.0-x86_64-linux-gnu-ubuntu-16.04 llvm\n",
    "```\n",
    " - Edit file `config.cmake`. Change `set(USE_CUDA OFF)` to `set(USE_CUDA ON)`, and change `set(USE_LLVM OFF)` to `set(USE_LLVM /home/ubuntu/tvm/build/llvm/bin/llvm-config)`\n",
    " - Then build tvm\n",
    "```\n",
    "cd build\n",
    "cmake ..\n",
    "make -j4\n",
    "```\n",
    "4. **Install python frontend**\n",
    "```\n",
    "cd ..\n",
    "cd python && pip install -e . && cd ..\n",
    "cd topi/python && pip install -e . && cd ../..\n",
    "cd nnvm/python && pip install -e . && cd ../..\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Enable TVMOp in MXNet\n",
    "\n",
    "**CMake:** -DUSE_TVM_OP=ON\n",
    "\n",
    "**Make:** USE_TVM_OP = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 Basic Usage\n",
    "\n",
    "TVM uses symbolic API like tensorflow, it just defines the computation flow. In TVM, every node is a `Tensor` and has its operator, contains `placeholder`(its output same as input).\n",
    "\n",
    "Especially, when we use tvm to implement an operator, we don't know any imformation of **shape**. But we know **ndim** and **dtypes**, actually the kernels are generated for each ndim/dtype.\n",
    "\n",
    "We use a simple vector add as example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "\n",
    "A = tvm.placeholder([tvm.var() for _ in range(1)], name='A', dtype=\"float32\")\n",
    "B = tvm.placeholder([tvm.var() for _ in range(1)], name='B', dtype=\"float32\")\n",
    "C = tvm.compute([tvm.var() for _ in range(1)], lambda *i: A[i] + B[i], name='C')\n",
    "# C = tvm.compute([tvm.var() for _ in range(2)], lambda *i: A(*i) + B(*i), name='C')\n",
    "# C = tvm.compute([tvm.var() for _ in range(2)], lambda i, j: A[i, j] + B[i, j], name='C')\n",
    "s = tvm.create_schedule(C.op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tvm.tensor.Tensor"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'float32'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tvm.tensor.PlaceholderOp"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(A.op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tvm.schedule.Schedule, tvm.schedule.Stage)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(s), type(s[C])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "produce C {\n",
      "  for (i0, 0, tindex) {\n",
      "    C[(i0*stride)] = (A[(i0*stride)] + B[(i0*stride)])\n",
      "  }\n",
      "}\n",
      "\n",
      "################################################################################\n",
      "; ModuleID = 'default_function'\n",
      "source_filename = \"default_function\"\n",
      "target datalayout = \"e-m:e-i64:64-f80:128-n8:16:32:64-S128\"\n",
      "target triple = \"x86_64-unknown-linux-gnu\"\n",
      "\n",
      "%0 = type { i8*, %1, i32, %2, i64*, i64*, i64 }\n",
      "%1 = type { i32, i32 }\n",
      "%2 = type { i8, i8, i16 }\n",
      "\n",
      "@__TVMAPISetLastError = linkonce dllexport local_unnamed_addr global void (i8*)* null, align 8\n",
      "@.str = private constant [69 x i8] c\"Assert fail: (num_args == 3), default_function: num_args should be 3\\00\", align 1\n",
      "@.str.1 = private constant [144 x i8] c\"Assert fail: ((((arg0.code == 3) || (arg0.code == 13)) || (arg0.code == 7)) || (arg0.code == 4)), default_function: Expect arg[0] to be pointer\\00\", align 1\n",
      "@.str.2 = private constant [144 x i8] c\"Assert fail: ((((arg1.code == 3) || (arg1.code == 13)) || (arg1.code == 7)) || (arg1.code == 4)), default_function: Expect arg[1] to be pointer\\00\", align 1\n",
      "@.str.3 = private constant [144 x i8] c\"Assert fail: ((((arg2.code == 3) || (arg2.code == 13)) || (arg2.code == 7)) || (arg2.code == 4)), default_function: Expect arg[2] to be pointer\\00\", align 1\n",
      "@.str.4 = private constant [55 x i8] c\"Assert fail: (dev_type == 1), device_type need to be 1\\00\", align 1\n",
      "@.str.5 = private constant [81 x i8] c\"Assert fail: (1 == tvm_struct_get(arg0, 0, 4)), arg0.ndim is expected to equal 1\\00\", align 1\n",
      "@.str.6 = private constant [186 x i8] c\"Assert fail: (((tvm_struct_get(arg0, 0, 5) == (uint8)2) && (tvm_struct_get(arg0, 0, 6) == (uint8)32)) && (tvm_struct_get(arg0, 0, 7) == (uint16)1)), arg0.dtype is expected to be float32\\00\", align 1\n",
      "@.str.7 = private constant [112 x i8] c\"Assert fail: ((uint64)0 == tvm_struct_get(arg0, 0, 8)), Argument arg0.byte_offset has an unsatisfied constraint\\00\", align 1\n",
      "@.str.8 = private constant [81 x i8] c\"Assert fail: (1 == tvm_struct_get(arg1, 0, 4)), arg1.ndim is expected to equal 1\\00\", align 1\n",
      "@.str.9 = private constant [186 x i8] c\"Assert fail: (((tvm_struct_get(arg1, 0, 5) == (uint8)2) && (tvm_struct_get(arg1, 0, 6) == (uint8)32)) && (tvm_struct_get(arg1, 0, 7) == (uint16)1)), arg1.dtype is expected to be float32\\00\", align 1\n",
      "@.str.10 = private constant [112 x i8] c\"Assert fail: ((uint64)0 == tvm_struct_get(arg1, 0, 8)), Argument arg1.byte_offset has an unsatisfied constraint\\00\", align 1\n",
      "@.str.11 = private constant [105 x i8] c\"Assert fail: (1 == tvm_struct_get(arg1, 0, 10)), Argument arg1.device_type has an unsatisfied constraint\\00\", align 1\n",
      "@.str.12 = private constant [107 x i8] c\"Assert fail: (dev_id == tvm_struct_get(arg1, 0, 9)), Argument arg1.device_id has an unsatisfied constraint\\00\", align 1\n",
      "@.str.13 = private constant [81 x i8] c\"Assert fail: (1 == tvm_struct_get(arg2, 0, 4)), arg2.ndim is expected to equal 1\\00\", align 1\n",
      "@.str.14 = private constant [186 x i8] c\"Assert fail: (((tvm_struct_get(arg2, 0, 5) == (uint8)2) && (tvm_struct_get(arg2, 0, 6) == (uint8)32)) && (tvm_struct_get(arg2, 0, 7) == (uint16)1)), arg2.dtype is expected to be float32\\00\", align 1\n",
      "@.str.15 = private constant [112 x i8] c\"Assert fail: ((uint64)0 == tvm_struct_get(arg2, 0, 8)), Argument arg2.byte_offset has an unsatisfied constraint\\00\", align 1\n",
      "@.str.16 = private constant [105 x i8] c\"Assert fail: (1 == tvm_struct_get(arg2, 0, 10)), Argument arg2.device_type has an unsatisfied constraint\\00\", align 1\n",
      "@.str.17 = private constant [107 x i8] c\"Assert fail: (dev_id == tvm_struct_get(arg2, 0, 9)), Argument arg2.device_id has an unsatisfied constraint\\00\", align 1\n",
      "@__tvm_main__ = weak local_unnamed_addr constant [17 x i8] c\"default_function\\00\", align 1\n",
      "\n",
      "define dllexport i32 @default_function(i8* noalias nocapture readonly, i8* noalias nocapture readonly, i32) local_unnamed_addr {\n",
      "entry:\n",
      "  %3 = icmp eq i32 %2, 3\n",
      "  br i1 %3, label %assert_end, label %assert_fail, !prof !5\n",
      "\n",
      "assert_fail:                                      ; preds = %entry\n",
      "  %4 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %4(i8* getelementptr inbounds ([69 x i8], [69 x i8]* @.str, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end:                                       ; preds = %entry\n",
      "  %5 = bitcast i8* %0 to %0**\n",
      "  %6 = load %0*, %0** %5, align 8\n",
      "  %7 = bitcast i8* %1 to i32*\n",
      "  %8 = load i32, i32* %7, align 4, !tbaa !9\n",
      "  %9 = getelementptr inbounds i8, i8* %0, i64 8\n",
      "  %10 = bitcast i8* %9 to %0**\n",
      "  %11 = load %0*, %0** %10, align 8\n",
      "  %12 = getelementptr inbounds i8, i8* %1, i64 4\n",
      "  %13 = bitcast i8* %12 to i32*\n",
      "  %14 = load i32, i32* %13, align 4, !tbaa !23\n",
      "  %15 = getelementptr inbounds i8, i8* %0, i64 16\n",
      "  %16 = bitcast i8* %15 to %0**\n",
      "  %17 = load %0*, %0** %16, align 8\n",
      "  %18 = getelementptr inbounds i8, i8* %1, i64 8\n",
      "  %19 = bitcast i8* %18 to i32*\n",
      "  %20 = load i32, i32* %19, align 4, !tbaa !25\n",
      "  %21 = getelementptr inbounds %0, %0* %6, i64 0, i32 0\n",
      "  %22 = load i8*, i8** %21, align 8\n",
      "  %23 = getelementptr inbounds %0, %0* %6, i64 0, i32 4\n",
      "  %24 = load i64*, i64** %23, align 8\n",
      "  %25 = load i64, i64* %24, align 8, !tbaa !28\n",
      "  %26 = trunc i64 %25 to i32\n",
      "  %27 = getelementptr inbounds %0, %0* %6, i64 0, i32 5\n",
      "  %28 = load i64*, i64** %27, align 8\n",
      "  %29 = icmp eq i32 %26, 1\n",
      "  br i1 %29, label %if_end, label %if_else\n",
      "\n",
      "if_else:                                          ; preds = %assert_end\n",
      "  %30 = icmp eq i64* %28, null\n",
      "  br i1 %30, label %if_end, label %if_else2\n",
      "\n",
      "if_end:                                           ; preds = %if_else2, %if_else, %assert_end\n",
      "  %31 = phi i32 [ 0, %assert_end ], [ %46, %if_else2 ], [ 1, %if_else ]\n",
      "  %32 = getelementptr inbounds %0, %0* %6, i64 0, i32 1, i32 0\n",
      "  %33 = load i32, i32* %32, align 4\n",
      "  %34 = getelementptr inbounds %0, %0* %6, i64 0, i32 1, i32 1\n",
      "  %35 = load i32, i32* %34, align 4\n",
      "  %36 = getelementptr inbounds %0, %0* %11, i64 0, i32 0\n",
      "  %37 = load i8*, i8** %36, align 8\n",
      "  %38 = getelementptr inbounds %0, %0* %11, i64 0, i32 4\n",
      "  %39 = load i64*, i64** %38, align 8\n",
      "  %40 = load i64, i64* %39, align 8, !tbaa !42\n",
      "  %41 = trunc i64 %40 to i32\n",
      "  %42 = getelementptr inbounds %0, %0* %11, i64 0, i32 5\n",
      "  %43 = load i64*, i64** %42, align 8\n",
      "  %44 = icmp eq i32 %41, 1\n",
      "  br i1 %44, label %if_end6, label %if_else5\n",
      "\n",
      "if_else2:                                         ; preds = %if_else\n",
      "  %45 = load i64, i64* %28, align 8, !tbaa !56\n",
      "  %46 = trunc i64 %45 to i32\n",
      "  br label %if_end\n",
      "\n",
      "if_else5:                                         ; preds = %if_end\n",
      "  %47 = icmp eq i64* %43, null\n",
      "  br i1 %47, label %if_end6, label %if_else8\n",
      "\n",
      "if_end6:                                          ; preds = %if_else8, %if_else5, %if_end\n",
      "  %48 = phi i32 [ 0, %if_end ], [ %59, %if_else8 ], [ 1, %if_else5 ]\n",
      "  %49 = getelementptr inbounds %0, %0* %17, i64 0, i32 0\n",
      "  %50 = load i8*, i8** %49, align 8\n",
      "  %51 = getelementptr inbounds %0, %0* %17, i64 0, i32 4\n",
      "  %52 = load i64*, i64** %51, align 8\n",
      "  %53 = load i64, i64* %52, align 8, !tbaa !70\n",
      "  %54 = trunc i64 %53 to i32\n",
      "  %55 = getelementptr inbounds %0, %0* %17, i64 0, i32 5\n",
      "  %56 = load i64*, i64** %55, align 8\n",
      "  %57 = icmp eq i32 %54, 1\n",
      "  br i1 %57, label %if_end12, label %if_else11\n",
      "\n",
      "if_else8:                                         ; preds = %if_else5\n",
      "  %58 = load i64, i64* %43, align 8, !tbaa !84\n",
      "  %59 = trunc i64 %58 to i32\n",
      "  br label %if_end6\n",
      "\n",
      "if_else11:                                        ; preds = %if_end6\n",
      "  %60 = icmp eq i64* %56, null\n",
      "  br i1 %60, label %if_end12, label %if_else14\n",
      "\n",
      "if_end12:                                         ; preds = %if_else14, %if_else11, %if_end6\n",
      "  %61 = phi i32 [ 0, %if_end6 ], [ %63, %if_else14 ], [ 1, %if_else11 ]\n",
      "  switch i32 %8, label %assert_fail16 [\n",
      "    i32 13, label %assert_end17\n",
      "    i32 7, label %assert_end17\n",
      "    i32 4, label %assert_end17\n",
      "    i32 3, label %assert_end17\n",
      "  ]\n",
      "\n",
      "if_else14:                                        ; preds = %if_else11\n",
      "  %62 = load i64, i64* %56, align 8, !tbaa !98\n",
      "  %63 = trunc i64 %62 to i32\n",
      "  br label %if_end12\n",
      "\n",
      "assert_fail16:                                    ; preds = %if_end12\n",
      "  %64 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %64(i8* getelementptr inbounds ([144 x i8], [144 x i8]* @.str.1, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end17:                                     ; preds = %if_end12, %if_end12, %if_end12, %if_end12\n",
      "  switch i32 %14, label %assert_fail18 [\n",
      "    i32 13, label %assert_end19\n",
      "    i32 7, label %assert_end19\n",
      "    i32 4, label %assert_end19\n",
      "    i32 3, label %assert_end19\n",
      "  ]\n",
      "\n",
      "assert_fail18:                                    ; preds = %assert_end17\n",
      "  %65 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %65(i8* getelementptr inbounds ([144 x i8], [144 x i8]* @.str.2, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end19:                                     ; preds = %assert_end17, %assert_end17, %assert_end17, %assert_end17\n",
      "  switch i32 %20, label %assert_fail20 [\n",
      "    i32 13, label %assert_end21\n",
      "    i32 7, label %assert_end21\n",
      "    i32 4, label %assert_end21\n",
      "    i32 3, label %assert_end21\n",
      "  ]\n",
      "\n",
      "assert_fail20:                                    ; preds = %assert_end19\n",
      "  %66 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %66(i8* getelementptr inbounds ([144 x i8], [144 x i8]* @.str.3, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end21:                                     ; preds = %assert_end19, %assert_end19, %assert_end19, %assert_end19\n",
      "  %67 = icmp eq i32 %33, 1\n",
      "  br i1 %67, label %assert_end23, label %assert_fail22, !prof !5\n",
      "\n",
      "assert_fail22:                                    ; preds = %assert_end21\n",
      "  %68 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %68(i8* getelementptr inbounds ([55 x i8], [55 x i8]* @.str.4, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end23:                                     ; preds = %assert_end21\n",
      "  %69 = getelementptr inbounds %0, %0* %6, i64 0, i32 2\n",
      "  %70 = load i32, i32* %69, align 4\n",
      "  %71 = icmp eq i32 %70, 1\n",
      "  br i1 %71, label %assert_end25, label %assert_fail24, !prof !5\n",
      "\n",
      "assert_fail24:                                    ; preds = %assert_end23\n",
      "  %72 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %72(i8* getelementptr inbounds ([81 x i8], [81 x i8]* @.str.5, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end25:                                     ; preds = %assert_end23\n",
      "  %73 = getelementptr inbounds %0, %0* %6, i64 0, i32 3, i32 2\n",
      "  %74 = load i16, i16* %73, align 2\n",
      "  %75 = icmp eq i16 %74, 1\n",
      "  %76 = getelementptr inbounds %0, %0* %6, i64 0, i32 3, i32 1\n",
      "  %77 = load i8, i8* %76, align 1\n",
      "  %78 = icmp eq i8 %77, 32\n",
      "  %79 = getelementptr inbounds %0, %0* %6, i64 0, i32 3, i32 0\n",
      "  %80 = load i8, i8* %79, align 1\n",
      "  %81 = icmp eq i8 %80, 2\n",
      "  %82 = and i1 %78, %81\n",
      "  %83 = and i1 %75, %82\n",
      "  br i1 %83, label %assert_end27, label %assert_fail26, !prof !5\n",
      "\n",
      "assert_fail26:                                    ; preds = %assert_end25\n",
      "  %84 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %84(i8* getelementptr inbounds ([186 x i8], [186 x i8]* @.str.6, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end27:                                     ; preds = %assert_end25\n",
      "  %85 = getelementptr inbounds %0, %0* %6, i64 0, i32 6\n",
      "  %86 = load i64, i64* %85, align 8\n",
      "  %87 = icmp eq i64 %86, 0\n",
      "  br i1 %87, label %assert_end29, label %assert_fail28, !prof !5\n",
      "\n",
      "assert_fail28:                                    ; preds = %assert_end27\n",
      "  %88 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %88(i8* getelementptr inbounds ([112 x i8], [112 x i8]* @.str.7, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end29:                                     ; preds = %assert_end27\n",
      "  %89 = getelementptr inbounds %0, %0* %11, i64 0, i32 2\n",
      "  %90 = load i32, i32* %89, align 4\n",
      "  %91 = icmp eq i32 %90, 1\n",
      "  br i1 %91, label %assert_end31, label %assert_fail30, !prof !5\n",
      "\n",
      "assert_fail30:                                    ; preds = %assert_end29\n",
      "  %92 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %92(i8* getelementptr inbounds ([81 x i8], [81 x i8]* @.str.8, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end31:                                     ; preds = %assert_end29\n",
      "  %93 = getelementptr inbounds %0, %0* %11, i64 0, i32 3, i32 2\n",
      "  %94 = load i16, i16* %93, align 2\n",
      "  %95 = icmp eq i16 %94, 1\n",
      "  %96 = getelementptr inbounds %0, %0* %11, i64 0, i32 3, i32 1\n",
      "  %97 = load i8, i8* %96, align 1\n",
      "  %98 = icmp eq i8 %97, 32\n",
      "  %99 = getelementptr inbounds %0, %0* %11, i64 0, i32 3, i32 0\n",
      "  %100 = load i8, i8* %99, align 1\n",
      "  %101 = icmp eq i8 %100, 2\n",
      "  %102 = and i1 %98, %101\n",
      "  %103 = and i1 %95, %102\n",
      "  br i1 %103, label %assert_end33, label %assert_fail32, !prof !5\n",
      "\n",
      "assert_fail32:                                    ; preds = %assert_end31\n",
      "  %104 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %104(i8* getelementptr inbounds ([186 x i8], [186 x i8]* @.str.9, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end33:                                     ; preds = %assert_end31\n",
      "  %105 = getelementptr inbounds %0, %0* %11, i64 0, i32 6\n",
      "  %106 = load i64, i64* %105, align 8\n",
      "  %107 = icmp eq i64 %106, 0\n",
      "  br i1 %107, label %assert_end35, label %assert_fail34, !prof !5\n",
      "\n",
      "assert_fail34:                                    ; preds = %assert_end33\n",
      "  %108 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %108(i8* getelementptr inbounds ([112 x i8], [112 x i8]* @.str.10, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end35:                                     ; preds = %assert_end33\n",
      "  %109 = getelementptr inbounds %0, %0* %11, i64 0, i32 1, i32 0\n",
      "  %110 = load i32, i32* %109, align 4\n",
      "  %111 = icmp eq i32 %110, 1\n",
      "  br i1 %111, label %assert_end37, label %assert_fail36, !prof !5\n",
      "\n",
      "assert_fail36:                                    ; preds = %assert_end35\n",
      "  %112 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %112(i8* getelementptr inbounds ([105 x i8], [105 x i8]* @.str.11, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end37:                                     ; preds = %assert_end35\n",
      "  %113 = getelementptr inbounds %0, %0* %11, i64 0, i32 1, i32 1\n",
      "  %114 = load i32, i32* %113, align 4\n",
      "  %115 = icmp eq i32 %35, %114\n",
      "  br i1 %115, label %assert_end39, label %assert_fail38, !prof !5\n",
      "\n",
      "assert_fail38:                                    ; preds = %assert_end37\n",
      "  %116 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %116(i8* getelementptr inbounds ([107 x i8], [107 x i8]* @.str.12, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end39:                                     ; preds = %assert_end37\n",
      "  %117 = getelementptr inbounds %0, %0* %17, i64 0, i32 2\n",
      "  %118 = load i32, i32* %117, align 4\n",
      "  %119 = icmp eq i32 %118, 1\n",
      "  br i1 %119, label %assert_end41, label %assert_fail40, !prof !5\n",
      "\n",
      "assert_fail40:                                    ; preds = %assert_end39\n",
      "  %120 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %120(i8* getelementptr inbounds ([81 x i8], [81 x i8]* @.str.13, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end41:                                     ; preds = %assert_end39\n",
      "  %121 = getelementptr inbounds %0, %0* %17, i64 0, i32 3, i32 2\n",
      "  %122 = load i16, i16* %121, align 2\n",
      "  %123 = icmp eq i16 %122, 1\n",
      "  %124 = getelementptr inbounds %0, %0* %17, i64 0, i32 3, i32 1\n",
      "  %125 = load i8, i8* %124, align 1\n",
      "  %126 = icmp eq i8 %125, 32\n",
      "  %127 = getelementptr inbounds %0, %0* %17, i64 0, i32 3, i32 0\n",
      "  %128 = load i8, i8* %127, align 1\n",
      "  %129 = icmp eq i8 %128, 2\n",
      "  %130 = and i1 %126, %129\n",
      "  %131 = and i1 %123, %130\n",
      "  br i1 %131, label %assert_end43, label %assert_fail42, !prof !5\n",
      "\n",
      "assert_fail42:                                    ; preds = %assert_end41\n",
      "  %132 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %132(i8* getelementptr inbounds ([186 x i8], [186 x i8]* @.str.14, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end43:                                     ; preds = %assert_end41\n",
      "  %133 = getelementptr inbounds %0, %0* %17, i64 0, i32 6\n",
      "  %134 = load i64, i64* %133, align 8\n",
      "  %135 = icmp eq i64 %134, 0\n",
      "  br i1 %135, label %assert_end45, label %assert_fail44, !prof !5\n",
      "\n",
      "assert_fail44:                                    ; preds = %assert_end43\n",
      "  %136 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %136(i8* getelementptr inbounds ([112 x i8], [112 x i8]* @.str.15, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end45:                                     ; preds = %assert_end43\n",
      "  %137 = getelementptr inbounds %0, %0* %17, i64 0, i32 1, i32 0\n",
      "  %138 = load i32, i32* %137, align 4\n",
      "  %139 = icmp eq i32 %138, 1\n",
      "  br i1 %139, label %assert_end47, label %assert_fail46, !prof !5\n",
      "\n",
      "assert_fail46:                                    ; preds = %assert_end45\n",
      "  %140 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %140(i8* getelementptr inbounds ([105 x i8], [105 x i8]* @.str.16, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end47:                                     ; preds = %assert_end45\n",
      "  %141 = getelementptr inbounds %0, %0* %17, i64 0, i32 1, i32 1\n",
      "  %142 = load i32, i32* %141, align 4\n",
      "  %143 = icmp eq i32 %35, %142\n",
      "  br i1 %143, label %assert_end49, label %assert_fail48, !prof !5\n",
      "\n",
      "assert_fail48:                                    ; preds = %assert_end47\n",
      "  %144 = load void (i8*)*, void (i8*)** @__TVMAPISetLastError, align 8, !tbaa !6\n",
      "  tail call void %144(i8* getelementptr inbounds ([107 x i8], [107 x i8]* @.str.17, i64 0, i64 0))\n",
      "  ret i32 -1\n",
      "\n",
      "assert_end49:                                     ; preds = %assert_end47\n",
      "  tail call fastcc void @default_function_compute_(i32 %54, i8* %50, i8* %22, i32 %31, i8* %37, i32 %48, i32 %61)\n",
      "  ret i32 0\n",
      "}\n",
      "\n",
      "; Function Attrs: noinline norecurse nounwind\n",
      "define private fastcc void @default_function_compute_(i32, i8* noalias nocapture, i8* noalias nocapture readonly, i32, i8* noalias nocapture readonly, i32, i32) unnamed_addr #0 {\n",
      "entry:\n",
      "  %7 = icmp sgt i32 %0, 0\n",
      "  br i1 %7, label %for_body.lr.ph, label %for_end, !prof !5\n",
      "\n",
      "for_body.lr.ph:                                   ; preds = %entry\n",
      "  %8 = bitcast i8* %4 to float*\n",
      "  %9 = bitcast i8* %2 to float*\n",
      "  %10 = bitcast i8* %1 to float*\n",
      "  %11 = sext i32 %3 to i64\n",
      "  %12 = sext i32 %5 to i64\n",
      "  %13 = sext i32 %6 to i64\n",
      "  %wide.trip.count = zext i32 %0 to i64\n",
      "  %min.iters.check = icmp ult i32 %0, 4\n",
      "  br i1 %min.iters.check, label %for_body.preheader, label %vector.scevcheck\n",
      "\n",
      "vector.scevcheck:                                 ; preds = %for_body.lr.ph\n",
      "  %ident.check = icmp ne i32 %5, 1\n",
      "  %ident.check1 = icmp ne i32 %3, 1\n",
      "  %14 = or i1 %ident.check, %ident.check1\n",
      "  %ident.check2 = icmp ne i32 %6, 1\n",
      "  %15 = or i1 %14, %ident.check2\n",
      "  br i1 %15, label %for_body.preheader, label %vector.ph\n",
      "\n",
      "vector.ph:                                        ; preds = %vector.scevcheck\n",
      "  %n.vec = and i64 %wide.trip.count, 4294967292\n",
      "  br label %vector.body\n",
      "\n",
      "vector.body:                                      ; preds = %vector.body, %vector.ph\n",
      "  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]\n",
      "  %16 = mul nsw i64 %index, %13\n",
      "  %17 = mul nsw i64 %index, %12\n",
      "  %18 = getelementptr inbounds float, float* %8, i64 %17\n",
      "  %19 = bitcast float* %18 to <4 x float>*\n",
      "  %wide.load = load <4 x float>, <4 x float>* %19, align 4, !tbaa !112\n",
      "  %20 = mul nsw i64 %index, %11\n",
      "  %21 = getelementptr inbounds float, float* %9, i64 %20\n",
      "  %22 = bitcast float* %21 to <4 x float>*\n",
      "  %wide.load3 = load <4 x float>, <4 x float>* %22, align 4, !tbaa !115\n",
      "  %23 = fadd <4 x float> %wide.load, %wide.load3\n",
      "  %24 = getelementptr inbounds float, float* %10, i64 %16\n",
      "  %25 = bitcast float* %24 to <4 x float>*\n",
      "  store <4 x float> %23, <4 x float>* %25, align 4, !tbaa !118\n",
      "  %index.next = add i64 %index, 4\n",
      "  %26 = icmp eq i64 %index.next, %n.vec\n",
      "  br i1 %26, label %middle.block, label %vector.body, !llvm.loop !121\n",
      "\n",
      "middle.block:                                     ; preds = %vector.body\n",
      "  %cmp.n = icmp eq i64 %n.vec, %wide.trip.count\n",
      "  br i1 %cmp.n, label %for_end, label %for_body.preheader\n",
      "\n",
      "for_body.preheader:                               ; preds = %middle.block, %vector.scevcheck, %for_body.lr.ph\n",
      "  %indvars.iv.ph = phi i64 [ 0, %vector.scevcheck ], [ 0, %for_body.lr.ph ], [ %n.vec, %middle.block ]\n",
      "  br label %for_body\n",
      "\n",
      "for_body:                                         ; preds = %for_body.preheader, %for_body\n",
      "  %indvars.iv = phi i64 [ %indvars.iv.next, %for_body ], [ %indvars.iv.ph, %for_body.preheader ]\n",
      "  %27 = mul nsw i64 %indvars.iv, %13\n",
      "  %28 = mul nsw i64 %indvars.iv, %12\n",
      "  %29 = getelementptr inbounds float, float* %8, i64 %28\n",
      "  %30 = load float, float* %29, align 4, !tbaa !112\n",
      "  %31 = mul nsw i64 %indvars.iv, %11\n",
      "  %32 = getelementptr inbounds float, float* %9, i64 %31\n",
      "  %33 = load float, float* %32, align 4, !tbaa !115\n",
      "  %34 = fadd float %30, %33\n",
      "  %35 = getelementptr inbounds float, float* %10, i64 %27\n",
      "  store float %34, float* %35, align 4, !tbaa !118\n",
      "  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1\n",
      "  %exitcond = icmp eq i64 %indvars.iv.next, %wide.trip.count\n",
      "  br i1 %exitcond, label %for_end, label %for_body, !prof !123, !llvm.loop !124\n",
      "\n",
      "for_end:                                          ; preds = %for_body, %middle.block, %entry\n",
      "  ret void\n",
      "}\n",
      "\n",
      "attributes #0 = { noinline norecurse nounwind }\n",
      "\n",
      "!llvm.dbg.cu = !{!0}\n",
      "!llvm.module.flags = !{!3, !4}\n",
      "\n",
      "!0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: \"TVM\", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, dwoId: 1)\n",
      "!1 = !DIFile(filename: \"model.tvm\", directory: \"/tmp/\")\n",
      "!2 = !{}\n",
      "!3 = !{i32 2, !\"tvm_target\", !\"llvm\"}\n",
      "!4 = !{i32 4, !\"Debug Info Version\", i32 3}\n",
      "!5 = !{!\"branch_weights\", i32 1048576, i32 1}\n",
      "!6 = !{!7, !7, i64 0}\n",
      "!7 = !{!\"ctx_ptr\", !8, i64 0}\n",
      "!8 = !{!\"tvm-tbaa\"}\n",
      "!9 = !{!10, !10, i64 0}\n",
      "!10 = !{!\"0x55ad6af78d70.w1.b0\", !11, i64 0}\n",
      "!11 = !{!\"0x55ad6af78d70.w2.b0\", !12, i64 0}\n",
      "!12 = !{!\"0x55ad6af78d70.w4.b0\", !13, i64 0}\n",
      "!13 = !{!\"0x55ad6af78d70.w8.b0\", !14, i64 0}\n",
      "!14 = !{!\"0x55ad6af78d70.w16.b0\", !15, i64 0}\n",
      "!15 = !{!\"0x55ad6af78d70.w32.b0\", !16, i64 0}\n",
      "!16 = !{!\"0x55ad6af78d70.w64.b0\", !17, i64 0}\n",
      "!17 = !{!\"0x55ad6af78d70.w128.b0\", !18, i64 0}\n",
      "!18 = !{!\"0x55ad6af78d70.w256.b0\", !19, i64 0}\n",
      "!19 = !{!\"0x55ad6af78d70.w512.b0\", !20, i64 0}\n",
      "!20 = !{!\"0x55ad6af78d70.w1024.b0\", !21, i64 0}\n",
      "!21 = !{!\"int32\", !22, i64 0}\n",
      "!22 = !{!\"0x55ad6af78d70\", !8, i64 0}\n",
      "!23 = !{!24, !24, i64 0}\n",
      "!24 = !{!\"0x55ad6af78d70.w1.b1\", !11, i64 0}\n",
      "!25 = !{!26, !26, i64 0}\n",
      "!26 = !{!\"0x55ad6af78d70.w1.b2\", !27, i64 0}\n",
      "!27 = !{!\"0x55ad6af78d70.w2.b2\", !12, i64 0}\n",
      "!28 = !{!29, !29, i64 0}\n",
      "!29 = !{!\"0x55ad6b3416a0.w1.b0\", !30, i64 0}\n",
      "!30 = !{!\"0x55ad6b3416a0.w2.b0\", !31, i64 0}\n",
      "!31 = !{!\"0x55ad6b3416a0.w4.b0\", !32, i64 0}\n",
      "!32 = !{!\"0x55ad6b3416a0.w8.b0\", !33, i64 0}\n",
      "!33 = !{!\"0x55ad6b3416a0.w16.b0\", !34, i64 0}\n",
      "!34 = !{!\"0x55ad6b3416a0.w32.b0\", !35, i64 0}\n",
      "!35 = !{!\"0x55ad6b3416a0.w64.b0\", !36, i64 0}\n",
      "!36 = !{!\"0x55ad6b3416a0.w128.b0\", !37, i64 0}\n",
      "!37 = !{!\"0x55ad6b3416a0.w256.b0\", !38, i64 0}\n",
      "!38 = !{!\"0x55ad6b3416a0.w512.b0\", !39, i64 0}\n",
      "!39 = !{!\"0x55ad6b3416a0.w1024.b0\", !40, i64 0}\n",
      "!40 = !{!\"int64\", !41, i64 0}\n",
      "!41 = !{!\"0x55ad6b3416a0\", !8, i64 0}\n",
      "!42 = !{!43, !43, i64 0}\n",
      "!43 = !{!\"0x55ad6b3444c0.w1.b0\", !44, i64 0}\n",
      "!44 = !{!\"0x55ad6b3444c0.w2.b0\", !45, i64 0}\n",
      "!45 = !{!\"0x55ad6b3444c0.w4.b0\", !46, i64 0}\n",
      "!46 = !{!\"0x55ad6b3444c0.w8.b0\", !47, i64 0}\n",
      "!47 = !{!\"0x55ad6b3444c0.w16.b0\", !48, i64 0}\n",
      "!48 = !{!\"0x55ad6b3444c0.w32.b0\", !49, i64 0}\n",
      "!49 = !{!\"0x55ad6b3444c0.w64.b0\", !50, i64 0}\n",
      "!50 = !{!\"0x55ad6b3444c0.w128.b0\", !51, i64 0}\n",
      "!51 = !{!\"0x55ad6b3444c0.w256.b0\", !52, i64 0}\n",
      "!52 = !{!\"0x55ad6b3444c0.w512.b0\", !53, i64 0}\n",
      "!53 = !{!\"0x55ad6b3444c0.w1024.b0\", !54, i64 0}\n",
      "!54 = !{!\"int64\", !55, i64 0}\n",
      "!55 = !{!\"0x55ad6b3444c0\", !8, i64 0}\n",
      "!56 = !{!57, !57, i64 0}\n",
      "!57 = !{!\"0x55ad6b342bc0.w1.b0\", !58, i64 0}\n",
      "!58 = !{!\"0x55ad6b342bc0.w2.b0\", !59, i64 0}\n",
      "!59 = !{!\"0x55ad6b342bc0.w4.b0\", !60, i64 0}\n",
      "!60 = !{!\"0x55ad6b342bc0.w8.b0\", !61, i64 0}\n",
      "!61 = !{!\"0x55ad6b342bc0.w16.b0\", !62, i64 0}\n",
      "!62 = !{!\"0x55ad6b342bc0.w32.b0\", !63, i64 0}\n",
      "!63 = !{!\"0x55ad6b342bc0.w64.b0\", !64, i64 0}\n",
      "!64 = !{!\"0x55ad6b342bc0.w128.b0\", !65, i64 0}\n",
      "!65 = !{!\"0x55ad6b342bc0.w256.b0\", !66, i64 0}\n",
      "!66 = !{!\"0x55ad6b342bc0.w512.b0\", !67, i64 0}\n",
      "!67 = !{!\"0x55ad6b342bc0.w1024.b0\", !68, i64 0}\n",
      "!68 = !{!\"int64\", !69, i64 0}\n",
      "!69 = !{!\"0x55ad6b342bc0\", !8, i64 0}\n",
      "!70 = !{!71, !71, i64 0}\n",
      "!71 = !{!\"0x55ad6b347210.w1.b0\", !72, i64 0}\n",
      "!72 = !{!\"0x55ad6b347210.w2.b0\", !73, i64 0}\n",
      "!73 = !{!\"0x55ad6b347210.w4.b0\", !74, i64 0}\n",
      "!74 = !{!\"0x55ad6b347210.w8.b0\", !75, i64 0}\n",
      "!75 = !{!\"0x55ad6b347210.w16.b0\", !76, i64 0}\n",
      "!76 = !{!\"0x55ad6b347210.w32.b0\", !77, i64 0}\n",
      "!77 = !{!\"0x55ad6b347210.w64.b0\", !78, i64 0}\n",
      "!78 = !{!\"0x55ad6b347210.w128.b0\", !79, i64 0}\n",
      "!79 = !{!\"0x55ad6b347210.w256.b0\", !80, i64 0}\n",
      "!80 = !{!\"0x55ad6b347210.w512.b0\", !81, i64 0}\n",
      "!81 = !{!\"0x55ad6b347210.w1024.b0\", !82, i64 0}\n",
      "!82 = !{!\"int64\", !83, i64 0}\n",
      "!83 = !{!\"0x55ad6b347210\", !8, i64 0}\n",
      "!84 = !{!85, !85, i64 0}\n",
      "!85 = !{!\"0x55ad6b344910.w1.b0\", !86, i64 0}\n",
      "!86 = !{!\"0x55ad6b344910.w2.b0\", !87, i64 0}\n",
      "!87 = !{!\"0x55ad6b344910.w4.b0\", !88, i64 0}\n",
      "!88 = !{!\"0x55ad6b344910.w8.b0\", !89, i64 0}\n",
      "!89 = !{!\"0x55ad6b344910.w16.b0\", !90, i64 0}\n",
      "!90 = !{!\"0x55ad6b344910.w32.b0\", !91, i64 0}\n",
      "!91 = !{!\"0x55ad6b344910.w64.b0\", !92, i64 0}\n",
      "!92 = !{!\"0x55ad6b344910.w128.b0\", !93, i64 0}\n",
      "!93 = !{!\"0x55ad6b344910.w256.b0\", !94, i64 0}\n",
      "!94 = !{!\"0x55ad6b344910.w512.b0\", !95, i64 0}\n",
      "!95 = !{!\"0x55ad6b344910.w1024.b0\", !96, i64 0}\n",
      "!96 = !{!\"int64\", !97, i64 0}\n",
      "!97 = !{!\"0x55ad6b344910\", !8, i64 0}\n",
      "!98 = !{!99, !99, i64 0}\n",
      "!99 = !{!\"0x55ad6b3476a0.w1.b0\", !100, i64 0}\n",
      "!100 = !{!\"0x55ad6b3476a0.w2.b0\", !101, i64 0}\n",
      "!101 = !{!\"0x55ad6b3476a0.w4.b0\", !102, i64 0}\n",
      "!102 = !{!\"0x55ad6b3476a0.w8.b0\", !103, i64 0}\n",
      "!103 = !{!\"0x55ad6b3476a0.w16.b0\", !104, i64 0}\n",
      "!104 = !{!\"0x55ad6b3476a0.w32.b0\", !105, i64 0}\n",
      "!105 = !{!\"0x55ad6b3476a0.w64.b0\", !106, i64 0}\n",
      "!106 = !{!\"0x55ad6b3476a0.w128.b0\", !107, i64 0}\n",
      "!107 = !{!\"0x55ad6b3476a0.w256.b0\", !108, i64 0}\n",
      "!108 = !{!\"0x55ad6b3476a0.w512.b0\", !109, i64 0}\n",
      "!109 = !{!\"0x55ad6b3476a0.w1024.b0\", !110, i64 0}\n",
      "!110 = !{!\"int64\", !111, i64 0}\n",
      "!111 = !{!\"0x55ad6b3476a0\", !8, i64 0}\n",
      "!112 = !{!113, !113, i64 0}\n",
      "!113 = !{!\"float32\", !114, i64 0}\n",
      "!114 = !{!\"0x55ad6af78db0\", !8, i64 0}\n",
      "!115 = !{!116, !116, i64 0}\n",
      "!116 = !{!\"float32\", !117, i64 0}\n",
      "!117 = !{!\"0x55ad6b324b80\", !8, i64 0}\n",
      "!118 = !{!119, !119, i64 0}\n",
      "!119 = !{!\"float32\", !120, i64 0}\n",
      "!120 = !{!\"0x55ad6b3211e0\", !8, i64 0}\n",
      "!121 = distinct !{!121, !122}\n",
      "!122 = !{!\"llvm.loop.isvectorized\", i32 1}\n",
      "!123 = !{!\"branch_weights\", i32 1, i32 1048576}\n",
      "!124 = distinct !{!124, !122}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pseudocode = tvm.lower(s, [A, B, C], simple_mode=True)\n",
    "mod_cpu = tvm.build(s, [A, B, C], target='llvm')\n",
    "print(pseudocode)\n",
    "print('#' * 80)\n",
    "print(mod_cpu.get_source())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For we use contiguous tensor, without broadcast, the stride can be computed by reverse exclusive scan using multiplication. For example, the stride of a ndarray with shape **(2, 3, 4)** is **(12, 4, 1)**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tvm.NDArray shape=(3,), cpu(0)>\n",
       "array([3., 3., 3.], dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = tvm.nd.array(np.ones((3,), dtype=\"float32\"))\n",
    "b = tvm.nd.array(np.ones((1,), dtype=\"float32\") * 2)\n",
    "c = tvm.nd.array(np.zeros((3,), dtype=\"float32\"))\n",
    "mod_cpu(a, b, c)\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "produce C {\n",
      "  // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = floordiv((tindex + 63), 64)\n",
      "  // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 64\n",
      "  if (likely((((blockIdx.x*64) + threadIdx.x) < tindex))) {\n",
      "    if (likely((((blockIdx.x*64) + threadIdx.x) < tindex))) {\n",
      "      C[(((blockIdx.x*64) + threadIdx.x)*stride)] = (A[(((blockIdx.x*64) + threadIdx.x)*stride)] + B[(((blockIdx.x*64) + threadIdx.x)*stride)])\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n",
      "################################################################################\n",
      "extern \"C\" __global__ void default_function_kernel0( float* __restrict__ C,  float* __restrict__ A,  float* __restrict__ B, int tindex, int stride, int stride1, int stride2) {\n",
      "  if (((int)blockIdx.x) < (tindex >> 6)) {\n",
      "    C[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride2)] = (A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] + B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);\n",
      "  } else {\n",
      "    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < tindex) {\n",
      "      if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < tindex) {\n",
      "        C[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride2)] = (A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] + B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "bx, tx = s[C].split(C.op.axis[0], factor=64)\n",
    "s[C].bind(bx, tvm.thread_axis(\"blockIdx.x\"))\n",
    "s[C].bind(tx, tvm.thread_axis(\"threadIdx.x\"))\n",
    "\n",
    "pseudocode = tvm.lower(s, [A, B, C], simple_mode=True)\n",
    "mod_gpu = tvm.build(s, [A, B, C], target='cuda')\n",
    "print(pseudocode)\n",
    "print('#' * 80)\n",
    "print(mod_gpu.imported_modules[0].get_source())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tvm.NDArray shape=(2,), gpu(0)>\n",
       "array([3., 3.], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ctx = tvm.context('cuda', 0)\n",
    "a = tvm.nd.array(np.ones((2,), dtype=\"float32\"), ctx=ctx)\n",
    "b = tvm.nd.array(np.ones((2,), dtype=\"float32\") * 2, ctx=ctx)\n",
    "c = tvm.nd.array(np.zeros((2,), dtype=\"float32\"), ctx=ctx)\n",
    "mod_gpu(a, b, c)\n",
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.4 Advanced Usage\n",
    "\n",
    "Do not support native python control flow, use `tvm.if_then_else` instead. Or try to use [hybrid script](https://docs.tvm.ai/langref/hybrid_script.html).\n",
    "\n",
    "How about broadcast?\n",
    "\n",
    "- https://docs.tvm.ai/api/python/tvm.html?highlight=auto_broadcast#tvm.decl_buffer\n",
    "\n",
    "How to optimize the schedule of an operator? Refer to:\n",
    "\n",
    "- https://docs.tvm.ai/tutorials/index.html#tensor-expression-and-schedules\n",
    "- https://docs.tvm.ai/tutorials/language/schedule_primitives.html\n",
    "\n",
    "Reshape, slicing, etc. See:\n",
    "\n",
    "- http://tvm.d2l.ai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Use TVM to implement operator in MXNet backend\n",
    "\n",
    "Also use `add` as example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use python implement core compute function\n",
    "\n",
    "**Core compute funtion, like kernel(s) in `.h` files:**\n",
    "\n",
    "- Inputs: arguments that need traversion is necessary\n",
    "- Outouts: schedule, input(s), output(s)\n",
    "\n",
    "```python\n",
    "# contrib/tvmop/basic/ufunc.py\n",
    "def compute_add(atype, btype, ndim):\n",
    "    A = tvm.placeholder([tvm.var() for _ in range(ndim)], name='A', dtype=atype)\n",
    "    B = tvm.placeholder([tvm.var() for _ in range(ndim)], name='B', dtype=btype)\n",
    "    C = tvm.compute([tvm.var() for _ in range(ndim)],\n",
    "                    lambda *index: A[index].astype() + B[index].astype(), name='C')\n",
    "    s = tvm.create_schedule(C.op)\n",
    "    return s, A, B, C\n",
    "```\n",
    "\n",
    "**Device related compute function, similar to `FCompute`:**\n",
    "\n",
    "- Function signature similar to above\n",
    "- In decorator: `name`, `target`, `dtype` and `ndim` are necessary\n",
    "\n",
    "```python\n",
    "# contrib/tvmop/basic/ufunc.py\n",
    "@defop(name=\"vadd\", target=\"cpu\", auto_broadcast=True,\n",
    "       atypes=AllTypes, btypes=FloatTypes, ndim=[5])\n",
    "def vadd(atype, btype, ndim):\n",
    "    s, A, B, C = compute_add(dtype, ndim, '-')\n",
    "    axes = [axis for axis in C.op.axis]\n",
    "    fused = s[C].fuse(*axes)\n",
    "    s[C].parallel(fused)\n",
    "\n",
    "    return s, [A, B, C]\n",
    "\n",
    "\n",
    "@defop(name=\"cuda_vadd\", target=\"cuda\", auto_broadcast=True,\n",
    "       dtype=[\"float32\", \"float64\"], ndim=[5])\n",
    "def vadd_gpu(dtype, ndim):\n",
    "    s, A, B, C = compute_add(dtype, ndim)\n",
    "    s = tvm.create_schedule(C.op)\n",
    "    axes = [axis for axis in C.op.axis]\n",
    "    fused = s[C].fuse(*axes)\n",
    "    bx, tx = s[C].split(fused, factor=64)\n",
    "    s[C].bind(bx, tvm.thread_axis(\"blockIdx.x\"))\n",
    "    s[C].bind(tx, tvm.thread_axis(\"threadIdx.x\"))\n",
    "    return s, [A, B, C]\n",
    "```\n",
    "\n",
    "**Backward is similar:**\n",
    "\n",
    "```python\n",
    "# contrib/tvmop/basic/ufunc.py\n",
    "def compute_backward_vadd(dtype, ndim, reduce1st, req):\n",
    "    axes = ([reduce1st, 1 - reduce1st] * ndim)[:ndim]\n",
    "    X = tvm.placeholder([tvm.var() for _ in range(ndim)], name='X', dtype=dtype)\n",
    "    reducer = tvm.comm_reducer(lambda x, y: x + y,\n",
    "        lambda t: tvm.const(0, dtype=t), name=\"sum\")\n",
    "    ret = reduce_axes(X, axes, reducer)\n",
    "    in_grad_a, in_grad = assign_by_req(ret, req)\n",
    "    s = tvm.create_schedule(in_grad.op)\n",
    "    return s, X, in_grad_a, in_grad, [ret, in_grad]\n",
    "\n",
    "\n",
    "@defop(name=\"backward_vadd\", target=\"cpu\", dtype=AllTypes,\n",
    "       ndim=[5], reduce1st=[0, 1],\n",
    "       req=[\"kWriteTo\", \"kAddTo\"], attrs=[\"reduce1st\", \"req\"])\n",
    "def backward_vadd(dtype, ndim, reduce1st, req):\n",
    "    s, X, in_grad_a, in_grad, c_list = compute_backward_vadd(dtype, ndim, reduce1st, req)\n",
    "    for t in c_list:\n",
    "        axes = [axis for axis in t.op.axis]\n",
    "        fused = s[t].fuse(*axes)\n",
    "        s[t].parallel(fused)\n",
    "    return s, [X, in_grad_a, in_grad]\n",
    "\n",
    "\n",
    "@defop(name=\"cuda_backward_vadd\", target=\"gpu\", dtype=[\"float32\", \"float64\"],\n",
    "       ndim=[5], reduce1st=[0, 1],\n",
    "       req=[\"kWriteTo\", \"kAddTo\"], attrs=[\"reduce1st\", \"req\"])\n",
    "def backward_vadd_gpu(dtype, ndim, reduce1st, req):\n",
    "    s, X, in_grad_a, in_grad, c_list = compute_backward_vadd(dtype, ndim, reduce1st, req)\n",
    "    num_thread = 64\n",
    "    for t in c_list:\n",
    "        block_x = tvm.thread_axis(\"blockIdx.x\")\n",
    "        thread_x = tvm.thread_axis(\"threadIdx.x\")\n",
    "        axes = [axis for axis in t.op.axis]\n",
    "        fused = s[t].fuse(*axes)\n",
    "        bx, tx = s[t].split(fused, factor=num_thread)\n",
    "        s[t].bind(bx, block_x)\n",
    "        s[t].bind(tx, thread_x)\n",
    "    return s, [X, in_grad_a, in_grad]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register operator in backend\n",
    "\n",
    "```c++\n",
    "// src/operator/contrib/tvmop/ufunc.cc\n",
    "static constexpr char func_vadd_cpu[] = \"vadd\";\n",
    "static constexpr char func_vadd_gpu[] = \"cuda_vadd\";\n",
    "static constexpr char func_bakcward_vadd_cpu[] = \"backward_vadd\";\n",
    "static constexpr char func_bakcward_vadd_gpu[] = \"cuda_backward_vadd\";\n",
    "\n",
    "template<const char* func>\n",
    "void TVMBinaryCompute(const nnvm::NodeAttrs& attrs,\n",
    "                      const mxnet::OpContext& ctx,\n",
    "                      const std::vector<TBlob>& inputs,\n",
    "                      const std::vector<OpReqType>& req,\n",
    "                      const std::vector<TBlob>& outputs) {\n",
    "  CHECK_EQ(inputs.size(), 2U);\n",
    "  CHECK_EQ(outputs.size(), 1U);\n",
    "  TBlob idata[2], odata;\n",
    "  for (int k = 0; k < 2; ++k) {\n",
    "    idata[k] = Padding(inputs[k], max_dim);\n",
    "  }\n",
    "  odata = Padding(outputs[0], max_dim);\n",
    "  tvm::runtime::TVMOpModule::Get()->Call(func, ctx, {idata[0], idata[1], odata});\n",
    "}\n",
    "\n",
    "NNVM_REGISTER_OP(_contrib_tvm_vadd)\n",
    "    .set_num_inputs(2)\n",
    "    .set_num_outputs(1)\n",
    "    .add_argument(\"a\", \"NDArray-or-Symbol\", \"first input\")\n",
    "    .add_argument(\"b\", \"NDArray-or-Symbol\", \"second input\")\n",
    "    .set_attr<nnvm::FListInputNames>(\"FListInputNames\",\n",
    "      [](const NodeAttrs& attrs) {\n",
    "        return std::vector<std::string>{\"a\", \"b\"};\n",
    "      })\n",
    "    .set_attr<mxnet::FInferShape>(\"FInferShape\", BinaryBroadcastShape)\n",
    "    .set_attr<nnvm::FInferType>(\"FInferType\", mxnet::op::ElemwiseType<2, 1>)\n",
    "#if MXNET_USE_CUDA\n",
    "    .set_attr<mxnet::FCompute>(\"FCompute<gpu>\", mxnet::op::TVMBinaryCompute<func_vadd_gpu>)\n",
    "#endif  // MXNET_USE_CUDA\n",
    "    .set_attr<mxnet::FCompute>(\"FCompute<cpu>\", mxnet::op::TVMBinaryCompute<func_vadd_cpu>)\n",
    "    .set_attr<nnvm::FGradient>(\"FGradient\", ElemwiseGradUseNone{\"_backward_contrib_tvm_vadd\"});\n",
    "\n",
    "NNVM_REGISTER_OP(_backward_contrib_tvm_vadd)\n",
    "    .set_num_inputs(1)\n",
    "    .set_num_outputs(2)\n",
    "    .set_attr<nnvm::TIsBackward>(\"TIsBackward\", true)\n",
    "#if MXNET_USE_CUDA\n",
    "    .set_attr<mxnet::FCompute>(\"FCompute<gpu>\",\n",
    "                               mxnet::op::TVMBinaryBackwardComputeUseNone<func_bakcward_vadd_gpu>)\n",
    "#endif  // MXNET_USE_CUDA\n",
    "    .set_attr<mxnet::FCompute>(\"FCompute<cpu>\",\n",
    "                               mxnet::op::TVMBinaryBackwardComputeUseNone<func_bakcward_vadd_cpu>);\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Handle Different DTypes\n",
    "\n",
    "https://docs.tvm.ai/tutorials/language/intrin_math.html#add-your-own-intrinsic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_xxx(dtype, ndim):\n",
    "    A = tvm.placeholder([tvm.var() for _ in range(ndim)], name='A', dtype=dtype)\n",
    "    B = tvm.placeholder([tvm.var() for _ in range(ndim)], name='B', dtype=dtype)\n",
    "    C = tvm.compute([tvm.var() for _ in range(ndim)],\n",
    "                    lambda *index: tvm.if_then_else(A.dtype.startswith(\"float\"), A[index] - B[index], A[index] + B[index]),\n",
    "                    name='C')\n",
    "    s = tvm.create_schedule(C.op)\n",
    "    return s, A, B, C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s, A, B, C = compute_xxx(\"float32\", 1)\n",
    "tvm.lower(s, [A, B, C], simple_mode=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if xxx:\n",
    "    C = aaa\n",
    "else:\n",
    "    C = bbb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@defop(name=\"vadd\", target=\"cpu\", auto_broadcast=True,\n",
    "       dtype=[\"float32\", \"float64\"], ndim=[5])\n",
    "def vadd1(dtype, ndim):\n",
    "    s, A, B, C = compute_add(dtype, ndim)\n",
    "    axes = [axis for axis in C.op.axis]\n",
    "    fused = s[C].fuse(*axes)\n",
    "    s[C].parallel(fused)\n",
    "\n",
    "    return s, [A, B, C]\n",
    "\n",
    "\n",
    "@defop(name=\"vadd\", target=\"cpu\", auto_broadcast=True,\n",
    "       dtype=[\"float16\"], ndim=[5])\n",
    "def vadd2(dtype, ndim):\n",
    "    s, A, B, C = compute_add_f16(dtype, ndim)\n",
    "    axes = [axis for axis in C.op.axis]\n",
    "    fused = s[C].fuse(*axes)\n",
    "    s[C].parallel(fused)\n",
    "\n",
    "    return s, [A, B, C]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Python frontend & unittest\n",
    "\n",
    "......"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}