Commit ca5929e9 by Tianqi Chen Committed by GitHub

[DOC] Initial doc system (#88)

* [DOC] Initial doc system

* Migrate API

* Update docs
parent 4d280905
[![Build Status](https://travis-ci.com/tqchen/tvm.svg?token=ZQpnpAReT4LHdjWAX8jR&branch=master)](https://travis-ci.com/tqchen/tvm)
[![GitHub license](http://dmlc.github.io/img/apache2.svg)](./LICENSE)
[Installation](docs/how_to/install.md) |
[Documentation](docs) |
[Tutorials](tutorials) |
[Contributor Guide](docs/how_to/contribute.md) |
[Release Notes](NEWS.md)
......
doxygen
modules
tutorials
# Makefile for Sphinx documentation
#
# You can set these variables from the command line.
SPHINXOPTS =
SPHINXBUILD = sphinx-build
PAPER =
BUILDDIR = _build
# User-friendly check for sphinx-build
ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1)
$(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/)
endif
# Internal variables.
PAPEROPT_a4 = -D latex_paper_size=a4
PAPEROPT_letter = -D latex_paper_size=letter
ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
# the i18n builder cannot share the environment and doctrees with the others
I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
.PHONY: help clean html dirhtml singlehtml pickle json htmlhelp qthelp devhelp epub latex latexpdf text man changes linkcheck doctest coverage gettext
help:
@echo "Please use \`make <target>' where <target> is one of"
@echo " html to make standalone HTML files"
@echo " dirhtml to make HTML files named index.html in directories"
@echo " singlehtml to make a single large HTML file"
@echo " pickle to make pickle files"
@echo " json to make JSON files"
@echo " htmlhelp to make HTML files and a HTML help project"
@echo " qthelp to make HTML files and a qthelp project"
@echo " applehelp to make an Apple Help Book"
@echo " devhelp to make HTML files and a Devhelp project"
@echo " epub to make an epub"
@echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter"
@echo " latexpdf to make LaTeX files and run them through pdflatex"
@echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx"
@echo " text to make text files"
@echo " man to make manual pages"
@echo " texinfo to make Texinfo files"
@echo " info to make Texinfo files and run them through makeinfo"
@echo " gettext to make PO message catalogs"
@echo " changes to make an overview of all changed/added/deprecated items"
@echo " xml to make Docutils-native XML files"
@echo " pseudoxml to make pseudoxml-XML files for display purposes"
@echo " linkcheck to check all external links for integrity"
@echo " doctest to run all doctests embedded in the documentation (if enabled)"
@echo " coverage to run coverage check of the documentation (if enabled)"
clean:
rm -rf $(BUILDDIR)/*
html:
$(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html
@echo
@echo "Build finished. The HTML pages are in $(BUILDDIR)/html."
dirhtml:
$(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml
@echo
@echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml."
singlehtml:
$(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml
@echo
@echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml."
pickle:
$(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle
@echo
@echo "Build finished; now you can process the pickle files."
json:
$(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json
@echo
@echo "Build finished; now you can process the JSON files."
htmlhelp:
$(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp
@echo
@echo "Build finished; now you can run HTML Help Workshop with the" \
".hhp project file in $(BUILDDIR)/htmlhelp."
qthelp:
$(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp
@echo
@echo "Build finished; now you can run "qcollectiongenerator" with the" \
".qhcp project file in $(BUILDDIR)/qthelp, like this:"
@echo "# qcollectiongenerator $(BUILDDIR)/qthelp/rabit.qhcp"
@echo "To view the help file:"
@echo "# assistant -collectionFile $(BUILDDIR)/qthelp/rabit.qhc"
applehelp:
$(SPHINXBUILD) -b applehelp $(ALLSPHINXOPTS) $(BUILDDIR)/applehelp
@echo
@echo "Build finished. The help book is in $(BUILDDIR)/applehelp."
@echo "N.B. You won't be able to view it unless you put it in" \
"~/Library/Documentation/Help or install it in your application" \
"bundle."
devhelp:
$(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp
@echo
@echo "Build finished."
@echo "To view the help file:"
@echo "# mkdir -p $$HOME/.local/share/devhelp/rabit"
@echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/rabit"
@echo "# devhelp"
epub:
$(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub
@echo
@echo "Build finished. The epub file is in $(BUILDDIR)/epub."
latex:
$(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
@echo
@echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex."
@echo "Run \`make' in that directory to run these through (pdf)latex" \
"(use \`make latexpdf' here to do that automatically)."
latexpdf:
$(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
@echo "Running LaTeX files through pdflatex..."
$(MAKE) -C $(BUILDDIR)/latex all-pdf
@echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
latexpdfja:
$(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
@echo "Running LaTeX files through platex and dvipdfmx..."
$(MAKE) -C $(BUILDDIR)/latex all-pdf-ja
@echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
text:
$(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text
@echo
@echo "Build finished. The text files are in $(BUILDDIR)/text."
man:
$(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man
@echo
@echo "Build finished. The manual pages are in $(BUILDDIR)/man."
texinfo:
$(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo
@echo
@echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo."
@echo "Run \`make' in that directory to run these through makeinfo" \
"(use \`make info' here to do that automatically)."
info:
$(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo
@echo "Running Texinfo files through makeinfo..."
make -C $(BUILDDIR)/texinfo info
@echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo."
gettext:
$(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale
@echo
@echo "Build finished. The message catalogs are in $(BUILDDIR)/locale."
changes:
$(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes
@echo
@echo "The overview file is in $(BUILDDIR)/changes."
linkcheck:
$(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck
@echo
@echo "Link check complete; look for any errors in the above output " \
"or in $(BUILDDIR)/linkcheck/output.txt."
doctest:
$(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest
@echo "Testing of doctests in the sources finished, look at the " \
"results in $(BUILDDIR)/doctest/output.txt."
coverage:
$(SPHINXBUILD) -b coverage $(ALLSPHINXOPTS) $(BUILDDIR)/coverage
@echo "Testing of coverage in the sources finished, look at the " \
"results in $(BUILDDIR)/coverage/python.txt."
xml:
$(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(BUILDDIR)/xml
@echo
@echo "Build finished. The XML files are in $(BUILDDIR)/xml."
pseudoxml:
$(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml
@echo
@echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml."
The documentation of tvm is generated with recommonmark and sphinx.
- pip install sphinx sphinx-gallery sphinx_rtd_theme matplotlib Image recommonmark
You can build it locally by typing "make html" in this folder.
# -*- coding: utf-8 -*-
#
# documentation build configuration file, created by
# sphinx-quickstart on Thu Jul 23 19:40:08 2015.
#
# This file is execfile()d with the current directory set to its
# containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.
import sys
import os, subprocess
import shlex
import recommonmark
import sphinx_gallery
from recommonmark.parser import CommonMarkParser
from recommonmark.transform import AutoStructify
# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
libpath = os.path.join(curr_path, '../python/')
sys.path.insert(0, libpath)
# -- General configuration ------------------------------------------------
# General information about the project.
project = u'tvm'
author = u'%s developers' % project
copyright = u'2017, %s' % author
github_doc_root = 'https://github.com/tqchen/tvm/tree/master/docs/'
# add markdown parser
CommonMarkParser.github_doc_root = github_doc_root
source_parsers = {
'.md': CommonMarkParser
}
os.environ['TVM_BUILD_DOC'] = '1'
# Version information.
import tvm
version = tvm.__version__
release = tvm.__version__
# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones
extensions = [
'sphinx.ext.autodoc',
'sphinx.ext.autosummary',
'sphinx.ext.napoleon',
'sphinx.ext.mathjax',
'sphinx_gallery.gen_gallery',
]
# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']
# The suffix(es) of source filenames.
# You can specify multiple suffix as a list of string:
# source_suffix = ['.rst', '.md']
source_suffix = ['.rst', '.md']
# The encoding of source files.
#source_encoding = 'utf-8-sig'
# generate autosummary even if no references
autosummary_generate = True
# The master toctree document.
master_doc = 'index'
# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#
# This is also used if you do content translation via gettext catalogs.
# Usually you set "language" from the command line for these cases.
language = None
# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
#today = ''
# Else, today_fmt is used as the format for a strftime call.
#today_fmt = '%B %d, %Y'
# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
exclude_patterns = ['_build']
# The reST default role (used for this markup: `text`) to use for all
# documents.
#default_role = None
# If true, '()' will be appended to :func: etc. cross-reference text.
#add_function_parentheses = True
# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
#add_module_names = True
# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
#show_authors = False
# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'sphinx'
# A list of ignored prefixes for module index sorting.
#modindex_common_prefix = []
# If true, keep warnings as "system message" paragraphs in the built documents.
#keep_warnings = False
# If true, `todo` and `todoList` produce output, else they produce nothing.
todo_include_todos = False
# -- Options for HTML output ----------------------------------------------
# The theme is set by the make target
html_theme = os.environ.get('TVM_THEME', 'rtd')
on_rtd = os.environ.get('READTHEDOCS', None) == 'True'
# only import rtd theme and set it if want to build docs locally
if not on_rtd and html_theme == 'rtd':
import sphinx_rtd_theme
html_theme = 'sphinx_rtd_theme'
html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
#html_static_path = ['_static']
# Output file base name for HTML help builder.
htmlhelp_basename = project + 'doc'
# -- Options for LaTeX output ---------------------------------------------
latex_elements = {
}
# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title,
# author, documentclass [howto, manual, or own class]).
latex_documents = [
(master_doc, '%s.tex' % project, project,
author, 'manual'),
]
# hook for doxygen
def run_doxygen(folder):
"""Run the doxygen make command in the designated folder."""
try:
retcode = subprocess.call("cd %s; make doxygen" % folder, shell=True)
if retcode < 0:
sys.stderr.write("doxygen terminated by signal %s" % (-retcode))
except OSError as e:
sys.stderr.write("doxygen execution failed: %s" % e)
examples_dirs = ['../tutorials/python']
gallery_dirs = ['tutorials']
def generate_doxygen_xml(app):
"""Run the doxygen make commands if we're on the ReadTheDocs server"""
read_the_docs_build = os.environ.get('READTHEDOCS', None) == 'True'
if read_the_docs_build:
run_doxygen('..')
def setup(app):
# Add hook for building doxygen xml when needed
# no c++ API for now
# app.connect("builder-inited", generate_doxygen_xml)
app.add_config_value('recommonmark_config', {
'url_resolver': lambda url: github_doc_root + url,
}, True)
app.add_transform(AutoStructify)
sphinx_gallery_conf = {
'backreferences_dir': 'gen_modules/backreferences',
'doc_module': ('tvm'),
'reference_url': {
'tvm': None
},
'examples_dirs': examples_dirs,
'gallery_dirs': gallery_dirs
}
......@@ -8,7 +8,7 @@ scratch on various systems. It consists of two steps:
To get started, clone tvm repo from github. It is important to clone the submodules along, with ```--recursive``` option.
```bash
git clone --recursive https://github.com/tqchen/tvm
git clone --recursive ssh://git@github.com/tqchen/tvm
```
For windows users who use github tools, you can open the git shell, and type the following command.
```bash
......
TVM Documentation
=================
Welcome to TVM documentation.
Contents
--------
.. toctree::
:maxdepth: 1
how_to/contribute
how_to/install
tutorials/index
python/api
Python API
==========
tvm
---
tvm is a library root namespace contains functions for
declaring computation.
.. autofunction:: tvm.load_json
.. autofunction:: tvm.save_json
.. autofunction:: tvm.var
.. autofunction:: tvm.convert
.. autofunction:: tvm.placeholder
.. autofunction:: tvm.compute
.. autofunction:: tvm.scan
.. autofunction:: tvm.extern
.. autofunction:: tvm.reduce_axis
.. autofunction:: tvm.sum
tvm.tensor
----------
The `tvm.tensor` module contains declaration of Tensor
and Operation class for computation declaration.
.. autoclass:: tvm.tensor.Tensor
:members:
:inherited-members:
.. autoclass:: tvm.tensor.Operation
:members:
:inherited-members:
tvm.schedule
------------
.. autofunction:: tvm.create_schedule
.. autoclass:: tvm.schedule.Schedule
:members:
.. autoclass:: tvm.schedule.Stage
:members:
tvm.build
---------
.. autofunction:: tvm.lower
.. autofunction:: tvm.build
tvm.ndarray
-----------
tvm.ndarray provides a minimum runtime array API to testing out
the correctness of the program.
.. autofunction:: tvm.cpu
.. autofunction:: tvm.gpu
.. autofunction:: tvm.vpi
.. autofunction:: tvm.opencl
.. autofunction:: tvm.ndarray.array
.. autoclass:: tvm.ndarray.TVMContext
:members:
.. autoclass:: tvm.ndarray.NDArray
:members:
:inherited-members:
tvm.Function
------------
.. autofunction:: tvm.register_func
.. autoclass:: tvm.Function
tvm.module
----------
.. autofunction:: tvm.module.load
.. autofunction:: tvm.module.load
.. autoclass:: tvm.module.Module
:members:
:inherited-members:
tvm.node
--------
tvm.node provides
.. autofunction:: tvm.register_node
.. autoclass:: tvm.node.NodeBase
:members:
.. autoclass:: tvm.node.Node
:members:
tvm.expr
--------
.. automodule:: tvm.expr
:members:
# pylint: disable=redefined-builtin, wildcard-import
"""C++ backend related python scripts"""
from __future__ import absolute_import as _abs
from ._ctypes._node import register_node
from . import tensor
from . import arith
......@@ -13,10 +12,16 @@ from . import codegen
from . import collections
from . import schedule
from . import module
from . import node
from . import ndarray as nd
from .ndarray import cpu, gpu, opencl, cl, vpi
from ._ctypes._function import Function
from ._base import TVMError
from ._base import __version__
from .api import *
from .node import register_node
from .schedule import create_schedule
from .build import build, lower
......@@ -136,7 +136,8 @@ def _make_tvm_args(args, temp_args):
class Function(object):
"""A function object at runtime."""
"""The Function object used in TVM.
"""
__slots__ = ["handle", "is_global"]
# pylint: disable=no-member
def __init__(self, handle, is_global=False):
......@@ -158,6 +159,11 @@ class Function(object):
check_call(_LIB.TVMFuncFree(self.handle))
def __call__(self, *args):
"""Call the function with positional arguments
args : list
The positional arguments to the function call.
"""
temp_args = []
values, tcodes, num_args = _make_tvm_args(args, temp_args)
ret_val = TVMValue()
......
......@@ -42,7 +42,7 @@ class SliceBase(object):
pass
class NodeBase(object):
"""Symbol is symbolic graph."""
"""NodeBase is the base class of all TVM language AST object."""
__slots__ = ["handle"]
# pylint: disable=no-member
def __init__(self, handle):
......@@ -119,7 +119,21 @@ class NodeBase(object):
def const(value, dtype=None):
"""construct a constant"""
"""Construct a constant value for a given type.
Parameters
----------
value : int or float
The input value
dtype : str
The data type.
Returns
-------
expr : Expr
Constant expression corresponds to the value.
"""
if dtype is None:
if isinstance(value, Integral):
dtype = 'int32'
......
......@@ -22,6 +22,9 @@ class TempDirectory(object):
"""
return os.path.join(self.temp_dir, name)
def listdir(self):
""""List contents in the dir"""
return os.listdir(self.temp_dir)
def tempdir():
"""Return a new temp dir which deletes the contents when exit
......
......@@ -63,7 +63,7 @@ def save_json(node):
return _api_internal._save_json(node)
def Var(name="tindex", dtype=int32):
def var(name="tindex", dtype=int32):
"""Create a new variable with specified name and dtype
Parameters
......@@ -73,6 +73,11 @@ def Var(name="tindex", dtype=int32):
dtype : int
The data type
Returns
-------
var : Var
The result symbolic variable.
"""
return _api_internal._Var(name, dtype)
......@@ -93,7 +98,7 @@ def placeholder(shape, dtype=None, name="placeholder"):
Returns
-------
tensor: tensor.Tensor
tensor: Tensor
The created tensor
"""
shape = (shape,) if isinstance(shape, _expr.Expr) else shape
......@@ -112,7 +117,7 @@ def compute(shape, fcompute, name="compute"):
shape: Tuple of Expr
The shape of the tensor
fcompute: lambda function of *indices-> value
fcompute: lambda function of indices-> value
Specifies the input source expression
name: str, optional
......@@ -120,7 +125,7 @@ def compute(shape, fcompute, name="compute"):
Returns
-------
tensor: tensor.Tensor
tensor: Tensor
The created tensor
"""
shape = (shape,) if isinstance(shape, _expr.Expr) else shape
......@@ -171,6 +176,8 @@ def scan(init, update, state_placeholder, inputs=None, name="scan"):
Example
-------
.. code-block:: python
# The following code is equivalent to numpy.cumsum
m = tvm.Var("m")
n = tvm.Var("n")
......@@ -234,7 +241,7 @@ def extern(shape, inputs, fcompute,
if not isinstance(t, _tensor.Tensor):
raise ValueError("expect inputs to be tensor")
input_placeholders.append(
Buffer(t.shape, t.dtype, t.op.name))
decl_buffer(t.shape, t.dtype, t.op.name))
types.add(t.dtype)
if dtype is None:
......@@ -244,7 +251,7 @@ def extern(shape, inputs, fcompute,
dtype = [infered_type for _ in shape]
for shp, dt in zip(shape, dtype):
output_placeholders.append(Buffer(shp, dt, name))
output_placeholders.append(decl_buffer(shp, dt, name))
body = fcompute(input_placeholders, output_placeholders)
if isinstance(body, _expr.Expr):
body = _make.Evaluate(body)
......@@ -268,13 +275,13 @@ def call_packed(*args):
int32, "tvm_call_packed", args, 4, None, 0)
def Buffer(shape, dtype=None,
def decl_buffer(shape, dtype=None,
name="buffer",
data=None,
strides=None,
byte_offset=None,
offset_alignment=0):
"""Create a new symbolic buffer
"""Decleare a new symbolic buffer
Parameters
----------
......@@ -308,7 +315,7 @@ def Buffer(shape, dtype=None,
dtype = float32 if dtype is None else dtype
strides = () if strides is None else strides
if data is None:
data = Var(name, "handle")
data = var(name, "handle")
return _api_internal._Buffer(
name, data, shape, strides, dtype, byte_offset, offset_alignment)
......@@ -345,8 +352,8 @@ def _IterVar(dom, name, iter_type, thread_tag=''):
if not isinstance(dom, _collections.Range):
raise ValueError("dom need to be Range")
name = name if name else 'iter'
var = Var(name)
return _api_internal._IterVar(dom, var, iter_type, thread_tag)
v = var(name)
return _api_internal._IterVar(dom, v, iter_type, thread_tag)
def thread_axis(dom=None, tag='', name=''):
......@@ -382,6 +389,11 @@ def reduce_axis(dom, name="rv"):
name : str
The name of the variable.
Returns
-------
axis : IterVar
An iteration variable representing the value.
"""
return _IterVar(dom, name, 2)
......@@ -399,6 +411,11 @@ def sum(expr, axis, where=None):
where : optional, Expr
Filtering predicate of the reduction.
Returns
-------
value : Expr
The result value.
"""
axis = axis if isinstance(axis, list) else [axis]
x = _make.Reduce("Add", expr, axis, where)
......@@ -421,6 +438,11 @@ def min(lhs, rhs=None, axis=None, where=None):
where : optional, Expr
Filtering predicate of the reduction.
Returns
-------
value : Expr
The result value.
"""
if rhs and axis:
raise ValueError("Can only take one argument, rhs or axis")
......@@ -449,6 +471,11 @@ def max(lhs, rhs=None, axis=None, where=None):
where : optional, Expr
Filtering predicate of the reduction.
Returns
-------
value : Expr
The result value.
"""
if rhs and axis:
raise ValueError("Can only take one argument, rhs or axis")
......@@ -461,19 +488,6 @@ def max(lhs, rhs=None, axis=None, where=None):
return x
def Schedule(ops):
"""Create a schedule for list of ops
Parameters
----------
ops : list of Operations
The source expression.
"""
if not isinstance(ops, (list, _collections.Array)):
ops = [ops]
return _api_internal._Schedule(ops)
def convert(value):
"""Convert value to TVM node or function.
......@@ -483,7 +497,7 @@ def convert(value):
Returns
-------
tvm_val : Node or function
tvm_val : Node or Function
Converted value in TVM
"""
if isinstance(value, (Function, NodeBase)):
......
......@@ -46,7 +46,7 @@ def lower(sch,
arg_list = []
for x in args:
if isinstance(x, tensor.Tensor):
buf = api.Buffer(x.shape, dtype=x.dtype, name=x.op.name)
buf = api.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
assert x not in binds
binds[x] = buf
arg_list.append(buf)
......
"""Expression class"""
"""Module to declare Expression class"""
# pylint: disable=missing-docstring
from __future__ import absolute_import as _abs
from ._ctypes._node import NodeBase, register_node
......@@ -58,6 +58,7 @@ class ExprOp(object):
class Expr(NodeBase, ExprOp):
"""Base class of all tvm Expressions"""
pass
class ConstExpr(Expr):
......@@ -72,10 +73,9 @@ class CmpExpr(Expr):
class LogicalExpr(Expr):
pass
@register_node("Variable")
class Var(Expr):
"""Symbolic variable expression."""
pass
@register_node
......
......@@ -9,7 +9,6 @@ import numpy as _np
from ._ctypes._ndarray import TVMContext, TVMType, NDArrayBase
from ._ctypes._ndarray import cpu, gpu, opencl, vpi, empty, sync
from ._ctypes._ndarray import _init_ndarray_module
from ._ctypes._function import Function
cl = opencl
......@@ -40,7 +39,7 @@ def array(arr, ctx=cpu(0)):
Returns
-------
ret : tvm.nd.NDArray
ret : NDArray
The created array
"""
if not isinstance(arr, _np.ndarray):
......
"""Namespace for base node class"""
# pylint: disable=unused-import
from __future__ import absolute_import as _abs
from ._ctypes._node import NodeBase, register_node
Node = NodeBase
......@@ -21,6 +21,24 @@ class Fuse(NodeBase):
pass
def create_schedule(ops):
"""Create a schedule for list of ops
Parameters
----------
ops : list of Operations
The source expression.
Returns
-------
sch : schedule.Schedule
The created schedule.
"""
if not isinstance(ops, (list, _collections.Array)):
ops = [ops]
return _api_internal._Schedule(ops)
@register_node
class Schedule(NodeBase):
"""Schedule for all the stages."""
......@@ -278,7 +296,8 @@ class Stage(NodeBase):
The original y dimension
x_factor : Expr
The stride factor on x axis
y_factor : Expr The stride factor on y axis
y_factor : Expr
The stride factor on y axis
Returns
-------
......
......@@ -57,6 +57,26 @@ class Tensor(NodeBase):
"""Dimension of the tensor."""
return len(self.shape)
@property
def axis(self):
"""Axis of the tensor."""
return self.__getattr__("axis")
@property
def op(self):
"""The corressponding :any:`Operation`."""
return self.__getattr__("op")
@property
def value_index(self):
"""The output value index the tensor corressponds to."""
return self.__getattr__("value_index")
@property
def shape(self):
"""The output shape of the tensor."""
return self.__getattr__("shape")
class Operation(NodeBase):
"""Represent an operation that generate a tensor"""
......
......@@ -61,7 +61,7 @@ void CodeGenC::AddFunction(LoweredFunc f) {
}
std::string CodeGenC::Finish() {
return stream.str();
return decl_stream.str() + stream.str();
}
void CodeGenC::PrintExpr(const Expr& n, std::ostream& os) { // NOLINT(*)
......
......@@ -7,7 +7,7 @@ def lower(s, args, name="mydot"):
for x in args:
assert isinstance(x, tvm.tensor.Tensor)
buf = tvm.Buffer(x.shape, dtype=x.dtype, name=x.op.name)
buf = tvm.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
binds[x] = buf
arg_list.append(buf)
s.normalize()
......@@ -31,7 +31,7 @@ def test_dot():
B = tvm.placeholder((n,), name='B')
k = tvm.reduce_axis((0, n), 'k')
C = tvm.compute((1,), lambda _: tvm.sum(A[k] * B[k], axis=k), name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
fapi = lower(s, [A, B, C])
def verify(target):
......
......@@ -8,7 +8,7 @@ def test_add():
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
# schedule
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
# create iter var and assign them tags.
num_thread = 256
bx, x = s[C].split(C.op.axis[0], factor=num_thread*4)
......
......@@ -5,7 +5,7 @@ import numpy as np
def test_gemm():
# graph
nn = 1024
n = tvm.Var('n')
n = tvm.var('n')
n = tvm.convert(nn)
m = n
l = n
......@@ -17,7 +17,7 @@ def test_gemm():
lambda ii, jj: tvm.sum(A[ii, k] * B[jj, k], axis=k),
name='CC')
# schedule
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
xtile, ytile = 32, 32
scale = 8
num_thread = 8
......
......@@ -3,13 +3,13 @@ import numpy as np
def test_sum():
# graph
n = tvm.Var('n')
m = tvm.Var('m')
n = tvm.var('n')
m = tvm.var('m')
A = tvm.placeholder((n, m), name='A')
k = tvm.reduce_axis((0, m))
B = tvm.compute((n,), lambda i: tvm.sum(A[i, k], axis=k, where=(i>1)), name='B')
# schedule
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
# create iter var and assign them tags.
num_thread = 1
xo, xi = s[B].split(B.op.axis[0], factor=num_thread)
......@@ -52,7 +52,7 @@ def test_rfactor():
k = tvm.reduce_axis((0, n))
B = tvm.compute((1,), lambda i: tvm.sum(A[k], axis=k), name='B')
# schedule
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
kf, ki = s[B].split(k, nparts=4)
BF = s.rfactor(B, kf)
s[BF].parallel(BF.op.axis[0])
......@@ -87,7 +87,7 @@ def test_rfactor_threads():
nthread = 16
B = tvm.compute((m,), lambda i: tvm.sum(A[i, k], axis=k, where=(i>1)), name='B')
# schedule
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
ko, kf = s[B].split(k, factor=nthread)
BF = s.rfactor(B, kf)
bx, ty = s[B].split(s[B].op.axis[0], factor=nthread)
......
......@@ -2,8 +2,8 @@ import tvm
import numpy as np
def test_scan():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
X = tvm.placeholder((m, n), name="X")
s_state = tvm.placeholder((m, n))
s_init = tvm.compute((1, n), lambda _, i: X[0, i])
......@@ -11,7 +11,7 @@ def test_scan():
res = tvm.scan(s_init, s_update, s_state)
# schedule
s = tvm.Schedule(res.op)
s = tvm.create_schedule(res.op)
num_thread = 256
block_x = tvm.thread_axis(None, "blockIdx.x")
thread_x = tvm.thread_axis((0, num_thread), "threadIdx.x")
......
......@@ -47,7 +47,7 @@ def rnn_matexp():
max_auto_unroll_step = 0
detect_global_barrier = DETECT_GLOBAL_BARRIER
num_step = tvm.Var("num_step")
num_step = tvm.var("num_step")
num_hidden = tvm.convert(n_num_hidden)
batch_size = tvm.convert(n_batch_size)
num_thread_y = 8
......@@ -65,7 +65,7 @@ def rnn_matexp():
name="update")
s_scan = tvm.scan(s_init, s_update, s_state)
# schedule
s = tvm.Schedule(s_scan.op)
s = tvm.create_schedule(s_scan.op)
CL = s_update
SS = s.cache_read(s_state, "shared", [CL])
SL = s.cache_read(SS, "local", [CL])
......
import tvm
def test_basic():
a = tvm.Var("a")
b = tvm.Var("b")
a = tvm.var("a")
b = tvm.var("b")
m = tvm.arith.DetectLinearEquation(a * 4 + b * 6 + 7, a)
assert m[1].value == 4
assert tvm.ir_pass.Simplify(m[0] - (b * 6 + 7)).value == 0
......
......@@ -6,10 +6,10 @@ def test_basic():
assert s.max().value == 3
def test_deduce():
a = tvm.Var('a')
b = tvm.Var('b')
c = tvm.Var('c')
d = tvm.Var('d')
a = tvm.var('a')
b = tvm.var('b')
c = tvm.var('c')
d = tvm.var('d')
b_s = tvm.arith.intset_interval(2, 3)
c_s = tvm.arith.intset_interval(10, 15)
......@@ -36,10 +36,10 @@ def test_deduce():
assert str(tvm.ir_pass.Simplify(res3.min())) == str(ans3)
def test_check():
a = tvm.Var('a')
b = tvm.Var('b')
c = tvm.Var('c')
d = tvm.Var('d')
a = tvm.var('a')
b = tvm.var('b')
c = tvm.var('c')
d = tvm.var('d')
b_s = tvm.arith.intset_interval(2, 3)
c_s = tvm.arith.intset_interval(5, 7)
......
import tvm
def test_basic():
a = tvm.Var()
b = tvm.Var()
a = tvm.var()
b = tvm.var()
m = tvm.arith.EvalModular(a * 4 + b * 6 + 7)
assert m.coeff == 2
assert m.base == 1
......
......@@ -3,11 +3,11 @@ from tvm.addon import testing
import numpy as np
def test_add_pipeline():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
# GPU schedule have to split by gridIdx and threadIdx
num_thread = 256
......@@ -18,9 +18,9 @@ def test_add_pipeline():
# compile to IR
bounds = tvm.schedule.InferBound(s)
stmt = tvm.schedule.ScheduleOps(s, bounds)
Ab = tvm.Buffer(A.shape, A.dtype, name='A')
Bb = tvm.Buffer(B.shape, B.dtype, name='B')
Cb = tvm.Buffer(C.shape, C.dtype, name='C')
Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
Bb = tvm.decl_buffer(B.shape, B.dtype, name='B')
Cb = tvm.decl_buffer(C.shape, C.dtype, name='C')
stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb})
stmt = tvm.ir_pass.Simplify(stmt)
fapi = tvm.ir_pass.MakeAPI(stmt, "myadd", [Ab, Bb, Cb], 0)
......
......@@ -7,7 +7,7 @@ def test_add_pipeline():
A = tvm.placeholder((n,), name='A')
def extern_generator(ins, outs):
"""Manually write the IR for the extern function, add pipeline"""
i = tvm.Var('i')
i = tvm.var('i')
stmt = tvm.make.For(
i, 0, n, 0, 0,
tvm.make.Store(outs[0].data,
......@@ -15,7 +15,7 @@ def test_add_pipeline():
1, i))
return stmt
C = tvm.extern(A.shape, [A], extern_generator, name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
def check_llvm():
if not tvm.codegen.enabled("llvm"):
......
......@@ -7,7 +7,7 @@ def test_llvm_add_pipeline():
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(C.op.axis[0], factor=4)
s[C].parallel(xo)
s[C].vectorize(xi)
......@@ -35,7 +35,7 @@ def test_llvm_flip_pipeline():
n = tvm.convert(nn)
A = tvm.placeholder((n + base), name='A')
C = tvm.compute((n,), lambda i: A(nn + base- i - 1), name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(C.op.axis[0], factor=4)
s[C].parallel(xo)
s[C].vectorize(xi)
......@@ -62,7 +62,7 @@ def test_llvm_madd_pipeline():
n = tvm.convert(nn)
A = tvm.placeholder((n + base, stride), name='A')
C = tvm.compute((n, stride), lambda i, j: A(base + i, j) + 1, name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(C.op.axis[0], factor=4)
s[C].parallel(xo)
s[C].vectorize(xi)
......@@ -86,7 +86,7 @@ def test_llvm_temp_space():
A = tvm.placeholder((n,), name='A')
B = tvm.compute(A.shape, lambda i: A(i) + 1, name='B')
C = tvm.compute(A.shape, lambda i: B(i) + 1, name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
def check_llvm():
if not tvm.codegen.enabled("llvm"):
......
......@@ -17,8 +17,8 @@ def test_stack_vm_basic():
print(shape0)
assert shape0 == a.shape[0]
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), tvm.float32)
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), tvm.float32)
stmt = tvm.make.Evaluate(tvm.call_packed("tvm_call_back_get_shape", Ab.shape[0]))
fapi = tvm.ir_pass.MakeAPI(stmt, "print_shape", [Ab], 0)
run_jit(fapi, lambda f: f(a))
......@@ -31,9 +31,9 @@ def tvm_stack_vm_print(*x):
def test_stack_vm_loop():
dtype = 'int64'
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), dtype)
i = tvm.Var('i')
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), dtype)
i = tvm.var('i')
# for i in 0 to n-1:
stmt = tvm.make.For(
i, 0, n - 1, 0, 0,
......@@ -52,9 +52,9 @@ def test_stack_vm_loop():
def test_stack_vm_cond():
dtype = 'int64'
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), dtype)
i = tvm.Var('i')
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), dtype)
i = tvm.var('i')
# for i in 0 to n-1:
stmt = tvm.make.For(
i, 0, n - 1, 0, 0,
......
......@@ -30,14 +30,14 @@ def test_ir():
assert isinstance(stmt, tvm.stmt.Evaluate)
def test_let():
x = tvm.Var('x')
y = tvm.Var('y')
x = tvm.var('x')
y = tvm.var('y')
stmt = tvm.make.LetStmt(
x, 10, tvm.make.Evaluate(x + 1));
def test_attr():
x = tvm.Var('x')
y = tvm.Var('y')
x = tvm.var('x')
y = tvm.var('y')
stmt = tvm.make.AttrStmt(
y, "stride", 10, tvm.make.Evaluate(x + 1));
assert stmt.node == y
......@@ -52,15 +52,15 @@ def test_attr():
def test_basic():
a = tvm.Var('a')
b = tvm.Var('b')
a = tvm.var('a')
b = tvm.var('b')
c = a + b
assert str(c) == '(%s + %s)' % (a.name, b.name)
def test_stmt():
x = tvm.make.Evaluate(0)
tvm.make.For(tvm.Var('i'), 0, 1,
tvm.make.For(tvm.var('i'), 0, 1,
tvm.stmt.For.Serial, 0,
x)
......
import tvm
def test_buffer():
m = tvm.Var('m')
n = tvm.Var('n')
l = tvm.Var('l')
Ab = tvm.Buffer((m, n), tvm.float32)
Bb = tvm.Buffer((n, l), tvm.float32)
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
Ab = tvm.decl_buffer((m, n), tvm.float32)
Bb = tvm.decl_buffer((n, l), tvm.float32)
assert isinstance(Ab, tvm.schedule.Buffer)
assert Ab.dtype == tvm.float32
......
......@@ -11,8 +11,8 @@ def test_array_save_load_json():
assert(a[1].value == 2)
def test_map():
a = tvm.Var('a')
b = tvm.Var('b')
a = tvm.var('a')
b = tvm.var('b')
amap = tvm.convert({a: 2,
b: 3})
assert a in amap
......@@ -22,8 +22,8 @@ def test_map():
assert a + 1 not in amap
def test_map_save_load_json():
a = tvm.Var('a')
b = tvm.Var('b')
a = tvm.var('a')
b = tvm.var('b')
amap = tvm.convert({a: 2,
b: 3})
json_str = tvm.save_json(amap)
......
......@@ -2,8 +2,8 @@
import tvm
def test_scan_group():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
s_state = tvm.placeholder((m, n))
s_init = tvm.compute((1, n), lambda _, i: x[0, i])
......@@ -13,7 +13,7 @@ def test_scan_group():
s_update3 = tvm.compute((m, n), lambda t, i: s_update2[t, i] + 1)
res = tvm.scan(s_init, s_update3, s_state, inputs=x)
s = tvm.Schedule(res.op)
s = tvm.create_schedule(res.op)
assert s[s_update1].group is not None
assert s[s_update2].group == s[s_update1].group
# Assign within group, is valid
......@@ -34,12 +34,12 @@ def test_scan_group():
pass
def test_compute_group():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
x1 = tvm.compute(x.shape, lambda *i: x(*i) + 1, name="x1")
x2 = tvm.compute(x.shape, lambda *i: x1(*i) + 2, name="x2")
s = tvm.Schedule(x2.op)
s = tvm.create_schedule(x2.op)
g = s.create_group(outputs=x1, inputs=x, include_inputs=True)
assert s[x1].group == g
assert s[x].group == g
......@@ -48,12 +48,12 @@ def test_compute_group():
assert g.num_child_stages == 2
def test_nest_group():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
x1 = tvm.compute(x.shape, lambda *i: x(*i) + 1, name="x1")
x2 = tvm.compute(x.shape, lambda *i: x1(*i) + 2, name="x2")
s = tvm.Schedule(x2.op)
s = tvm.create_schedule(x2.op)
g1 = s.create_group(outputs=x1, inputs=x)
g2 = s.create_group(outputs=x1, inputs=x, include_inputs=True)
assert set(s.groups) == set([g1, g2])
......
......@@ -2,14 +2,14 @@ import tvm
import pickle as pkl
def test_schedule_create():
m = tvm.Var('m')
n = tvm.Var('n')
l = tvm.Var('l')
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
AA = tvm.compute((m, l), lambda i, j: A[i, j])
T = tvm.compute((m, n, l), lambda i, j, k: AA(i, k) * B(j, k))
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
s[AA].set_scope("shared")
xo, xi = s[T].split(T.op.axis[0], factor=10)
xi1, xi2 = s[T].split(xi, factor=2)
......@@ -31,11 +31,11 @@ def test_schedule_create():
assert(str(s_loaded.outputs[0].body) == str(s.outputs[0].body))
def test_reorder():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m,), name='A')
T = tvm.compute(m, lambda i: A[i+1])
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=10)
xi1, xi2 = s[T].split(xi, factor=2)
order = (xi2, xi1, xo)
......@@ -44,45 +44,45 @@ def test_reorder():
assert tuple(s[T].leaf_iter_vars) == order
def test_split():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m,), name='A')
T = tvm.compute((m,), lambda i: A[i])
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=10)
assert tuple(s[T].leaf_iter_vars) == (xo, xi)
def test_tile():
m = tvm.Var('m')
n = tvm.Var('n')
m = tvm.var('m')
n = tvm.var('n')
A = tvm.placeholder((m, n), name='A')
T = tvm.compute((m, n), lambda i, j: A[i, j])
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
xo, yo, xi, yi = s[T].tile(T.op.axis[0], T.op.axis[1], x_factor=10, y_factor=5)
assert tuple(s[T].leaf_iter_vars) == (xo, yo, xi, yi)
def test_fuse():
m = tvm.Var('m')
n = tvm.Var('n')
m = tvm.var('m')
n = tvm.var('n')
A = tvm.placeholder((m, n), name='A')
T = tvm.compute((m, n), lambda i, j: A[i, j])
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
xo, yo, xi, yi = s[T].tile(T.op.axis[0], T.op.axis[1], x_factor=10, y_factor=5)
fused = s[T].fuse(yo, xo)
assert any(isinstance(x, tvm.schedule.Fuse) for x in s[T].relations)
assert tuple(s[T].leaf_iter_vars) == (fused, xi, yi)
def test_vectorize():
m = tvm.Var('m')
n = tvm.Var('n')
m = tvm.var('m')
n = tvm.var('n')
A = tvm.placeholder((m, n), name='A')
T = tvm.compute((m, n), lambda i, j: A[i, j])
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
xo, yo, xi, yi = s[T].tile(T.op.axis[0], T.op.axis[1], x_factor=10, y_factor=5)
s[T].vectorize(yi)
s[T].unroll(xi)
......@@ -92,20 +92,20 @@ def test_vectorize():
assert s[T].iter_var_attrs[yi].iter_type == VECTORIZE
def test_rfactor():
n = tvm.Var('n')
n = tvm.var('n')
k1 = tvm.reduce_axis((0, n), name="k1")
k2 = tvm.reduce_axis((0, n), name="k2")
A = tvm.placeholder((n, n, n), name='A')
B = tvm.compute((n, ), lambda i: tvm.sum(A[i, k1, k2], axis=[k1, k2]))
# normal schedule
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
BF = s.rfactor(B, k1)
assert(tuple(BF.shape) == (n, n))
assert(set(BF.op.body.axis) == set([k2]))
assert(s[B].op.body.axis[0].dom.extent == n)
assert(len(s[B].all_iter_vars) == 2)
# schedule with splot
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
ko, ki = s[B].split(k1, factor=4)
xo, xi = s[B].split(B.op.axis[0], factor=8)
BF = s.rfactor(B, ki)
......
import tvm
def test_tensor():
m = tvm.Var('m')
n = tvm.Var('n')
l = tvm.Var('l')
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
T = tvm.compute((m, n, l), lambda i, j, k: A[i, k] * B[j, k])
......@@ -19,7 +19,7 @@ def test_tensor():
def test_conv1d():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.placeholder((n+2), name='A')
def computeB(ii):
i = ii + 1
......@@ -28,15 +28,15 @@ def test_conv1d():
def test_tensor_slice():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.compute((n, n), lambda i, j: 1)
B = tvm.compute((n,), lambda i: A[0][i] + A[0][i])
def test_tensor_reduce():
m = tvm.Var('m')
n = tvm.Var('n')
l = tvm.Var('l')
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
T = tvm.compute((m, n, l), lambda i, j, k: A[i, k] * B[j, k])
......@@ -50,8 +50,8 @@ def test_tensor_reduce():
def test_tensor_scan():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.placeholder((m, n))
s = tvm.placeholder((m, n))
res = tvm.scan(tvm.compute((1, n), lambda _, i: x[0, i]),
......@@ -60,8 +60,8 @@ def test_tensor_scan():
assert tuple(res.shape) == (m, n)
def test_scan_multi_out():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x1 = tvm.placeholder((m, n))
s1 = tvm.placeholder((m, n))
x2 = tvm.placeholder((m, n))
......@@ -81,7 +81,7 @@ def test_scan_multi_out():
assert isinstance(zz, tvm.tensor.ScanOp)
def test_extern():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m,), name='A')
def extern_func(ins, outs):
......@@ -92,7 +92,7 @@ def test_extern():
def test_extern_multi_out():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i] * 10)
......
......@@ -26,9 +26,9 @@ def test_dso_module_load():
temp = testing.tempdir()
def save_object(names):
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), dtype)
i = tvm.Var('i')
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), dtype)
i = tvm.var('i')
# for i in 0 to n-1:
stmt = tvm.make.For(
i, 0, n - 1, 0, 0,
......
import tvm
def test_simplify():
x = tvm.Var('x')
x = tvm.var('x')
e1 = tvm.ir_pass.Simplify(x + 2 + 1)
assert(tvm.ir_pass.Equal(e1, x + 3))
e2 = tvm.ir_pass.Simplify(x * 3 + 5 * x)
......@@ -14,15 +14,15 @@ def test_simplify():
def test_verify_ssa():
x = tvm.Var('x')
y = tvm.Var()
x = tvm.var('x')
y = tvm.var()
z = tvm.make.Evaluate(x + y)
assert(tvm.ir_pass.VerifySSA(z))
def test_convert_ssa():
x = tvm.Var('x')
y = tvm.Var()
x = tvm.var('x')
y = tvm.var()
let1 = tvm.make.Let(x, 1, x + 1)
let2 = tvm.make.Let(x, 1, x + y)
z = tvm.make.Evaluate(let1 + let2)
......@@ -32,7 +32,7 @@ def test_convert_ssa():
def test_expr_use_var():
x = tvm.Var('x')
x = tvm.var('x')
assert(tvm.ir_pass.ExprUseVar(x+1, x))
assert(not tvm.ir_pass.ExprUseVar(1+10, x))
......
import tvm
def test_inline():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m,), name='A')
T = tvm.compute((m,), lambda i,: A[i] + 10, name='T')
stmt = tvm.make.Evaluate(T[10] + 11 * T[100])
......
import tvm
def test_basic():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.placeholder((n, ), name='A')
B = tvm.placeholder((n, ), name='B')
T = tvm.compute((n, ), lambda i: A[i]+B[i])
s = tvm.Schedule(T.op)
s = tvm.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=4)
bounds = tvm.schedule.InferBound(s)
......@@ -16,11 +16,11 @@ def test_basic():
print(stmt)
def test_multi_loop():
i = tvm.Var('i')
j = tvm.Var('j')
k = tvm.Var('k')
m = tvm.Var('m')
n = tvm.Var('n')
i = tvm.var('i')
j = tvm.var('j')
k = tvm.var('k')
m = tvm.var('m')
n = tvm.var('n')
stmt = tvm.make.For(
i, 0, 4, 0, 0,
tvm.make.For(
......@@ -34,11 +34,11 @@ def test_multi_loop():
print(stmt)
def test_multi_if():
i = tvm.Var('i')
j = tvm.Var('j')
k = tvm.Var('k')
m = tvm.Var('m')
n = tvm.Var('n')
i = tvm.var('i')
j = tvm.var('j')
k = tvm.var('k')
m = tvm.var('m')
n = tvm.var('n')
stmt = tvm.make.For(
i, 0, 4, 0, 0,
tvm.make.For(
......@@ -54,12 +54,12 @@ def test_multi_if():
print(stmt)
def test_thread_axis():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.compute((m, l), lambda i, j: A[i, j] + 3, name='B')
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
s[B].set_scope("shared")
num_thread = 16
......
......@@ -3,18 +3,18 @@ import numpy
def test_makeapi():
"""Not yet working, mock design"""
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
bounds = tvm.schedule.InferBound(s)
stmt = tvm.schedule.ScheduleOps(s, bounds)
Ab = tvm.Buffer(A.shape, A.dtype, name='A')
Bb = tvm.Buffer(B.shape, B.dtype, name='B')
Cb = tvm.Buffer(C.shape, C.dtype, name='C')
Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
Bb = tvm.decl_buffer(B.shape, B.dtype, name='B')
Cb = tvm.decl_buffer(C.shape, C.dtype, name='C')
stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb})
num_unpacked_args = 2
......
import tvm
def test_remove_no_op():
i = tvm.Var('i')
j = tvm.Var('j')
k = tvm.Var('k')
m = tvm.Var('m')
n = tvm.Var('n')
i = tvm.var('i')
j = tvm.var('j')
k = tvm.var('k')
m = tvm.var('m')
n = tvm.var('n')
dtype = 'int64'
Ab = tvm.Buffer((n, ), dtype)
Ab = tvm.decl_buffer((n, ), dtype)
stmt = tvm.make.For(
i, 0, 4, 0, 0,
tvm.make.For(
......
......@@ -4,10 +4,10 @@ import numpy
def test_simplify():
"""Not yet working, mock design"""
dtype = 'int64'
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), dtype)
i = tvm.Var('i')
j = tvm.Var('j')
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), dtype)
i = tvm.var('i')
j = tvm.var('j')
# for i in 0 to n-1:
stmt = tvm.make.For(
i, 2, n, 0, 0,
......@@ -22,7 +22,7 @@ def test_simplify():
def test_basic():
m = tvm.Var('m')
m = tvm.var('m')
ret = tvm.ir_pass.CanonicalSimplify(tvm.make.Evaluate(m-1))
assert str(ret.value) == "(m - 1)"
......
......@@ -6,7 +6,7 @@ def lower(s, args):
for x in args:
assert isinstance(x, tvm.tensor.Tensor)
buf = tvm.Buffer(x.shape, dtype=x.dtype, name=x.op.name)
buf = tvm.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
binds[x] = buf
arg_list.append(buf)
s.normalize()
......@@ -28,7 +28,7 @@ def test_basic_pipeline():
stages.append(B)
B = tvm.compute((n,), lambda i: B[i] + k, name="A%s" % k)
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
xo, xi = s[B].split(B.op.axis[0], nparts=1)
s[B].bind(xo, tvm.thread_axis("pipeline"))
xo, xi = s[B].split(xi, factor=4)
......@@ -43,13 +43,13 @@ def test_basic_pipeline():
assert(tvm.ir_pass.VerifySSA(stmt))
def test_conv1d():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.compute((n+2), lambda i: 1, name='A')
def computeB(ii):
i = ii + 1
return A[i-1] + A[i] + A[i+1]
B = tvm.compute(n, computeB, name='B')
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
px, xi = s[B].split(B.op.axis[0], nparts=1)
s[B].bind(px, tvm.thread_axis("pipeline"))
s[A].compute_at(s[B], px)
......
import tvm
def test_flatten2():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3, name='A2')
s = tvm.Schedule(A2.op)
s = tvm.create_schedule(A2.op)
xo, xi = s[A2].split(A2.op.axis[0], 8)
s[A1].compute_at(s[A2], xo)
bounds = tvm.schedule.InferBound(s)
......@@ -15,8 +15,8 @@ def test_flatten2():
stmt = tvm.schedule.ScheduleOps(s, bounds)
print(stmt)
Ab = tvm.Buffer(A.shape, A.dtype, name='A')
A2b = tvm.Buffer(A2.shape, A2.dtype, name='A2')
Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
A2b = tvm.decl_buffer(A2.shape, A2.dtype, name='A2')
stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b})
stmt = tvm.ir_pass.Simplify(stmt)
print(stmt)
......
import tvm
def test_storage_sync():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3, name='A2')
s = tvm.Schedule(A2.op)
s = tvm.create_schedule(A2.op)
xo, xi = s[A2].split(A2.op.axis[0], factor=8)
s[A2].bind(xo, tvm.thread_axis("blockIdx.x"))
s[A1].compute_at(s[A2], xo)
......@@ -18,8 +18,8 @@ def test_storage_sync():
assert isinstance(bounds, tvm.collections.Map)
stmt = tvm.schedule.ScheduleOps(s, bounds)
print(stmt)
Ab = tvm.Buffer(A.shape, A.dtype, name='A')
A2b = tvm.Buffer(A2.shape, A2.dtype, name='A2')
Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
A2b = tvm.decl_buffer(A2.shape, A2.dtype, name='A2')
stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b})
f = tvm.ir_pass.MakeAPI(stmt, "test", [Ab, A2b], 0)
flist = tvm.ir_pass.SplitHostDevice(f)
......
......@@ -2,10 +2,10 @@ import tvm
def test_unroll_loop():
dtype = 'int64'
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), dtype)
i = tvm.Var('i')
j = tvm.Var('j')
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), dtype)
i = tvm.var('i')
j = tvm.var('j')
# for i in 0 to n-1:
stmt = tvm.make.For(
i, n, 2, 0, 0,
......
......@@ -2,10 +2,10 @@ import tvm
def test_vectorize_loop():
dtype = 'int64'
n = tvm.Var('n')
Ab = tvm.Buffer((n, ), dtype)
i = tvm.Var('i')
j = tvm.Var('j')
n = tvm.var('n')
Ab = tvm.decl_buffer((n, ), dtype)
i = tvm.var('i')
j = tvm.var('j')
VECTORIZE = 2
# for i in 0 to n-1:
stmt = tvm.make.For(
......
import tvm
def test_virtual_thread():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m, ), name='A')
A1 = tvm.compute((m,), lambda i: A[i], name='A1')
A2 = tvm.compute((m,), lambda i: A1[i] + 3, name='A2')
s = tvm.Schedule(A2.op)
s = tvm.create_schedule(A2.op)
vx = tvm.thread_axis("vthread", name="vx")
xo, xi = s[A2].split(A2.op.axis[0], nparts=2)
s[A2].bind(xo, vx)
......@@ -17,8 +17,8 @@ def test_virtual_thread():
assert isinstance(bounds, tvm.collections.Map)
stmt = tvm.schedule.ScheduleOps(s, bounds)
Ab = tvm.Buffer(A.shape, A.dtype, name='A')
A2b = tvm.Buffer(A2.shape, A2.dtype, name='A2')
Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
A2b = tvm.decl_buffer(A2.shape, A2.dtype, name='A2')
stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b})
stmt = tvm.ir_pass.Simplify(stmt)
stmt = tvm.ir_pass.InjectVirtualThread(stmt)
......
......@@ -10,7 +10,7 @@ def test_get_global():
return 10
# get it out from global function table
f = tvm.get_global_func("my_packed_func")
assert isinstance(f, tvm.nd.Function)
assert isinstance(f, tvm.Function)
y = f(*targs)
assert y == 10
......@@ -32,7 +32,7 @@ def test_convert():
assert(tuple(args) == targs)
f = tvm.convert(myfunc)
assert isinstance(f, tvm.nd.Function)
assert isinstance(f, tvm.Function)
f(*targs)
def test_byte_array():
......
import tvm
def test_bound1():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3, name='A2')
s = tvm.Schedule([A2.op])
s = tvm.create_schedule([A2.op])
xo, xi = s[A2].split(s[A2].op.axis[0], 8)
s[A1].compute_at(s[A2], xo)
bounds = tvm.schedule.InferBound(s)
......@@ -15,12 +15,12 @@ def test_bound1():
assert(bounds[A1.op.axis[0]].extent.value == 8)
def test_bound2():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3, name='A2')
s = tvm.Schedule(A2.op)
s = tvm.create_schedule(A2.op)
xo, yo, xi, yi = s[A2].tile(A2.op.axis[0], A2.op.axis[1], 8, 8)
s[A1].compute_at(s[A2], yo)
bounds = tvm.schedule.InferBound(s)
......@@ -29,13 +29,13 @@ def test_bound2():
assert(bounds[A1.op.axis[1]].extent.value == 8)
def test_bound3():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3, name='A2')
s = tvm.Schedule(A2.op)
s = tvm.create_schedule(A2.op)
s[A1].set_scope("shared")
xo, xi = s[A2].split(A2.op.axis[0], 32)
xi0, xi1 = s[A2].split(xi, nparts=16)
......@@ -50,8 +50,8 @@ def test_bound3():
assert(bounds[A1.op.axis[1]].extent.value==16)
def test_bound_scan():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
X = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
s_state = tvm.placeholder((m, n))
s_init = tvm.compute((1, n), lambda _, i: X[0, i])
......@@ -59,7 +59,7 @@ def test_bound_scan():
s_scan = tvm.scan(s_init, s_update, s_state)
assert tuple(s_scan.shape) == (m, n)
s = tvm.Schedule(s_scan.op)
s = tvm.create_schedule(s_scan.op)
XX = s.cache_read(X, "local", s_update)
xo, xi = s[s_update].split(s_update.op.axis[1], factor=4)
s[XX].compute_at(s[s_update], xo)
......@@ -69,13 +69,13 @@ def test_bound_scan():
assert bounds[XX.op.axis[1]].extent.value == 4
def test_bound_conv1d():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.compute((n+2), lambda i: 1, name='A')
def computeB(ii):
i = ii + 1
return A[i-1] + A[i] + A[i+1]
B = tvm.compute(n, computeB, name='B')
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
s[A].compute_at(s[B], B.op.axis[0])
s.normalize()
bounds = tvm.schedule.InferBound(s)
......@@ -90,7 +90,7 @@ def test_bound_blur():
j = jj + 1
return A[i][j] + A[i-1][j] + A[i+1][j] + A[i][j+1] + A[i][j-1]
B = tvm.compute((n-2, n-2), computeB, name='B')
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
s[A].compute_at(s[B], B.op.axis[1])
s.normalize()
bounds = tvm.schedule.InferBound(s)
......@@ -98,12 +98,12 @@ def test_bound_blur():
assert(bounds[A.op.axis[1]].extent.value == 3)
def test_bound_rfactor():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.placeholder((n,), name='A')
k = tvm.reduce_axis((0, n))
B = tvm.compute((1,), lambda i: tvm.sum(A[k], axis=k, where=(i>1)), name='B')
# schedule
s = tvm.Schedule(B.op)
s = tvm.create_schedule(B.op)
kf, ki = s[B].split(k, nparts=4)
BF = s.rfactor(B, kf)
s.normalize()
......@@ -113,12 +113,12 @@ def test_bound_rfactor():
assert(bounds[BF.op.axis[1]].extent.value == 1)
def test_bound_group_schedule():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
x1 = tvm.compute(x.shape, lambda *i: x(*i) + 1, name="x1")
x2 = tvm.compute(x.shape, lambda *i: x1(*i) + 2, name="x2")
s = tvm.Schedule(x2.op)
s = tvm.create_schedule(x2.op)
g = s.create_group(outputs=x1, inputs=x, include_inputs=True)
g.compute_at(s[x2], x2.op.axis[0])
assert s[x1].group == g
......@@ -129,12 +129,12 @@ def test_bound_group_schedule():
assert bounds[x.op.axis[1]].extent == n
def test_bound_nest_group():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
x1 = tvm.compute(x.shape, lambda *i: x(*i) + 1, name="x1")
x2 = tvm.compute(x.shape, lambda *i: x1(*i) + 2, name="x2")
s = tvm.Schedule(x2.op)
s = tvm.create_schedule(x2.op)
g1 = s.create_group(outputs=x, inputs=x, include_inputs=True)
g2 = s.create_group(outputs=x1, inputs=x, include_inputs=True)
assert s[x].group == g1
......@@ -150,13 +150,13 @@ def test_bound_nest_group():
def test_bound_nest_thread():
m = tvm.Var('m')
m = tvm.var('m')
A = tvm.placeholder((m), name='A')
A1 = tvm.compute((m,), lambda i: A[i], name='A1')
A2 = tvm.compute((m,), lambda i: A1[i] + 2, name='A2')
A3 = tvm.compute((m,), lambda i: A2[i] + 3, name='A3')
s = tvm.Schedule(A3.op)
s = tvm.create_schedule(A3.op)
s[A2].set_scope("shared")
s[A1].set_scope("local")
......@@ -186,7 +186,7 @@ def test_gemm_bound():
lambda ii, jj: tvm.sum(A[ii, k] * B[jj, k], axis=k),
name='CC')
# schedule
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
xtile, ytile = 32, 32
scale = 8
num_thread = 8
......
import tvm
def test_scan():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
s_state = tvm.placeholder((m, n))
s_init = tvm.compute((1, n), lambda _, i: x[0, i], name="s_init")
......@@ -16,7 +16,7 @@ def test_scan():
assert set(body) == set([s_scan.op, s_update.op, s_up1.op])
def test_attach_path():
s = tvm.Schedule(s_scan.op)
s = tvm.create_schedule(s_scan.op)
s[x_trans].compute_at(s[s_update], s_update.op.axis[0])
apath = tvm.schedule.CreateAttachPath(s)
assert(tuple(apath[s_update.op]) == tuple([s_scan.op.scan_axis]))
......@@ -28,9 +28,9 @@ def test_scan():
assert(fxpt[s_scan.spatial_axis_[0]].value != 0)
def test_scan_fix_point():
m = tvm.Var("m")
n = tvm.Var("n")
l = tvm.Var("l")
m = tvm.var("m")
n = tvm.var("n")
l = tvm.var("l")
x = tvm.compute((l, m, n), lambda *i: tvm.const(1, "float32"), name="x")
s_state = tvm.placeholder((l, m, n))
s_init = tvm.compute((1, m, n), lambda _, i, j: x[0, i, j], name="s_init")
......@@ -74,8 +74,8 @@ def test_scan_fix_point():
assert(fxpt[s_scan.op.spatial_axis_[1]].value == 0)
def test_scan5_multi_output():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x1 = tvm.placeholder((m, n))
s1 = tvm.placeholder((m, n))
x2 = tvm.placeholder((m, n))
......@@ -98,8 +98,8 @@ def test_scan_fix_point():
test_scan5_multi_output()
def test_create_read_graph():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j])
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3)
......
......@@ -2,23 +2,23 @@ import tvm
def test_schedule0():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
s = tvm.Schedule(A1.op)
s = tvm.create_schedule(A1.op)
bounds = tvm.schedule.InferBound(s)
assert isinstance(bounds, tvm.collections.Map)
stmt = tvm.schedule.ScheduleOps(s, bounds)
def test_schedule1():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
s = tvm.Schedule(A1.op)
s = tvm.create_schedule(A1.op)
xo, xi = s[A1].split(A1.op.axis[0], 8)
bounds = tvm.schedule.InferBound(s)
assert isinstance(bounds, tvm.collections.Map)
......@@ -26,13 +26,13 @@ def test_schedule1():
def test_schedule2():
m = tvm.Var('m')
l = tvm.Var('l')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
A1 = tvm.compute((m, l), lambda i, j: A[i, j], name='A1')
A2 = tvm.compute((m, l), lambda i, j: A1[i, j] + 3, name='A2')
s = tvm.Schedule(A2.op)
s = tvm.create_schedule(A2.op)
xo, xi = s[A2].split(A2.op.axis[0], 8)
s[A1].compute_at(s[A2], xo)
bounds = tvm.schedule.InferBound(s)
......@@ -41,8 +41,8 @@ def test_schedule2():
def test_schedule_scan():
m = tvm.Var("m")
n = tvm.Var("n")
m = tvm.var("m")
n = tvm.var("n")
x = tvm.compute((m, n), lambda i, j: tvm.const(1, "float32"), name="x")
s_state = tvm.placeholder((m, n))
s_init = tvm.compute((1, n), lambda _, i: x[0, i])
......@@ -50,7 +50,7 @@ def test_schedule_scan():
res = tvm.scan(s_init, s_update, s_state)
assert tuple(res.shape) == (m, n)
s = tvm.Schedule(res.op)
s = tvm.create_schedule(res.op)
s.normalize()
bounds = tvm.schedule.InferBound(s)
assert(bounds[res.op.scan_axis].min.value == 1)
......@@ -58,28 +58,28 @@ def test_schedule_scan():
print(stmt)
def test_auto_inline():
m = tvm.Var('m')
n = tvm.Var('n')
m = tvm.var('m')
n = tvm.var('n')
A = tvm.placeholder((m, n), name='A')
B = tvm.placeholder((m, n), name='B')
C = tvm.placeholder((m, n), name='C')
T1 = tvm.compute((m, n), lambda i, j: A(i, j) * B(i, j), name='T1')
T2 = tvm.compute((m, n), lambda i, j: T1(i, j) + C(i, j), name='T2')
s = tvm.Schedule(T2.op)
s = tvm.create_schedule(T2.op)
tvm.schedule.AutoInlineElemWise(s)
s.normalize()
bounds = tvm.schedule.InferBound(s)
stmt = tvm.schedule.ScheduleOps(s, bounds)
def test_inline_mixed():
n = tvm.Var('n')
n = tvm.var('n')
A = tvm.placeholder((n, ), name='A')
A1 = tvm.compute(A.shape, lambda *i: A(*i) + 1, name='A1')
A2 = tvm.compute(A.shape, lambda *i: A1(*i) + 2, name='A2')
C = tvm.compute((n,), lambda i: A2[i] + A1[i], name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(C.op.axis[0], factor=8)
s[A1].compute_at(s[C], xo)
s[A2].compute_inline()
......@@ -90,13 +90,13 @@ def test_inline_mixed():
def test_schedule_cache():
m = tvm.Var('m')
n = tvm.Var('n')
m = tvm.var('m')
n = tvm.var('n')
A = tvm.placeholder((m, n), name='A')
B = tvm.placeholder((m, n), name='B')
C = tvm.compute((m, n), lambda i, j: A(i, j) * B(i, j), name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
AA = s.cache_read(A, "shared", readers=[C])
CC = s.cache_write(C, "shared")
s[AA].compute_at(s[CC], CC.op.axis[0])
......
......@@ -8,7 +8,7 @@ def lower(s, args, name):
for x in args:
assert isinstance(x, tvm.tensor.Tensor)
buf = tvm.Buffer(x.shape, dtype=x.dtype, name=x.op.name)
buf = tvm.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
binds[x] = buf
arg_list.append(buf)
s.normalize()
......@@ -33,7 +33,7 @@ def test_add_pipeline():
A = tvm.placeholder((n,), name='A', dtype='int32')
B = tvm.placeholder((n,), name='B', dtype='int32')
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name='C')
s = tvm.Schedule(C.op)
s = tvm.create_schedule(C.op)
px, x = s[C].split(C.op.axis[0], nparts=1)
s[C].bind(px, tvm.thread_axis("pipeline"))
......
.. _tutorials-index:
TVM Tutorials
=============
.. _notebook_examples:
"""
Get Started with TVM
====================
**Author**: `Tianqi Chen <https://tqchen.github.io>`_
This is an introduction tutorial to TVM.
TVM is a domain specifric language for efficient kernel construction.
In this tutorial, we will demonstrate the basic workflow in TVM.
"""
from __future__ import absolute_import, print_function
import tvm
import numpy as np
######################################################################
# Vector Add Example
# ------------------
# In this tutorial, we will use a vector addition example to demonstrate
# the workflow in TVM. We will demonstrate how we can describe and compile
# vector addition code that runs on GPU.
#
######################################################################
# Describe the Computation
# ------------------------
# As a first step, we need to describe our computation.
# TVM adopts tensor semantics, with each intermediate result
# represented as multi-dimensional array. The user need to describe
# the computation rule that generate the tensors.
#
# We first define a symbolic variable n to represent the shape.
# We then define two placeholder Tensors, A and B, with given shape (n,)
#
# We then describe the result tensor C, with a compute operation.
# The compute function takes the shape of the tensor, as well as a lambda function
# that describes the computation rule for each position of the tensor.
#
# No computation happens during this phase, as we are only declaring how
# the computation should be done.
#
n = tvm.var("n")
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")
print(type(C))
######################################################################
# Schedule the Computation
# ------------------------
# While the above lines describes the computation rule, we can compute
# C in many ways since the axis of C can be computed in data parallel manner.
# TVM asks user to provide a description of computation called schedule.
#
# A schedule is a set of transformation of computation that transforms
# the loop of computations in the program.
#
# After we construct the schedule, by default the schedule computes
# C in a serial manner in a row-major order.
#
# .. code-block:: c
#
# for (int i = 0; i < n; ++i) {
# C[i] = A[i] + B[i];
# }
#
s = tvm.create_schedule(C.op)
######################################################################
# We used the split construct to split the first axis of C,
# this will split the original iteration axis into product of
# two iterations. This is equivalent to the following code.
#
# .. code-block:: c
#
# for (int bx = 0; bx < ceil(n / 64); ++bx) {
# for (int tx = 0; tx < 64; ++tx) {
# int i = bx * 64 + tx;
# if (i < n) {
# C[i] = A[i] + B[i];
# }
# }
# }
#
bx, tx = s[C].split(C.op.axis[0], factor=64)
######################################################################
# Finally we bind the iteratio axis bx and tx to threads in the GPU
# compute grid. These are GPU specific constructs that allows us
# to generate code that runs on GPU.
#
s[C].bind(bx, tvm.thread_axis("blockIdx.x"))
s[C].bind(tx, tvm.thread_axis("threadIdx.x"))
######################################################################
# Compilation
# -----------
# After we have finished specifying the schedule, we can compile it
# into a TVM function. By default TVM compiles into a type-erased
# function that can be directly called from python side.
#
# In the following line, we use tvm.build to create a function.
# The build function takes the schedule, the desired signature of the
# function(including the inputs and outputs) as well as target language
# we want to compile to.
#
# The result of compilation fadd is a CUDA device function that can
# as well as a host wrapper that calls into the CUDA function.
# fadd is the generated host wrapper function, it contains reference
# to the generated device function internally.
#
fadd_cuda = tvm.build(s, [A, B, C], "cuda", target_host="llvm", name="myadd")
######################################################################
# Run the Function
# ----------------
# The compiled function TVM function is designed to be a concise C API
# that can be invoked from any languages.
#
# We provide an minimum array API in python to aid quick testing and prototyping.
# The array API is based on `DLPack <https://github.com/dmlc/dlpack>`_ standard.
#
# - We first create a gpu context.
# - Then tvm.nd.array copies the data to cpu.
# - fadd runs the actual computation.
# - asnumpy() copies the gpu array back to cpu and we can use this to verify correctness
#
ctx = tvm.gpu(0)
n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd_cuda(a, b, c)
np.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
######################################################################
# Inspect the Generated Code
# --------------------------
# You can inspect the generated code in TVM. The result of tvm.build
# is a tvm Module. fadd is the host module that contains the host wrapper,
# it also contains a device module for the CUDA function.
#
# The following code fetchs the device module and prints the content code.
#
dev_module = fadd_cuda.imported_modules[0]
print("-----CUDA code-----")
print(dev_module.get_source())
######################################################################
# .. note:: Code Specialization
#
# As you may noticed, during the declaration, A, B and C both
# takes the same shape argument n. TVM will take advantage of this
# to pass only single shape argument to the kernel, as you will find in
# the printed device code. This is one form of specialization.
#
# On the host side, TVM will automatically generate check codes
# that checks the constraints in the parameters. So if you pass
# arrays with different shape into the fadd, an error will be raised.
#
# We can do more specializations. For example, we can write
# :code:`n = tvm.convert(1024)` instead of :code:`n = tvm.var("n")`,
# in the computation declaration. The generated function will
# only take vectors with length 1024.
#
######################################################################
# Save Compiled Module
# --------------------
# Besides runtime compilation, we can save the compiled module into
# file and load them back later. This is called ahead of time compilation.
#
# The following code first does the following step:
#
# - It saves the compiled host module into an object file.
# - Then it saves the device module into a ptx file.
# - cc.create_shared calls a env compiler(gcc) to create a shared library
#
from tvm.addon import cc_compiler as cc
from tvm.addon import testing
temp = testing.tempdir()
fadd_cuda.save(temp.relpath("myadd.o"))
fadd_cuda.imported_modules[0].save(temp.relpath("myadd.ptx"))
cc.create_shared(temp.relpath("myadd.so"), [temp.relpath("myadd.o")])
print(temp.listdir())
######################################################################
# .. note:: Module Storage Format
#
# The CPU(host) module is directly saved as a shared library(so).
# There can be multiple customed format on the device code.
# In our example, device code is stored in ptx, as well as a meta
# data json file. In the future we can consider pack every binary
# into one shared library.
#
######################################################################
# Load Compiled Module
# --------------------
# We can load the compiled module from the file system and run the code.
# The following code load the host and device module seperatedly and
# re-link them together. We can verify that the newly loaded function works.
#
fadd1 = tvm.module.load(temp.relpath("myadd.so"))
fadd1_dev = tvm.module.load(temp.relpath("myadd.ptx"))
fadd1.import_module(fadd1_dev)
fadd1(a, b, c)
np.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
######################################################################
# .. note:: Runtime API and Thread-Safety
#
# The compiled modules of TVM do not depend on the TVM compiler.
# Instead, it only depends on a minimum runtime library.
# TVM runtime library wraps the device drivers and provide
# thread-safe and device agnostic call into the compiled functions.
#
# This means you can call the compiled TVM function from any thread,
# on any GPUs.
#
######################################################################
# Generate OpenCL Code
# --------------------
# TVM provides code generation features into multiple backends,
# we can also generate OpenCL code or LLVM code that runs on CPU backends.
#
# The following codeblocks generate opencl code, creates array on opencl
# device, and verifies the correctness of the code.
#
tvm.module.init_opencl()
fadd_cl = tvm.build(s, [A, B, C], "opencl", name="myadd")
print("------opencl code------")
print(fadd_cl.imported_modules[0].get_source())
ctx = tvm.cl(0)
n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd_cl(a, b, c)
np.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
######################################################################
# Summary
# -------
# This tutorial provides a walk through of TVM workflow using
# a vector add example. The general workflow is
#
# - Describe your computation via series of operations.
# - Describe how we want to compute use schedule primitives.
# - Compile to the target function we want.
# - Optionally, save the function to be loaded later.
#
# You are more than welcomed to checkout other examples and
# tutorials to learn more about the supported operations, schedule primitives
# and other features in TVM.
#
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment