TVM User Tutorial -- Working with Operators Using Tensor Expression

最新推荐文章于 2024-07-10 17:19:49 发布

姆克儿

最新推荐文章于 2024-07-10 17:19:49 发布

阅读量194

点赞数

分类专栏： TVM文档翻译文章标签： ubuntu linux TVM

原文链接：https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html#sphx-glr-tutorial-tensor-expr-get-started-py

版权

TVM文档翻译专栏收录该内容

12 篇文章 2 订阅

订阅专栏

Author: Tianqi Chen

在本教程中，我们将关注 TVM 如何与张量表达式 (Tensor Expression, TE) 一起定义张量计算并应用循环优化。 TE 用纯函数式语言描述张量计算（即每个表达式都没有副作用）。从 TVM 的整体来看，Relay 将计算描述为一组算子，每个算子都可以表示为一个 TE 表达式，其中每个 TE 表达式接受输入张量并产生一个输出张量。

这是 TVM 中张量表达式语言的介绍性教程。 TVM 使用特定领域的张量表达式来进行有效的内核构建。我们将通过两个使用张量表达式语言的示例来演示基本工作流程。第一个例子介绍了 TE 和带有向量加法的调度。第二个扩展了这些概念，逐步优化了与 TE 的矩阵乘法。这个矩阵乘法示例将作为未来教程的比较基础，涵盖 TVM 的更高级功能。

例 1：在TE中，CPU向量加法的编写与调度方法

让我们看一个 Python 中的示例，在该示例中，我们将实现一个用于向量加法的 TE，然后是针对 CPU 的调度。我们首先初始化一个 TVM 环境。

import tvm
import tvm.testing
from tvm import te
import numpy as np

如果您可以识别您的目标 CPU 并指定它，您将获得更好的性能。如果您使用 LLVM，您可以从命令 llc --version 获取此信息以获取 CPU 类型，并且您可以检查 /proc/cpuinfo 以获取您的处理器可能支持的其他扩展。例如，您可以将 llvm -mcpu=skylake-avx512 用于具有 AVX-512 指令的 CPU。

tgt = tvm.target.Target(target="llvm", host="llvm")

描述向量计算

我们描述了向量加法计算。 TVM 采用张量语义，每个中间结果表示为一个多维数组。用户需要描述生成张量的计算规则。我们首先定义一个符号变量 n 来表示形状。然后我们定义两个占位符张量，A 和 B，具有给定的形状 (n,)。然后我们用compute操作描述结果张量 C。 compute定义了一个计算，输出符合指定的张量形状，计算将在 lambda 函数定义的张量中的每个位置执行。请注意，虽然 n 是一个变量，但它定义了 A、B 和 C 张量之间的一致形状。请记住，在此阶段不会发生实际的计算，因为我们只是在声明应该如何进行计算。

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")

lambda函数
`te.compute` 方法的第二个参数是执行计算的函数。在此示例中，我们使用匿名函数（也称为 `lambda` 函数）来定义计算，在本例中是对 `A` 和 `B` 的第 `i` 个元素进行加法运算。

为计算创建默认调度

虽然以上几行描述了计算规则，但我们可以以许多不同的方式计算 C 以适应不同的设备。对于具有多个维度的张量，您可以选择首先迭代哪个维度，或者可以将计算拆分到不同的线程中。 TVM 要求用户提供调度，这是对如何执行计算的描述。 TE 中的调度操作可以更改循环顺序、跨不同线程拆分计算以及将数据块组合在一起，以及其他操作。调度背后的一个重要概念是它们只描述如何执行计算，因此同一 TE 的不同调度将产生相同的结果。

TVM 允许您创建一个简单的计划，该计划将通过按行优先顺序迭代来计算 C。

for (int i = 0; i < n; ++i) {
  C[i] = A[i] + B[i];
}

s = te.create_schedule(C.op)

编译和评估默认调度

使用 TE 表达式和调度，我们可以为我们的目标语言和架构（在本例中为 LLVM 和 CPU）生成可运行的代码。我们向 TVM 提供调度、调度中的 TE 表达式列表、目标和主机，以及我们正在生成的函数的名称。输出的结果是一个类型擦除的函数，可以直接从 Python 调用。

在以下行中，我们使用 tvm.build 创建一个函数。 build 函数采用调度、所需的函数签名（包括输入和输出）以及我们要编译到的目标语言。

fadd = tvm.build(s, [A, B, C], tgt, name="myadd")

