<jupyter_start><jupyter_text>Schedule Primitives in TVM==========================**Author**: `<NAME> `_TVM is a domain specific language for efficient kernel construction.In this tutorial, we will show you how to schedule the computation byvarious primitives provided by TVM.<jupyter_code>from __future__ import absolute_import, print_function

import tvm
from tvm import te
import numpy as np<jupyter_output><empty_output><jupyter_text>There often exist several methods to compute the same result,however, different methods will result in different locality andperformance. So TVM asks user to provide how to execute thecomputation called **Schedule**.A **Schedule** is a set of transformation of computation thattransforms the loop of computations in the program.<jupyter_code># declare some variables for use later
n = te.var("n")
m = te.var("m")<jupyter_output><empty_output><jupyter_text>A schedule can be created from a list of ops, by default theschedule computes tensor in a serial manner in a row-major order.<jupyter_code># declare a matrix element-wise multiply
A = te.placeholder((m, n), name="A")
B = te.placeholder((m, n), name="B")
C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name="C")

s = te.create_schedule([C.op])
# lower will transform the computation from definition to the real
# callable function. With argument `simple_mode=True`, it will
# return you a readable C like statement, we use it here to print the
# schedule result.
print(tvm.lower(s, [A, B, C], simple_mode=True))<jupyter_output><empty_output><jupyter_text>One schedule is composed by multiple stages, and one**Stage** represents schedule for one operation. We provide variousmethods to schedule every stage. split-----:code:`split` can split a specified axis into two axes by:code:`factor`.<jupyter_code>A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] * 2, name="B")

s = te.create_schedule(B.op)
xo, xi = s[B].split(B.op.axis[0], factor=32)
print(tvm.lower(s, [A, B], simple_mode=True))<jupyter_output><empty_output><jupyter_text>You can also split a axis by :code:`nparts`, which splits the axiscontrary with :code:`factor`.<jupyter_code>A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i], name="B")

s = te.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], nparts=32)
print(tvm.lower(s, [A, B], simple_mode=True))<jupyter_output><empty_output><jupyter_text>tile----:code:`tile` help you execute the computation tile by tile over twoaxes.<jupyter_code>A = te.placeholder((m, n), name="A")
B = te.compute((m, n), lambda i, j: A[i, j], name="B")

s = te.create_schedule(B.op)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
print(tvm.lower(s, [A, B], simple_mode=True))<jupyter_output><empty_output><jupyter_text>fuse----:code:`fuse` can fuse two consecutive axes of one computation.<jupyter_code>A = te.placeholder((m, n), name="A")
B = te.compute((m, n), lambda i, j: A[i, j], name="B")

s = te.create_schedule(B.op)
# tile to four axes first: (i.outer, j.outer, i.inner, j.inner)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
# then fuse (i.inner, j.inner) into one axis: (i.inner.j.inner.fused)
fused = s[B].fuse(xi, yi)
print(tvm.lower(s, [A, B], simple_mode=True))<jupyter_output><empty_output><jupyter_text>reorder-------:code:`reorder` can reorder the axes in the specified order.<jupyter_code>A = te.placeholder((m, n), name="A")
B = te.compute((m, n), lambda i, j: A[i, j], name="B")

s = te.create_schedule(B.op)
# tile to four axes first: (i.outer, j.outer, i.inner, j.inner)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
# then reorder the axes: (i.inner, j.outer, i.outer, j.inner)
s[B].reorder(xi, yo, xo, yi)
print(tvm.lower(s, [A, B], simple_mode=True))<jupyter_output><empty_output><jupyter_text>bind----:code:`bind` can bind a specified axis with a thread axis, often usedin gpu programming.<jupyter_code>A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: A[i] * 2, name="B")

s = te.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], factor=64)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
print(tvm.lower(s, [A, B], simple_mode=True))<jupyter_output><empty_output><jupyter_text>compute_at----------For a schedule that consists of multiple operators, TVM will computetensors at the root separately by default.<jupyter_code>A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")

s = te.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))<jupyter_output><empty_output><jupyter_text>:code:`compute_at` can move computation of `B` into the first axisof computation of `C`.<jupyter_code>A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")

s = te.create_schedule(C.op)
s[B].compute_at(s[C], C.op.axis[0])
print(tvm.lower(s, [A, B, C], simple_mode=True))<jupyter_output><empty_output><jupyter_text>compute_inline--------------:code:`compute_inline` can mark one stage as inline, then the body ofcomputation will be expanded and inserted at the address where thetensor is required.<jupyter_code>A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")

s = te.create_schedule(C.op)
s[B].compute_inline()
print(tvm.lower(s, [A, B, C], simple_mode=True))<jupyter_output><empty_output><jupyter_text>compute_root------------:code:`compute_root` can move computation of one stage to the root.<jupyter_code>A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")

s = te.create_schedule(C.op)
s[B].compute_at(s[C], C.op.axis[0])
s[B].compute_root()
print(tvm.lower(s, [A, B, C], simple_mode=True))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18342504585626146, 'fraction_numerical': 0.006003001500750375, 'mean_word_length': 5.777401129943502, 'pattern_counts': {'":': 0, '<': 54, '<?xml version=': 0, '>': 54, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '368893', 'license': 'Apache-2.0', 'n_tokens_mistral': 2284, 'n_tokens_neox': 2243, 'n_words': 604, 'path': '_downloads/b78f1a6e1b2c2fb073a791dc258a1d7d/schedule_primitives.ipynb', 'repo_name': 'TVMChinese/TVMChinese.github.io', 'size': '6003'}