自动程序优化
============
前言
----
在过去的章节中，我们学习了如何构建元张量函数并将它们连接起来以进行端到端的模型执行。到目前为止，我们使用了三种主要的抽象类型。
- 驱动高层执行的计算图抽象
- 元张量函数的抽象
- 通过注册环境函数从而能被调用的库函数
所有这些元素都封装在一个 IRModule 中。大多数 MLC
过程可以看作是元张量函数之间的变换。
有许多不同的方法可以变换同一个程序。本章将讨论自动化一些流程的方法。
准备工作
--------
首先，我们将导入必要的依赖项并创建辅助函数。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
import numpy as np
import tvm
from tvm import relax
from tvm.ir.module import IRModule
from tvm.script import relax as R
from tvm.script import tir as T
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
import IPython
def code2html(code):
"""Helper function to use pygments to turn the code string into highlighted html."""
import pygments
from pygments.formatters import HtmlFormatter
from pygments.lexers import Python3Lexer
formatter = HtmlFormatter()
html = pygments.highlight(code, Python3Lexer(), formatter)
return "%s\n" % (formatter.get_style_defs(".highlight"), html)
回顾：变换元张量函数
--------------------
让我们首先回顾一下我们在前几章中所做的事情——变换单个元张量函数。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
@tvm.script.ir_module
class MyModule:
@T.prim_func
def main(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
for i, j, k in T.grid(128, 128, 128):
with T.block("C"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
C[vi, vj] = 0.0
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
首先，让我们定义一组用于测试评估的输入输出。
.. 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)
c_mm = a_np @ b_np
我们可以如下构建和运行 ``MyModule``\ 。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty((128, 128), dtype="float32")
lib = tvm.build(MyModule, target="llvm")
f_timer_before = lib.time_evaluator("main", tvm.cpu())
print("Time cost of MyModule: %.3f ms" % (f_timer_before(a_nd, b_nd, c_nd).mean * 1000))
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
Time cost of MyModule: 2.076 ms
接下来，我们通过重新组织循环访问模式来稍微变换 ``MyModule``\ 。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
def schedule_mm(sch: tvm.tir.Schedule, jfactor=4):
block_C = sch.get_block("C", "main")
i, j, k = sch.get_loops(block=block_C)
j_0, j_1 = sch.split(loop=j, factors=[None, jfactor])
sch.reorder(i, j_0, k, j_1)
sch.decompose_reduction(block_C, k)
return sch
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
sch = tvm.tir.Schedule(MyModule)
sch = schedule_mm(sch)
IPython.display.HTML(code2html(sch.mod.script()))
.. 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 main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
for i, j_0 in T.grid(128, 32):
for j_1_init in range(4):
with T.block("C_init"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1_init)
T.reads()
T.writes(C[vi, vj])
C[vi, vj] = T.float32(0)
for k, j_1 in T.grid(128, 4):
with T.block("C_update"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1)
vk = T.axis.reduce(128, k)
T.reads(C[vi, vj], A[vi, vk], B[vk, vj])
T.writes(C[vi, vj])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
for i, j_0 in T.grid(128, 32):
for j_1_init in range(4):
with T.block("C_init"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1_init)
T.reads()
T.writes(C[vi, vj])
C[vi, vj] = T.float32(0)
for k, j_1 in T.grid(128, 4):
with T.block("C_update"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 4 + j_1)
vk = T.axis.reduce(128, k)
T.reads(C[vi, vj], A[vi, vk], B[vk, vj])
T.writes(C[vi, vj])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
for i, j, k in T.grid(128, 128, 128):
with T.block("C"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
T.reads(A[vi, vk], B[vk, vj])
T.writes(C[vi, vj])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
for i, j_0, k, j_1 in T.grid(128, 8, 128, 16):
with T.block("C"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 16 + j_1)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(C[vi, vj])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
for i, j_0 in T.grid(128, 8):
for j_1_init in range(16):
with T.block("C_init"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 16 + j_1_init)
T.reads()
T.writes(C[vi, vj])
C[vi, vj] = T.float32(0)
for k, j_1 in T.grid(128, 16):
with T.block("C_update"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 16 + j_1)
vk = T.axis.reduce(128, k)
T.reads(C[vi, vj], A[vi, vk], B[vk, vj])
T.writes(C[vi, vj])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
Name |
FLOP |
Weight |
Speed (GFLOPS) |
Latency (us) |
Weighted Latency (us) |
Trials |
Done |
| 0 |
main |
4194304 |
1 |
3.4059 |
1231.4906 |
1231.4906 |
5 |
Y |
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
Total trials: 5
Total latency (us): 1231.49
2023-03-18 11:48:04 [DEBUG] [task_scheduler.cc:318]
ID | Name | FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done
------------------------------------------------------------------------------------------------------
0 | main | 4194304 | 1 | 3.4059 | 1231.4906 | 1231.4906 | 5 | Y
------------------------------------------------------------------------------------------------------
Total trials: 5
Total latency (us): 1231.49
``tune_tir`` 函数返回在调优过程中找到的优化后的调度。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
sch.trace.show()
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
/usr/share/miniconda/envs/mlc/lib/python3.8/site-packages/tvm/script/highlight.py:117: UserWarning: No module named 'black'
To print formatted TVM script, please install the formatter 'Black':
/usr/share/miniconda/envs/mlc/bin/python -m pip install "black==22.3.0" --upgrade --user
warnings.warn(
.. raw:: html
# from tvm import tir
def apply_trace(sch: tir.Schedule) -> None:
b0 = sch.get_block(name="C", func_name="main")
l1, l2, l3 = sch.get_loops(block=b0)
v4, v5 = sch.sample_perfect_tile(loop=l2, n=2, max_innermost_factor=16, decision=[16, 8])
l6, l7 = sch.split(loop=l2, factors=[v4, v5], preserve_unit_iters=True)
sch.reorder(l1, l6, l3, l7)
b8 = sch.decompose_reduction(block=b0, loop=l3)
sch.enter_postproc()
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
for i, j_0 in T.grid(128, 16):
for j_1_init in range(8):
with T.block("C_init"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 8 + j_1_init)
T.reads()
T.writes(C[vi, vj])
C[vi, vj] = T.float32(0)
for k, j_1 in T.grid(128, 8):
with T.block("C_update"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j_0 * 8 + j_1)
vk = T.axis.reduce(128, k)
T.reads(C[vi, vj], A[vi, vk], B[vk, vj])
T.writes(C[vi, vj])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
Name |
FLOP |
Weight |
Speed (GFLOPS) |
Latency (us) |
Weighted Latency (us) |
Trials |
Done |
| 0 |
main |
4194304 |
1 |
22.3003 |
188.0832 |
188.0832 |
64 |
Y |
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
Total trials: 64
Total latency (us): 188.083
2023-03-18 11:48:59 [DEBUG] [task_scheduler.cc:318]
ID | Name | FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done
------------------------------------------------------------------------------------------------------
0 | main | 4194304 | 1 | 22.3003 | 188.0832 | 188.0832 | 64 | Y
------------------------------------------------------------------------------------------------------
Total trials: 64
Total latency (us): 188.083
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
lib = tvm.build(sch.mod, target="llvm")
f_timer_after = lib.time_evaluator("main", tvm.cpu())
print("Time cost of MyModule after tuning: %.3f ms" % (f_timer_after(a_nd, b_nd, c_nd).mean * 1000))
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
Time cost of MyModule after tuning: 0.260 ms
结果比我们的原始代码快得多。我们可以查看历史轨迹和最终代码。就本章而言，你不需要了解所有变换。在高层次的理解中，历史轨迹包含：
- 更多级的循环转换
- 中间计算的矢量化
- 并行化和循环展开
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
sch.trace.show()
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
/usr/share/miniconda/envs/mlc/lib/python3.8/site-packages/tvm/script/highlight.py:117: UserWarning: No module named 'black'
To print formatted TVM script, please install the formatter 'Black':
/usr/share/miniconda/envs/mlc/bin/python -m pip install "black==22.3.0" --upgrade --user
warnings.warn(
.. raw:: html
# from tvm import tir
def apply_trace(sch: tir.Schedule) -> None:
b0 = sch.get_block(name="C", func_name="main")
b1 = sch.get_block(name="root", func_name="main")
sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")
l2, l3, l4 = sch.get_loops(block=b0)
v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64, decision=[16, 4, 1, 2])
l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8], preserve_unit_iters=True)
v13, v14, v15, v16 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64, decision=[1, 2, 4, 16])
l17, l18, l19, l20 = sch.split(loop=l3, factors=[v13, v14, v15, v16], preserve_unit_iters=True)
v21, v22 = sch.sample_perfect_tile(loop=l4, n=2, max_innermost_factor=64, decision=[16, 8])
l23, l24 = sch.split(loop=l4, factors=[v21, v22], preserve_unit_iters=True)
sch.reorder(l9, l17, l10, l18, l23, l11, l19, l24, l12, l20)
b25 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")
sch.reverse_compute_at(block=b25, loop=l18, preserve_unit_loops=True, index=-1)
sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=16)
sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=64)
v26 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25], decision=2)
sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v26)
sch.enter_postproc()
b27 = sch.get_block(name="root", func_name="main")
sch.unannotate(block_or_loop=b27, ann_key="meta_schedule.parallel")
sch.unannotate(block_or_loop=b27, ann_key="meta_schedule.vectorize")
sch.unannotate(block_or_loop=b27, ann_key="meta_schedule.unroll_explicit")
b28, b29 = sch.get_child_blocks(b27)
l30, l31, l32, l33, l34, l35, l36, l37, l38, l39 = sch.get_loops(block=b28)
l40 = sch.fuse(l30, l31, l32, preserve_unit_iters=True)
sch.parallel(loop=l40)
l41 = sch.fuse(l39, preserve_unit_iters=True)
sch.vectorize(loop=l41)
sch.annotate(block_or_loop=l40, ann_key="pragma_auto_unroll_max_step", ann_val=64)
sch.annotate(block_or_loop=l40, ann_key="pragma_unroll_explicit", ann_val=1)
l42, l43, l44, l45 = sch.get_loops(block=b29)
l46 = sch.fuse(l45, preserve_unit_iters=True)
sch.vectorize(loop=l46)
sch.annotate(block_or_loop=l42, ann_key="pragma_auto_unroll_max_step", ann_val=64)
sch.annotate(block_or_loop=l42, ann_key="pragma_unroll_explicit", ann_val=1)
b47 = sch.get_block(name="C", func_name="main")
l48, l49, l50, l51, l52, l53, l54, l55 = sch.get_loops(block=b47)
b56 = sch.decompose_reduction(block=b47, loop=l50)
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
C_global = T.alloc_buffer((128, 128))
for i_0_j_0_i_1_fused in T.parallel(64, annotations={"pragma_auto_unroll_max_step": 64, "pragma_unroll_explicit": 1}):
for j_1 in range(2):
for i_2_init, j_2_init, i_3_init in T.grid(1, 4, 2):
for j_3_fused_init in T.vectorized(16):
with T.block("C_init"):
vi = T.axis.spatial(128, i_0_j_0_i_1_fused // 4 * 8 + i_0_j_0_i_1_fused % 4 * 2 + i_2_init * 2 + i_3_init)
vj = T.axis.spatial(128, j_1 * 64 + j_2_init * 16 + j_3_fused_init)
T.reads()
T.writes(C_global[vi, vj])
T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
C_global[vi, vj] = T.float32(0)
for k_0, i_2, j_2, k_1, i_3 in T.grid(16, 1, 4, 8, 2):
for j_3_fused in T.vectorized(16):
with T.block("C_update"):
vi = T.axis.spatial(128, i_0_j_0_i_1_fused // 4 * 8 + i_0_j_0_i_1_fused % 4 * 2 + i_2 * 2 + i_3)
vj = T.axis.spatial(128, j_1 * 64 + j_2 * 16 + j_3_fused)
vk = T.axis.reduce(128, k_0 * 8 + k_1)
T.reads(C_global[vi, vj], A[vi, vk], B[vk, vj])
T.writes(C_global[vi, vj])
T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
C_global[vi, vj] = C_global[vi, vj] + A[vi, vk] * B[vk, vj]
for ax0 in range(2):
for ax1_fused in T.vectorized(64):
with T.block("C_global"):
v0 = T.axis.spatial(128, i_0_j_0_i_1_fused // 4 * 8 + i_0_j_0_i_1_fused % 4 * 2 + ax0)
v1 = T.axis.spatial(128, j_1 * 64 + ax1_fused)
T.reads(C_global[v0, v1])
T.writes(C[v0, v1])
C[v0, v1] = C_global[v0, v1]
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(X: T.Buffer((1, 784), "float32"), W: T.Buffer((128, 784), "float32"), B: T.Buffer((128,), "float32"), Z: T.Buffer((1, 128), "float32")):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((1, 128))
for i, j, k in T.grid(1, 128, 784):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
T.reads(X[vi, vk], W[vj, vk])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]
for i, j in T.grid(1, 128):
with T.block("Z"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj], B[vj])
T.writes(Z[vi, vj])
Z[vi, vj] = Y[vi, vj] + B[vj]
|
Name |
FLOP |
Weight |
Speed (GFLOPS) |
Latency (us) |
Weighted Latency (us) |
Trials |
Done |
| 0 |
main |
200832 |
1 |
6.9648 |
28.8354 |
28.8354 |
64 |
Y |
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
2023-03-18 11:49:55 [DEBUG] [task_scheduler.cc:318]
ID | Name | FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done
-----------------------------------------------------------------------------------------------------
0 | main | 200832 | 1 | 6.9648 | 28.8354 | 28.8354 | 64 | Y
-----------------------------------------------------------------------------------------------------
Total trials: 64
Total latency (us): 28.8354
Total trials: 64
Total latency (us): 28.8354
现在我们需要在调优后用新函数替换原来的
``linear0``\ 。我们可以通过首先获得一个 ``global_var``\ （一个指向
IRModule 中函数的 ``pointer`` 引用），然后调用 ``update_func``
来用新的函数替换原本的函数。
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
MyModuleWithParams2 = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
new_func = sch.mod["main"].with_attr("global_symbol", "linear0")
gv = MyModuleWithParams2.get_global_var("linear0")
MyModuleWithParams2.update_func(gv, new_func)
IPython.display.HTML(code2html(MyModuleWithParams2.script()))
.. raw:: html
# from tvm.script import ir as I
# from tvm.script import tir as T
# from tvm.script import relax as R
@I.ir_module
class Module:
@T.prim_func
def linear0(X: T.Buffer((1, 784), "float32"), W: T.Buffer((128, 784), "float32"), B: T.Buffer((128,), "float32"), Z: T.Buffer((1, 128), "float32")):
T.func_attr({"global_symbol": "linear0", "tir.noalias": True})
# with T.block("root"):
Y = T.alloc_buffer((1, 128))
for i_0, j_0 in T.grid(1, 1):
for i_1, j_1 in T.grid(1, 8):
for i_2_init, j_2_init, i_3_init in T.grid(1, 2, 1):
for j_3_fused_init in T.vectorized(8):
with T.block("Y_init"):
vi = T.axis.spatial(1, i_0 + i_1 + i_2_init + i_3_init)
vj = T.axis.spatial(128, j_0 * 128 + j_1 * 16 + j_2_init * 8 + j_3_fused_init)
T.reads()
T.writes(Y[vi, vj])
T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
Y[vi, vj] = T.float32(0)
for k_0, i_2, j_2, k_1, i_3 in T.grid(14, 1, 2, 56, 1):
for j_3_fused in T.vectorized(8):
with T.block("Y_update"):
vi = T.axis.spatial(1, i_0 + i_1 + i_2 + i_3)
vj = T.axis.spatial(128, j_0 * 128 + j_1 * 16 + j_2 * 8 + j_3_fused)
vk = T.axis.reduce(784, k_0 * 56 + k_1)
T.reads(Y[vi, vj], X[vi, vk], W[vj, vk])
T.writes(Y[vi, vj])
T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]
for ax0, ax1 in T.grid(1, 128):
with T.block("Z"):
vi, vj = T.axis.remap("SS", [ax0, ax1])
T.reads(Y[vi, vj], B[vj])
T.writes(Z[vi, vj])
Z[vi, vj] = Y[vi, vj] + B[vj]
@R.function
def main(x: R.Tensor((1, 784), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
with R.dataflow():
lv0 = R.call_dps_packed("linear0", (x, metadata["relax.expr.Constant"][0], metadata["relax.expr.Constant"][1]), out_sinfo=R.Tensor((1, 128), dtype="float32"))
lv1 = R.call_dps_packed("env.relu", (lv0,), out_sinfo=R.Tensor((1, 128), dtype="float32"))
out = R.call_dps_packed("env.linear", (lv1, metadata["relax.expr.Constant"][2], metadata["relax.expr.Constant"][3]), out_sinfo=R.Tensor((1, 10), dtype="float32"))
R.output(out)
return out
# Metadata omitted. Use show_meta=True in script() method to show it.