让我们运行该函数，并将输出与 numpy 中的相同计算进行比较。编译后的 TVM 函数公开了一个简洁的 C API，可以从任何语言调用。我们首先创建一个设备，这是一个 TVM 可以将调度编译到的设备（本例中为 CPU）。在这种情况下，设备是 LLVM CPU 目标。然后我们可以初始化设备中的张量并执行自定义加法操作。为了验证计算是否正确，我们可以将 c 张量的输出结果与 numpy 执行的相同计算进行比较。

dev = tvm.device(tgt.kind.name, 0)

n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
fadd(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

要比较此版本与 numpy 相比的速度，请创建一个辅助函数来运行 TVM 生成的代码的配置文件。

import timeit

np_repeat = 100
np_running_time = timeit.timeit(
    setup="import numpy\n"
    "n = 32768\n"
    'dtype = "float32"\n'
    "a = numpy.random.rand(n, 1).astype(dtype)\n"
    "b = numpy.random.rand(n, 1).astype(dtype)\n",
    stmt="answer = a + b",
    number=np_repeat,
)
print("Numpy running time: %f" % (np_running_time / np_repeat))


def evaluate_addition(func, target, optimization, log):
    dev = tvm.device(target.kind.name, 0)
    n = 32768
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)

    evaluator = func.time_evaluator(func.entry_name, dev, number=10)
    mean_time = evaluator(a, b, c).mean
    print("%s: %f" % (optimization, mean_time))

    log.append((optimization, mean_time))


log = [("numpy", np_running_time / np_repeat)]
evaluate_addition(fadd, tgt, "naive", log=log)

Out:
Numpy running time: 0.000007
naive: 0.000005

使用并行性调度

现在我们已经说明了 TE 的基础知识，让我们更深入地了解调度的作用，以及如何使用它们来优化不同架构的张量表达式。调度是应用于表达式以以多种不同方式对其进行转换的一系列步骤。当调度应用于 TE 中的表达式时，输入和输出保持不变，但在编译时，表达式的实现可能会发生变化。在默认调度中，这种张量添加是串行运行的，但很容易在所有处理器线程上并行化。我们可以将并行调度操作应用于我们的计算。

s[C].parallel(C.op.axis[0])

tvm.lower 命令将生成 TE 的中间表示 (IR) 以及相应的调度。通过在应用不同的调度操作时降低表达式，我们可以看到调度对计算顺序的影响。我们使用标志 simple_mode=True 来返回可读的 C 风格语句。

print(tvm.lower(s, [A, B, C], simple_mode=True))

Out:
@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),
             B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, n) "parallel" {
    C_2[(i*stride)] = ((float32*)A_2[(i*stride_1)] + (float32*)B_2[(i*stride_2)])
  }
}

TVM 现在可以在独立线程上运行这些块。让我们在应用并行操作的情况下编译并运行这个新计划：

fadd_parallel = tvm.build(s, [A, B, C], tgt, name="myadd_parallel")
fadd_parallel(a, b, c)

tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

evaluate_addition(fadd_parallel, tgt, "parallel", log=log)

Out:
parallel: 0.000004

使用向量化调度

现代 CPU 还具有对浮点值执行 SIMD 操作的能力，我们可以将另一个调度应用于我们的计算表达式以利用这一点。实现这一点需要多个步骤：首先，我们必须使用拆分调度原语将调度拆分为内部和外部循环。内部循环可以使用向量化来使用向量化调度原语使用 SIMD 指令，然后可以使用并行调度原语对外部循环进行并行化。选择拆分因子作为 CPU 上的线程数。

# Recreate the schedule, since we modified it with the parallel operation in
# the previous example
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")

s = te.create_schedule(C.op)

# This factor should be chosen to match the number of threads appropriate for
# your CPU. This will vary depending on architecture, but a good rule is
# setting this factor to equal the number of available CPU cores.
factor = 4

outer, inner = s[C].split(C.op.axis[0], factor=factor)
s[C].parallel(outer)
s[C].vectorize(inner)

fadd_vector = tvm.build(s, [A, B, C], tgt, name="myadd_parallel")

evaluate_addition(fadd_vector, tgt, "vector", log=log)

print(tvm.lower(s, [A, B, C], simple_mode=True))

Out:
vector: 0.000022
@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),
             B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i.outer: int32, 0, floordiv((n + 3), 4)) "parallel" {
    for (i.inner.s: int32, 0, 4) {
      if @tir.likely((((i.outer*4) + i.inner.s) < n), dtype=bool) {
        C_2[(((i.outer*4) + i.inner.s)*stride)] = ((float32*)A_2[(((i.outer*4) + i.inner.s)*stride_1)] + (float32*)B_2[(((i.outer*4) + i.inner.s)*stride_2)])
      }
    }
  }
}

