wehub-resource-sync
277 lines
12 KiB
ReStructuredText
277 lines
12 KiB
ReStructuredText
.. Licensed to the Apache Software Foundation (ASF) under one
|
|
or more contributor license agreements. See the NOTICE file
|
|
distributed with this work for additional information
|
|
regarding copyright ownership. The ASF licenses this file
|
|
to you under the Apache License, Version 2.0 (the
|
|
"License"); you may not use this file except in compliance
|
|
with the License. You may obtain a copy of the License at
|
|
|
|
.. http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
.. Unless required by applicable law or agreed to in writing,
|
|
software distributed under the License is distributed on an
|
|
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
KIND, either express or implied. See the License for the
|
|
specific language governing permissions and limitations
|
|
under the License.
|
|
|
|
.. _relax-learning:
|
|
|
|
Understand Relax Abstraction
|
|
============================
|
|
Relax is a graph abstraction used in Apache TVM, which
|
|
helps to end-to-end optimize ML models. The principal objective of Relax
|
|
is to depict the structure and data flow of ML models, including the
|
|
dependencies and relationships between different parts of the model, as
|
|
well as how to execute the model on hardware.
|
|
|
|
End to End Model Execution
|
|
--------------------------
|
|
|
|
In this chapter, we will use the following model as an example. This is
|
|
a two-layer neural network that consists of two linear operations with
|
|
relu activation.
|
|
|
|
.. image:: /_static/img/e2e_fashionmnist_mlp_model.png
|
|
:width: 85%
|
|
:align: center
|
|
|
|
|
|
High-Level Operations Representation
|
|
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
Let us begin by reviewing a Numpy implementation of the model.
|
|
|
|
.. code:: python
|
|
|
|
def numpy_mlp(data, w0, b0, w1, b1):
|
|
lv0 = data @ w0 + b0
|
|
lv1 = np.maximum(lv0, 0)
|
|
lv2 = lv1 @ w1 + b1
|
|
return lv2
|
|
|
|
The above example code shows the high-level array operations to perform the end-to-end model
|
|
execution. Of course, we can rewrite the above code using Relax as follows:
|
|
|
|
.. code:: python
|
|
|
|
from tvm.script import relax as R
|
|
|
|
@R.function
|
|
def relax_mlp(
|
|
data: R.Tensor(("n", 784), dtype="float32"),
|
|
w0: R.Tensor((784, 128), dtype="float32"),
|
|
b0: R.Tensor((128,), dtype="float32"),
|
|
w1: R.Tensor((128, 10), dtype="float32"),
|
|
b1: R.Tensor((10,), dtype="float32"),
|
|
) -> R.Tensor(("n", 10), dtype="float32"):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(data, w0) + b0
|
|
lv1 = R.nn.relu(lv0)
|
|
lv2 = R.matmul(lv1, w1) + b1
|
|
R.output(lv2)
|
|
return lv2
|
|
|
|
Low-Level Integration
|
|
~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
However, again from the pov of machine learning compilation (MLC), we would like to see
|
|
through the details under the hood of these array computations.
|
|
|
|
For the purpose of illustrating details under the hood, we will again write examples in low-level numpy:
|
|
|
|
We will use a loop instead of array functions when necessary to demonstrate the possible loop computations.
|
|
When possible, we always explicitly allocate arrays via numpy.empty and pass them around.
|
|
The code block below shows a low-level numpy implementation of the same model.
|
|
|
|
.. code:: python
|
|
|
|
def lnumpy_linear(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray):
|
|
n, m, K = X.shape[0], W.shape[1], X.shape[1]
|
|
Y = np.empty((n, m), dtype="float32")
|
|
for i in range(n):
|
|
for j in range(m):
|
|
for k in range(K):
|
|
if k == 0:
|
|
Y[i, j] = 0
|
|
Y[i, j] = Y[i, j] + X[i, k] * W[k, j]
|
|
|
|
for i in range(n):
|
|
for j in range(m):
|
|
Z[i, j] = Y[i, j] + B[j]
|
|
|
|
|
|
def lnumpy_relu0(X: np.ndarray, Y: np.ndarray):
|
|
n, m = X.shape
|
|
for i in range(n):
|
|
for j in range(m):
|
|
Y[i, j] = np.maximum(X[i, j], 0)
|
|
|
|
def lnumpy_mlp(data, w0, b0, w1, b1):
|
|
n = data.shape[0]
|
|
lv0 = np.empty((n, 128), dtype="float32")
|
|
lnumpy_linear(data, w0, b0, lv0)
|
|
|
|
lv1 = np.empty((n, 128), dtype="float32")
|
|
lnumpy_relu0(lv0, lv1)
|
|
|
|
out = np.empty((n, 10), dtype="float32")
|
|
lnumpy_linear(lv1, w1, b1, out)
|
|
return out
|
|
|
|
With the low-level NumPy example in mind, now we are ready to introduce an Relax abstraction
|
|
for the end-to-end model execution. The code block below shows a TVMScript implementation of the model.
|
|
|
|
.. code:: python
|
|
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.script import relax as R
|
|
|
|
@I.ir_module
|
|
class Module:
|
|
@T.prim_func(private=True)
|
|
def linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
|
|
M, N, K = T.int64(), T.int64(), T.int64()
|
|
X = T.match_buffer(x, (M, K), "float32")
|
|
W = T.match_buffer(w, (K, N), "float32")
|
|
B = T.match_buffer(b, (N,), "float32")
|
|
Z = T.match_buffer(z, (M, N), "float32")
|
|
Y = T.alloc_buffer((M, N), "float32")
|
|
for i, j, k in T.grid(M, N, K):
|
|
with T.sblock("Y"):
|
|
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
|
|
with T.init():
|
|
Y[v_i, v_j] = T.float32(0.0)
|
|
Y[v_i, v_j] = Y[v_i, v_j] + X[v_i, v_k] * W[v_k, v_j]
|
|
for i, j in T.grid(M, N):
|
|
with T.sblock("Z"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
Z[v_i, v_j] = Y[v_i, v_j] + B[v_j]
|
|
|
|
@T.prim_func(private=True)
|
|
def relu(x: T.handle, y: T.handle):
|
|
M, N = T.int64(), T.int64()
|
|
X = T.match_buffer(x, (M, N), "float32")
|
|
Y = T.match_buffer(y, (M, N), "float32")
|
|
for i, j in T.grid(M, N):
|
|
with T.sblock("Y"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
Y[v_i, v_j] = T.max(X[v_i, v_j], T.float32(0.0))
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("n", 784), dtype="float32"),
|
|
w0: R.Tensor((784, 256), dtype="float32"),
|
|
b0: R.Tensor((256,), dtype="float32"),
|
|
w1: R.Tensor((256, 10), dtype="float32"),
|
|
b1: R.Tensor((10,), dtype="float32")
|
|
) -> R.Tensor(("n", 10), dtype="float32"):
|
|
cls = Module
|
|
n = T.int64()
|
|
with R.dataflow():
|
|
lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
|
|
lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
|
|
lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
|
|
R.output(lv2)
|
|
return lv2
|
|
|
|
The above code contains kinds of functions: the primitive tensor functions (``T.prim_func``) and a
|
|
``R.function`` (relax function). Relax function is a new type of abstraction representing
|
|
high-level neural network executions.
|
|
|
|
Note that the above relax module natively supports symbolic shapes, see the ``"n"`` in the
|
|
tensor shapes in ``main`` function and ``M``, ``N``, ``K`` in the ``linear`` function. This is
|
|
a key feature of Relax abstraction, which enables the compiler to track dynamic shape relations
|
|
globally across tensor operators and function calls.
|
|
|
|
Again it is helpful to see the TVMScript code and low-level numpy code side-by-side and check the
|
|
corresponding elements, and we are going to walk through each of them in detail. Since we already
|
|
learned about primitive tensor functions, we are going to focus on the high-level execution part.
|
|
|
|
Key Elements of Relax
|
|
---------------------
|
|
This section will introduce the key elements of Relax abstraction and how it enables optimization
|
|
in ML compilers.
|
|
|
|
Type
|
|
~~~~
|
|
Type is the Relax representation of expression type information. It can
|
|
be ``TensorType``, ``TupleType``, etc. In the above example, we use ``TensorType``
|
|
(short in ``R.Tensor`` in TVMScript) to represent the shape and dtype of the tensor of the inputs,
|
|
outputs, and intermediate results.
|
|
|
|
R.call_tir
|
|
~~~~~~~~~~
|
|
The ``R.call_tir`` function is a new abstraction in Relax that allows calling primitive tensor
|
|
functions in the same IRModule. This is a key feature of Relax that enables cross-level
|
|
abstractions, from high-level neural network layers to low-level tensor operations.
|
|
Taking one line from the above code as an example:
|
|
|
|
.. code:: python
|
|
|
|
lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
|
|
|
|
To explain what does ``R.call_tir`` work, let us review an equivalent low-level numpy
|
|
implementation of the operation, as follows:
|
|
|
|
.. code:: python
|
|
|
|
lv0 = np.empty((n, 256), dtype="float32")
|
|
lnumpy_linear(x, w0, b0, lv0)
|
|
|
|
Specifically, ``call_tir`` allocates an output tensor res, then pass the inputs and the output
|
|
to the prim_func. After executing prim_func the result is populated in res, then we can return
|
|
the result.
|
|
|
|
This convention is called **destination passing**, The idea is that input and output are explicitly
|
|
allocated outside and passed to the low-level primitive function. This style is commonly used
|
|
in low-level library designs, so higher-level frameworks can handle that memory allocation
|
|
decision. Note that not all tensor operations can be presented in this style (specifically,
|
|
there are operations whose output shape depends on the input). Nevertheless, in common practice,
|
|
it is usually helpful to write the low-level function in this style when possible.
|
|
|
|
Dataflow Block
|
|
~~~~~~~~~~~~~~
|
|
Another important element in a relax function is the R.dataflow() scope annotation.
|
|
|
|
.. code:: python
|
|
|
|
with R.dataflow():
|
|
lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
|
|
lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
|
|
lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
|
|
R.output(lv2)
|
|
|
|
Before we talk about the dataflow block, let us first introduce the concept of **pure** and
|
|
**side-effect**. A function is **pure** or **side-effect free** if:
|
|
|
|
- it only reads from its inputs and returns the result via its output
|
|
- it will not change other parts of the program (such as incrementing a global counter).
|
|
|
|
For example, all ``R.call_tir`` functions are pure functions, as they only read from their inputs
|
|
and write the output to another new allocated tensor. However, the **inplace operations** are not
|
|
pure functions, in other words, they are side-effect functions, because they will change the existing
|
|
intermediate or input tensors.
|
|
|
|
A dataflow block is a way for us to mark the computational graph regions of the program.
|
|
Specifically, within a dataflow block, all the operations need to be **side-effect free**.
|
|
Outside a dataflow block, the operations can contain side-effect.
|
|
|
|
.. note::
|
|
|
|
A common question that arises is why we need to manually mark dataflow blocks instead of
|
|
automatically inferring them. There are two main reasons for this approach:
|
|
|
|
- Automatic inference of dataflow blocks can be challenging and imprecise, particularly
|
|
when dealing with calls to packed functions (such as cuBLAS integrations). By manually
|
|
marking dataflow blocks, we enable the compiler to accurately understand and optimize
|
|
the program's dataflow.
|
|
- Many optimizations can only be applied within dataflow blocks. For instance, fusion
|
|
optimization is limited to operations within a single dataflow block. If the compiler
|
|
were to incorrectly infer dataflow boundaries, it might miss crucial optimization
|
|
opportunities, potentially impacting the program's performance.
|
|
|
|
By allowing manual marking of dataflow blocks, we ensure that the compiler has the most
|
|
accurate information to work with, leading to more effective optimizations.
|