编程模型和语言层

运行过程与结果与NV相同

TVM 是一个专注于深度学习模型优化和编译的开源框架，它的编程模型基于 Tensor 表达、算子（Operator）定义 和 调度（Schedule），并且通过高效代码生成实现硬件上的高性能计算。TVM 编程模型的关键步骤包括：

  1. 定义计算：使用 TVM 的计算表达式定义要执行的计算任务。
  2. 调度（Scheduling）：为计算任务安排执行顺序和资源分配，指定如何并行化、向量化等。
  3. 编译：将计算和调度方案编译为可在目标硬件上运行的代码。
  4. 执行：在目标设备上运行编译后的代码。

下面是一个简单的示例，它展示了TVM的基本编程流程，包括定义矩阵加法运算并在CPU或者GPU上执行。

TVM 的编程模型示例

import tvm
from tvm import te
import numpy as np

# 1. 定义计算：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")

# 2. 创建调度：默认调度会按顺序执行计算
s = te.create_schedule(C.op)

# 3. 编译：为 GPU 生成低级代码
target = "llvm"  # CPU 目标 如果是hi是GPU上执行则改为“cuda”
fadd = tvm.build(s, [A, B, C], target, name="matrix_add")

# 4. 在 TVM 运行时中执行
ctx = tvm.cpu(0)
n_value = 1024
a = tvm.nd.array(np.random.uniform(size=n_value).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n_value).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n_value, dtype=C.dtype), ctx)

# 调用编译好的函数
fadd(a, b, c)

# 检查结果
np.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
print("Result matches with NumPy calculation.")

代码解释

  1. 定义计算：

    • 使用 te.placeholder 创建两个占位符 A 和 B，分别代表输入的矩阵。
    • 使用 te.compute 定义计算表达式，这里表示逐元素对 A 和 B 执行加法操作，结果存储在 C 中。

  2. 调度计算：

    • 使用 te.create_schedule(C.op) 创建调度。这里使用的是默认的顺序执行调度，也可以通过优化调度提升性能。

  3. 编译代码：

    • 使用 tvm.build 函数，将计算和调度编译成针对指定目标（如 CPU 或 GPU）的可执行代码。

  4. 运行并验证结果：

    • 创建 TVM 的 nd.array 将 NumPy 数据传入 TVM 中运行。
    • 使用编译好的函数 fadd 进行计算，并验证结果是否与 NumPy 计算的结果一致。 结果
Result matches with NumPy calculation.

这个简单的例子展示了 TVM 的核心编程流程。在实际的深度学习模型优化中，TVM 提供了更多高级特性，例如自动调度（AutoScheduler）、多目标硬件支持（CPU、GPU、TPU）等，可以极大提升模型在不同硬件平台上的运行效率。

下面再给一个 TVM 示例，这次展示如何使用 TVM 优化二维矩阵乘法（矩阵乘法是深度学习中常见的操作之一），并进行简单的调度优化。

TVM 矩阵乘法示例

import tvm
from tvm import te
import numpy as np

# 1. 定义矩阵乘法计算: C[i, j] = sum(A[i, k] * B[k, j] for k in range(K))
N = te.var("N")
M = te.var("M")
K = te.var("K")

# 定义矩阵 A, B
A = te.placeholder((N, K), name="A")
B = te.placeholder((K, M), name="B")

# 定义矩阵乘法的计算
C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j] for k in range(K)), name="C")

# 2. 创建调度
s = te.create_schedule(C.op)

# 简单调度优化：并行化行方向上的计算
# 这是 TVM 调度中常见的优化方式
s[C].parallel(C.op.axis[0])

# 3. 编译：为 CPU 生成代码
target = "llvm"  # CPU 目标
fmatmul = tvm.build(s, [A, B, C], target, name="matrix_multiply")

# 4. 在 TVM 运行时中执行
ctx = tvm.cpu(0)

# 定义矩阵的大小
N_value = 1024
M_value = 1024
K_value = 1024

# 创建随机的输入矩阵 A 和 B
a = tvm.nd.array(np.random.uniform(size=(N_value, K_value)).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=(K_value, M_value)).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros((N_value, M_value), dtype=C.dtype), ctx)

# 调用编译好的矩阵乘法函数
fmatmul(a, b, c)

# 使用 NumPy 计算参考结果
np_c = np.dot(a.asnumpy(), b.asnumpy())

# 验证 TVM 计算的结果是否与 NumPy 的结果一致
np.testing.assert_allclose(c.asnumpy(), np_c, rtol=1e-5)
print("Matrix multiplication result matches with NumPy.")

代码解释

  1. 定义矩阵乘法：

    • 使用 te.placeholder 创建两个占位符 A 和 B，分别代表输入的二维矩阵。
    • 使用 te.compute 定义矩阵乘法，te.sum 用于对中间维度 K 进行求和，从而实现矩阵乘法的核心计算。

  2. 调度优化：

    • 使用 te.create_schedule(C.op) 创建调度方案。
    • 通过 s[C].parallel(C.op.axis[0]) 让 TVM 并行化行方向上的计算，这是一个简单的优化方法，用于利用多核 CPU 提升矩阵乘法的性能。

  3. 编译代码：

    • 使用 tvm.build 将计算和调度方案编译成可执行的 CPU 代码。

  4. 运行并验证结果：

    • 创建随机的矩阵 A 和 B，在 TVM 的运行时环境中执行编译好的矩阵乘法函数。
    • 使用 NumPy 的 dot 函数计算参考矩阵乘法的结果，并与 TVM 的结果进行比较，确保其一致性。

结果

Matrix multiplication result matches with NumPy.

这表明 TVM 编译后的矩阵乘法操作正确地执行，并且与 NumPy 的计算结果一致。