比较不同调度

baseline = log[0][1]
print("%s\t%s\t%s" % ("Operator".rjust(20), "Timing".rjust(20), "Performance".rjust(20)))
for result in log:
    print(
        "%s\t%s\t%s"
        % (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20))
    )

Out:
Operator                  Timing             Performance
   numpy    6.619859486818314e-06                    1.0
   naive              4.9303e-06      0.7447741164019234
parallel              4.4504e-06       0.672280130546847
  vector             2.20605e-05      3.3324725462719575

代码Specialization

代码Specialization
您可能已经注意到，`A`、`B` 和 `C` 的声明都采用相同的形状参数 `n`。 TVM 将利用这一点仅将单个形状参数传递给内核，正如您将在打印的设备代码中找到的那样。这是Specialization的一种形式。在主机端，TVM 将自动生成检查代码来检查参数中的约束。因此，如果将不同形状的数组传递给 fadd，则会引发错误。我们可以做更多的Specialization。例如，我们可以在计算声明中写 `n = tvm.runtime.convert(1024)` 而不是 `n = te.var("n")`。生成的函数只会采用长度为 1024 的向量。

您可能已经注意到，A、B 和 C 的声明都采用相同的形状参数 n。 TVM 将利用这一点仅将单个形状参数传递给内核，正如您将在打印的设备代码中找到的那样。这是Specialization的一种形式。在主机端，TVM 将自动生成检查代码来检查参数中的约束。因此，如果将不同形状的数组传递给 fadd，则会引发错误。我们可以做更多的Specialization。例如，我们可以在计算声明中写 n = tvm.runtime.convert(1024) 而不是 n = te.var("n")。生成的函数只会采用长度为 1024 的向量。

我们已经定义、调度和编译了一个向量加法运算符，然后我们可以在 TVM 运行时执行它。我们可以将算子保存为一个库，稍后我们可以使用 TVM 运行时加载它。

GPU 的向量加法（可选）

TVM 能够针对多种架构。在下一个示例中，我们将针对 GPU 的向量添加编译。

# If you want to run this code, change ``run_cuda = True``
# Note that by default this example is not run in the docs CI.

