{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 测试 Tedd" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import set_env" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "# import numpy as np\n", "from IPython.display import display_svg\n", "from tvm import te, build, lower\n", "# from tvm_book.testing.relay.viz import graphviz_relay\n", "from tvm.contrib import tedd" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
# from tvm.script import tir as T\n",
"\n",
"@T.prim_func\n",
"def main(A: T.Buffer((1,), "float32"), B: T.Buffer((1,), "float32"), C: T.Buffer((1,), "float32")):\n",
" T.func_attr({"tir.noalias": T.bool(True)})\n",
" # with T.block("root"):\n",
" for i in range(1):\n",
" with T.block("C"):\n",
" v_i = T.axis.spatial(1, i)\n",
" T.reads(A[v_i], B[v_i])\n",
" T.writes(C[v_i])\n",
" C[v_i] = A[v_i] + B[v_i]\n",
"