TensorIR: 张量程序抽象案例研究
------------------------------
安装相关的包
~~~~~~~~~~~~
为了本课程的目标，我们会使用 TVM
（一个开源的机器学习编译框架）中一些正在持续开发的部分。我们提供了下面的命令用于为
MLC 课程安装一个包装好的版本。
.. raw:: latex
\diilbookstyleinputcell
.. code:: bash
python3 -m pip install mlc-ai-nightly -f https://mlc.ai/wheels
序言
~~~~
.. figure:: ../img/tensor_func_linear_relu.png
在开始今天这一节之前，让我们回忆机器学习编译过程的关键原则。大多数的机器学习编译可以被视为张量函数之间的变换。在接下来的内容中我们主要想回答的以下几个问题：
- 什么是表示张量函数可能的抽象？
- 什么是张量函数之间可能的变换？
在这里我们将聚焦于元张量函数，部分讨论这些问题。
学习一类张量程序抽象 – TensorIR
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
在之前的章节中，我们介绍过了元张量函数，讨论了张量程序抽象的总体概念。
现在我们已经准备好学习一个特定的张量程序抽象的实例：TensorIR。 TensorIR
是标准机器学习编译框架 Apache TVM 中使用的张量程序抽象。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
import numpy as np
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
使用张量程序抽象的主要目的是表示循环和相关的硬件加速选择，如多线程、特殊硬件指令的使用和内存访问。
为了帮助我们更好地解释，我们用下面的张量计算作为示例。
具体地，对于两个大小为 :math:`128 \times 128` 的矩阵 A 和
B，我们进行如下两步的张量计算。
- :math:`Y_{i, j} = \sum_k A_{i, k} \times B_{k, j}`
- :math:`C_{i, j} = \mathbb{relu}(Y_{i, j}) = \mathbb{max}(Y_{i, j}, 0)`
上面的计算很像在我们神经网络中经常看到的典型的元张量函数：一个线性层与一个
ReLU 激活层。首先，我们使用如下 NumPy 中的数组计算实现这两个操作。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
dtype = "float32"
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
# a @ b is equivalent to np.matmul(a, b)
c_mm_relu = np.maximum(a_np @ b_np, 0)
在底层，NumPy 调用库（例如 OpenBLAS）和它自己在低级 C
语言中的一些实现来执行这些计算。
从张量程序抽象的角度来看，我们想彻底理解这些数组计算的\ **背后**\ 的细节。具体来说，我们想问：实现相应计算的可能方式是什么？
为了说明底层细节，我们将在 NumPy API 的一个受限子集中编写示例 ——
我们称之为 **低级 NumPy**\ 。它使用以下的约定：
- 我们将在必要时使用循环而不是数组函数来展示可能的循环计算。
- 如果可能，我们总是通过 ``numpy.empty`` 显式地分配数组并传递它们。
需要注意的是，这不是人们通常编写 NumPy
程序的方式。不过，它们仍然与幕后发生的事情非常相似 ——
大多数现实世界的部署解决方案都将分配与计算分开处理。特定的库使用不同形式的循环和算术计算来执行计算。当然首先它们是使用诸如
``C`` 之类的低级语言实现的。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
Y = np.empty((128, 128), dtype="float32")
for i in range(128):
for j in range(128):
for k in range(128):
if k == 0:
Y[i, j] = 0
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
for i in range(128):
for j in range(128):
C[i, j] = max(Y[i, j], 0)
上面的程序是实现 ``mm_relu``
操作的一种方式。该程序包含两个阶段：首先我们分配一个中间存储
:math:`Y`\ ，将矩阵乘法的结果存储在那里。然后我们在第二个 for
循环序列中计算 ReLU。你可能会注意到，这肯定不是实现 ``mm_relu``
的唯一方法，当然这可能也不是你想到的第一件事。
无论如何，这确实是实现 ``mm_relu``
的方式之一。我们可以通过将我们的结果与使用数组计算的原始结果进行比较来验证代码的正确性。我们将在本节的后面部分回到这里，并重新讨论其他可能的方法。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)
上面的示例代码展示了我们如何在\ **幕后**\ 实现 ``mm_relu``\ 。当然，由于
Python 解释器，代码本身会运行得很慢。尽管如此，示例 NumPy
代码包含我们将在这些计算的实际实现中使用的所有可能元素。
- 多维缓冲区（数组）。
- 在数组维度上的循环。
- 在循环下执行的计算语句。
在看过低级 NumPy 示例后，现在我们准备介绍 TensorIR。下面的代码块展示了
``mm_relu`` 的 TensorIR 实现。这里的代码是用一种名为 TVMScript
的语言实现的，它是一种嵌入在 Python AST 中的特定领域方言。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
@tvm.script.ir_module
class MyModule:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
Y = T.alloc_buffer((128, 128), dtype="float32")
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
能够同时查看 NumPy 代码和 TensorIR
代码并检查相应的元素是很有帮助的。我们将详细介绍它们。
.. figure:: ../img/tensor_func_and_numpy.png
我们首先回顾一下在 NumPy 和 TensorIR
之间具有直接对应关系的元素。接下来我们将回过来查看不属于 NumPy
程序的其他元素。
函数参数与缓冲区
^^^^^^^^^^^^^^^^
首先，让我们看一下函数参数。函数参数对应于 NumPy 函数上的同一组参数。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# TensorIR
def mm_relu(A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")):
...
# numpy
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
...
这里 A、B 和 C 采用名为 ``T.Buffer`` 的类型，其形状参数为
``(128, 128)``\ ，数据类型为
``float32``\ 。这些附加信息有助于可能的机器学习编译过程以生成专门针对形状和数据类型的代码。
同样，TensorIR 在中间结果分配中也使用了缓冲区类型。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# TensorIR
Y = T.alloc_buffer((128, 128), dtype="float32")
# numpy
Y = np.empty((128, 128), dtype="float32")
For 循环迭代
^^^^^^^^^^^^
循环迭代也有直接对应关系。\ ``T.grid`` 是 TensorIR
中的语法糖，供我们书写多个嵌套的迭代器。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# TensorIR
for i, j, k in T.grid(128, 128, 128):
# numpy
for i in range(128):
for j in range(128):
for k in range(128):
计算块
^^^^^^
主要区别之一来自计算语句。TensorIR 包含一个名为 ``T.block`` 的额外结构。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# TensorIR
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
# coressponding numpy code
vi, vj, vk = i, j, k
if vk == 0:
Y[vi, vj] = 0
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
**块** 是 TensorIR 中的基本计算单位。值得注意的是，该块包含比普通 NumPy
代码更多的信息。一个块包含一组块轴（\ ``vi、vj、vk``\ ）和围绕它们定义的计算。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
上面三行声明了关于块轴的\ **关键性质**\ ，语法如下。
::
[block_axis] = T.axis.[axis_type]([axis_range], [mapped_value])
这三行包含以下信息：
- 定义了 ``vi``\ 、\ ``vj``\ 、\ ``vk`` 应被绑定到的位置（在本例中为
``i``\ 、\ ``j`` 和 ``k``\ ）；
- 声明了 ``vi``\ 、\ ``vj``\ 、\ ``vk``
的原始范围（\ ``T.axis.spatial(128, i)`` 中的 ``128``\ ）；
- 声明了块轴的属性（\ ``spatial``, ``reduce``\ ）。
我们一一浏览这些属性。首先，就边界关系而言，\ ``vi = T.axis.spatial(128, i)``
有效地蕴含了 ``vi = i``\ 。\ ``[axis_range]`` 值提供了 ``[block_axis]``
的预期范围。例如，\ ``vi = T.axis.spatial(128, i)`` 中的 ``128`` 表明
``vi`` 应该在 ``range(0, 128)`` 中。
块轴的属性
^^^^^^^^^^
我们现在开始进一步研究块轴属性。这些轴属性标记了轴与正在执行的计算之间的关系。
下图总结了块（迭代）轴和块 ``Y``
的读写关系。注意到，严格来说这个块正在对缓冲区 ``Y``
进行（规约）更新，我们暂时将其标记为\ *写入*\ ，因为我们不需要来自另一个块的缓冲区
``Y`` 的值。
.. figure:: ../img/tensor_ir_block_axis.png
在我们的示例中，块 ``Y`` 通过读取来自 ``A[vi, vk]`` 和 ``B[vk, vj]``
的值来计算结果 ``Y[vi, vj]``\ ，并对所有可能的 ``vk`` 执行求和。
在这个特定示例中，如果我们将 ``vi``\ 、\ ``vj`` 固定为
``(0, 1)``\ ，并对 ``vk in range(0, 128)`` 执行块
``Y``\ ，我们可以独立于其他可能的位置（具有不同 ``vi``, ``vj``
值的位置）有效地计算 ``C[0, 1]``\ 。
值得注意的是，对于一组固定的 ``vi`` 和 ``vj``\ ，计算块在 ``Y``
的空间位置 (``Y[vi, vj]``) 处生成一个点值，该点值独立于 ``Y``
中的其他位置（具有不同的\ ``vi``, ``vj`` 值的位置）。我们可以称
``vi``\ 、\ ``vj``
为\ **空间轴**\ ，因为它们直接对应于块写入的缓冲区空间区域的开始。
涉及归约的轴（\ ``vk``\ ）被命名为\ **归约轴**\ 。
为什么块需要额外附加的信息
^^^^^^^^^^^^^^^^^^^^^^^^^^
一个重要的观察结果是，附加信息（块轴范围及其属性）使块轴\ **独立**\ 于外部循环嵌套
``i``, ``j``, ``k``\ 。
块轴信息还提供了额外的属性，帮助我们验证用于执行计算的外部循环的正确性。例如，下面的代码块会导致错误，因为循环需要一个大小为
128 的迭代器，但我们只将它绑定到一个大小为 127 的 for 循环。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# wrong program due to loop and block iteration mismatch
for i in range(127):
with T.block("C"):
vi = T.axis.spatial(128, i)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
error here due to iterator size mismatch
...
这些附加信息也有助于我们进行机器学习编译分析。例如，虽然我们总是可以在空间轴上做并行化，在规约轴上进行并行化将需要特定的策略。
块轴绑定的语法糖
^^^^^^^^^^^^^^^^
在每个块轴直接映射到外部循环迭代器的情况下，我们可以使用
``T.axis.remap`` 在一行中声明所有块轴。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# SSR means the properties of each axes are "spatial", "spatial", "reduce"
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
相当于
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
vi = T.axis.spatial(range_of_i, i)
vj = T.axis.spatial(range_of_j, j)
vk = T.axis.reduce(range_of_k, k)
所以我们也可以编写如下程序。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
@tvm.script.ir_module
class MyModuleWithAxisRemapSugar:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
Y = T.alloc_buffer((128, 128), dtype="float32")
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
函数属性和装饰器
^^^^^^^^^^^^^^^^
到目前为止，我们已经涵盖了 TensorIR
中的大部分元素。在这一部分中，我们将介绍 TVMScript 的其余元素。
函数属性信息包含关于函数的额外信息。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
这里的 ``global_symbol`` 对应函数名，\ ``tir.noalias``
是一个属性，表示所有的缓冲存储器不重叠。你现在可以放心地跳过这些属性，因为它们不会影响对概念的整体理解。
``@tvm.script.ir_module`` 和 ``@T.prim_func``
这两个装饰器用于表示对应部分的类型。
``@tvm.script.ir_module`` 表示 ``MyModule`` 是一个 IRModule。IRModule
是在机器学习编译中保存张量函数集合的容器对象。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
type(MyModule)
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
tvm.ir.module.IRModule
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
type(MyModule["mm_relu"])
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
tvm.tir.function.PrimFunc
到目前为止，我们只看到包含单个张量函数的 IRModule。
机器学习编译过程中的一个 IRModule 可以包含多个张量函数。
以下代码块显示了具有两个函数的 IRModule 示例。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
@tvm.script.ir_module
class MyModuleWithTwoFunctions:
@T.prim_func
def mm(A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
Y: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm", "tir.noalias": True})
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
@T.prim_func
def relu(A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "relu", "tir.noalias": True})
for i, j in T.grid(128, 128):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.max(A[vi, vj], T.float32(0))
章节检查点
^^^^^^^^^^
到目前为止，我们一同看过了一个 TensorIR
程序示例，并涵盖了大部分元素，包括：
- 参数和中间临时内存中的缓冲区声明；
- For 循环迭代；
- **块**\ 和块轴属性。
在本节中，我们介绍了一个 TensorIR
示例，它涵盖了机器学习编译中最常见的元素。
TensorIR
包含的元素比我们在本节中介绍的更多，但本节涵盖了可以让我们开始机器学习编译之旅的大部分关键部分。我们将在后面的章节中介绍遇到的新元素。
变换
~~~~
在上一节中，我们了解了 TensorIR 及其关键元素。
现在，让我们了解所有机器学习编译工作流的主要成分——元张量函数的变换。
在上一节中，我们给出了如何使用低级 NumPy 编写 ``mm_relu`` 的示例。
在实践中，可以有多种方法来实现相同的功能，并且每种实现都可能导致不同的性能。
我们将讨论性能背后的原因以及如何在以后的内容中利用这些变体。
在本节中，让我们关注使用变换获得不同实现变体的能力。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
def lnumpy_mm_relu_v2(A: np.ndarray, B: np.ndarray, C: np.ndarray):
Y = np.empty((128, 128), dtype="float32")
for i in range(128):
for j0 in range(32):
for k in range(128):
for j1 in range(4):
j = j0 * 4 + j1
if k == 0:
Y[i, j] = 0
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
for i in range(128):
for j in range(128):
C[i, j] = max(Y[i, j], 0)
c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu_v2(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)
上面的代码块显示了 ``mm_relu``
的一个稍微不同的变体。它与原始程序的不同是
- 我们用两个循环 ``j0`` 和 ``j1`` 替换了 ``j`` 循环；
- 迭代顺序略有变化。
为了获得
``lnumpy_mm_relu_v2``\ ，我们必须重写一个新函数（或手动复制粘贴和编辑）。TensorIR
引入了一个名为 Schedule 的辅助结构，它允许我们务实地做到这一点。
为了提醒我们自己，让我们再看看当前的 ``MyModule`` 内容。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
import IPython
IPython.display.Code(MyModule.script(), language="python")
.. raw:: html
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j_0, j_1, k in T.grid(128, 32, 4, 128):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j_0, k, j_1 in T.grid(128, 32, 128, 4):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j_0 in T.grid(128, 32):
for k, j_1 in T.grid(128, 4):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for ax0 in range(4):
with T.block("C"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + ax0)
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j_0 in T.grid(128, 32):
for k, j_1 in T.grid(128, 4):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for ax0 in range(4):
with T.block("C"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + ax0)
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j_0 in T.grid(128, 16):
for k, j_1 in T.grid(128, 8):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 8 + j_1)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for ax0 in range(8):
with T.block("C"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 8 + ax0)
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
T.reads(A[v_i, v_k], B[v_k, v_j])
T.writes(Y[v_i, v_j])
with T.init():
Y[v_i, v_j] = T.float32(0)
Y[v_i, v_j] = Y[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
for i, j in T.grid(128, 128):
with T.block("C"):
v_i, v_j = T.axis.remap("SS", [i, j])
T.reads(Y[v_i, v_j])
T.writes(C[v_i, v_j])
C[v_i, v_j] = T.max(Y[v_i, v_j], T.float32(0))