run_cuda = False
if run_cuda:
    # Change this target to the correct backend for you gpu. For example: cuda (NVIDIA GPUs),
    # rocm (Radeon GPUS), OpenCL (opencl).
    tgt_gpu = tvm.target.Target(target="cuda", host="llvm")

    # Recreate the schedule
    n = te.var("n")
    A = te.placeholder((n,), name="A")
    B = te.placeholder((n,), name="B")
    C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
    print(type(C))

    s = te.create_schedule(C.op)

    bx, tx = s[C].split(C.op.axis[0], factor=64)

    ################################################################################
    # Finally we must bind the iteration axis bx and tx to threads in the GPU
    # compute grid. The naive schedule is not valid for GPUs, and these are
    # specific constructs that allow us to generate code that runs on a GPU.

    s[C].bind(bx, te.thread_axis("blockIdx.x"))
    s[C].bind(tx, te.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 the 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 GPU device function (if GPU is
    # involved) as well as a host wrapper that calls into the GPU
    # function. fadd is the generated host wrapper function, it contains
    # a reference to the generated device function internally.

    fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name="myadd")

    ################################################################################
    # The compiled TVM function exposes a concise C API that can be invoked from
    # any language.
    #
    # We provide a minimal array API in python to aid quick testing and prototyping.
    # The array API is based on the `DLPack <https://github.com/dmlc/dlpack>`_ standard.
    #
    # - We first create a GPU device.
    # - Then tvm.nd.array copies the data to the GPU.
    # - ``fadd`` runs the actual computation
    # - ``numpy()`` copies the GPU array back to the CPU (so we can verify correctness).
    #
    # Note that copying the data to and from the memory on the GPU is a required step.

    dev = tvm.device(tgt_gpu.kind.name, 0)

    n = 1024
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
    fadd(a, b, c)
    tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

    ################################################################################
    # Inspect the Generated GPU 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 (GPU) function.
    #
    # The following code fetches the device module and prints the content code.

    if (
        tgt_gpu.kind.name == "cuda"
        or tgt_gpu.kind.name == "rocm"
        or tgt_gpu.kind.name.startswith("opencl")
    ):
        dev_module = fadd.imported_modules[0]
        print("-----GPU code-----")
        print(dev_module.get_source())
    else:
        print(fadd.get_source())

保存和加载已编译的模块

除了运行时编译之外，我们还可以将编译后的模块保存到一个文件中，稍后再加载它们。
以下代码首先执行以下步骤：

它将编译的主机模块保存到目标文件中。
然后它将设备模块保存到 ptx 文件中。
cc.create_shared 调用编译器 (gcc) 来创建共享库

from tvm.contrib import cc
from tvm.contrib import utils

temp = utils.tempdir()
fadd.save(temp.relpath("myadd.o"))
if tgt.kind.name == "cuda":
    fadd.imported_modules[0].save(temp.relpath("myadd.ptx"))
if tgt.kind.name == "rocm":
    fadd.imported_modules[0].save(temp.relpath("myadd.hsaco"))
if tgt.kind.name.startswith("opencl"):
    fadd.imported_modules[0].save(temp.relpath("myadd.cl"))
cc.create_shared(temp.relpath("myadd.so"), [temp.relpath("myadd.o")])
print(temp.listdir())

Out:
['myadd.o', 'myadd.so']

模块存储格式
CPU（主机）模块直接保存为共享库（.so）。设备代码可以有多种自定义格式。在我们的示例中，设备代码存储在 ptx 以及元数据 json 文件中。它们可以通过导入单独加载和链接。

加载已编译模块

我们可以从文件系统中加载编译好的模块并运行代码。以下代码分别加载主机和设备模块并将它们链接在一起。我们可以验证新加载的函数是否有效。

fadd1 = tvm.runtime.load_module(temp.relpath("myadd.so"))
if tgt.kind.name == "cuda":
    fadd1_dev = tvm.runtime.load_module(temp.relpath("myadd.ptx"))
    fadd1.import_module(fadd1_dev)

if tgt.kind.name == "rocm":
    fadd1_dev = tvm.runtime.load_module(temp.relpath("myadd.hsaco"))
    fadd1.import_module(fadd1_dev)

if tgt.kind.name.startswith("opencl"):
    fadd1_dev = tvm.runtime.load_module(temp.relpath("myadd.cl"))
    fadd1.import_module(fadd1_dev)

fadd1(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

将所有内容打包到一个库中

在上面的示例中，我们分别存储了设备和主机代码。 TVM 还支持将所有内容导出为一个共享库。在后台，我们将设备模块打包成二进制 blob，并将它们与主机代码链接在一起。目前我们支持 Metal、OpenCL 和 CUDA 模块的打包。

fadd.export_library(temp.relpath("myadd_pack.so"))
fadd2 = tvm.runtime.load_module(temp.relpath("myadd_pack.so"))
fadd2(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

运行时 API 和线程安全
TVM 的编译模块不依赖于 TVM 编译器。相反，它们只依赖于最小的runtime库。 TVM 运行时库封装了设备驱动程序，并在编译后的函数中提供线程安全和设备无关的调用。这意味着您可以从任何线程、任何 GPU 上调用已编译的 TVM 函数，前提是您已经为该 GPU 编译了代码。

生成 OpenCL 代码

TVM 为多个后端提供代码生成功能。我们还可以生成在 CPU 后端运行的 OpenCL 代码或 LLVM 代码。
以下代码块生成 OpenCL 代码，在 OpenCL 设备上创建数组，并验证代码的正确性。

if tgt.kind.name.startswith("opencl"):
    fadd_cl = tvm.build(s, [A, B, C], tgt, name="myadd")
    print("------opencl code------")
    print(fadd_cl.imported_modules[0].get_source())
    dev = tvm.cl(0)
    n = 1024
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
    fadd_cl(a, b, c)
    tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

TE调度原语
这些原语的完整描述可以在 Schedule Primitives 文档页面中找到。

姆克儿

关注

0
点赞
踩
0

收藏

觉得还不错? 一键收藏
0
评论
TVM User Tutorial -- Working with Operators Using Tensor Expression

Author: Tianqi Chen在本教程中，我们将关注 TVM 如何与张量表达式 (Tensor Expression, TE) 一起定义张量计算并应用循环优化。 TE 用纯函数式语言描述张量计算（即每个表达式都没有副作用）。从 TVM 的整体来看，Relay 将计算描述为一组算子，每个算子都可以表示为一个 TE 表达式，其中每个 TE 表达式接受输入张量并产生一个输出张量。这是 TVM 中张量表达式语言的介绍性教程。 TVM 使用特定领域的张量表达式来进行有效的内核构建。我们将通过两个使
复制链接

扫一扫