基于调度模板和 AutoTVM 的算子优化

最新推荐文章于 2024-09-15 22:31:42 发布

6小贱

最新推荐文章于 2024-09-15 22:31:42 发布

阅读量128

点赞数

分类专栏： TVM 文章标签： python 机器学习深度学习

本文链接：https://blog.csdn.net/u013947807/article/details/125000003

版权

本文介绍了如何利用调度模板和AutoTVM进行算子优化，特别是矩阵相乘。内容包括安装依赖，基于TE的基本矩阵相乘，以及使用AutoTVM优化矩阵相乘的步骤，如定义搜索空间，使用高阶参数API，以及选择不同的优化器策略进行调优。

摘要由CSDN通过智能技术生成

自动优化的步骤：

定义一个搜索空间
在搜索空间中执行搜索算法

安装依赖

为了使用自动优化，需要安装如下的依赖:

pip3 install --user psutil xgboost cloudpickle

为了使得 TVM 在优化过程中运行更快，推荐使用 cython：

pip3 install --user cython
sudo make cython3

基于 TE 的基本矩阵相乘

import logging
import sys

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

# the module is called `autotvm`
from tvm import autotvm

# 基于 TE 实现的基本矩阵相乘
def matmul_basic(N, L, M, dtype):

    A = te.placeholder((N, L), name="A", dtype=dtype)
    B = te.placeholder((L, M), name="B", dtype=dtype)

    k = te.reduce_axis((0, L), name="k")
    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")
    s = te.create_schedule(C.op)

    # schedule
    y, x = s[C].op.axis
    k = s[C].op.reduce_axis[0]
    
    # 使用 8 作为 tiling factor,但这个不一定是最佳的，这个值需要根据实际的硬件架构优化
    yo, yi = s[C].split(y, 8)
    xo, xi = s[C].split(x, 8)

    s[C].reorder(yo, xo, k, yi, xi)

    return s, [A, B, C]

基于 AutoTVM 的矩阵相乘

基础的矩阵相乘模板

# Matmul V1: List candidate values
@autotvm.template("tutorial/matmul_v1")  # 1. use a decorator
def matmul_v1(N, L, M, dtype):
    A = te.placeholder((N, L), name="A", dtype=dtype)
    B = te.placeholder((L, M), name="B", dtype=dtype)

    k = te.reduce_axis((0, L), name="k")
    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")
    s = te.create_s