机器学习编译总结

Posted on 2022-11-21 Edited on 2025-03-21 In 机器学习编译 Valine:

TVM 是什么：是 Tensor Virtual Machine 的缩写，是一个 Open Deep Learning Compiler Stack（深度学习开源编译栈）
TVM 想干什么：将机器学习算法从开发阶段形态，通过变换和优化算法，使其变成部署形态
TVM 的原则：
- 集成与最小化依赖
- 利用硬件加速
- 通用优化
TVM Module 层次结构：
- IRModule：包含一个或多个 元张量函数 和一个 高层神经网络执行的抽象。通常用 @tvm.script.ir_module 装饰器装饰
- tensorIR: 元张量函数。通常表示一个算子实例的计算过程，包含多个 计算块。通常用 @T.prim_func 装饰器装饰
- 高层神经网络执行的抽象：IRModule 的程序入口。通常用 @R.function
- block: 计算块。张量的基本计算单位，通常包含多个 计算轴 上的循环。通常用 with T.block(block_name) 来标明作用域
- 计算轴：
  - 空间轴（spatial axis）：空间轴上循环的每个位置的计算独立于其他位置
  - 规约轴（reduce axis）：规约轴上的位置不会反映到最后的计算输出上
TVM Module 变换过程：
- 自动程序优化
- cuda 多线程优化
- 内存优化
- 图优化
- 等等

TVM Module 执行过程：

ex = relax.vm.build(MyModule, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
nd_res = vm["main"](
    data_nd, nd_params["w0"], nd_params["b0"], nd_params["w1"], nd_params["b1"]
)

可执行程序 = build(IR_Module)
虚拟机执行器 = 虚拟机(可执行程序)
运行结果 = 虚拟机执行器(模型输入 + 模型权重)

机器学习编译(6)——GPU硬件加速

Posted on 2022-11-09 Edited on 2025-03-21 In 机器学习编译 Valine:

URL

https://mlc.ai/zh/chapter_gpu_acceleration/index.html

背景知识

CUDA（Compute Unified Device Architecture，同一计算设备架构）教程：

GPU 体系结构

物理模型
- 典型的 GPU 包含一组流处理器 (stream multi-processors, SM)，每个流处理器都有许多核心，硬件实现上这些核心之间可共享内存（shared memory）
逻辑模型
- 逻辑模型中，引入了 Grid / Block / Thread 三级概念，逻辑模型与物理的对应关系如下图所示：
  
  因此：同一个 Block 中的 Thread 可共享 shared memory
Memory Hierarchy

shared memory 速度几乎和 L1 cache 一样，比 local memory 和 global memory 都快的多（在物理上，local memory 和 global memory 是同一块 DRAM）
在对 GPU 进行编程时，需要创建一组进程块 (thread blocks)，每个 thread 映射到单个核心，而 block 映射到流式多处理器 (SM)，如下图所示：
每个线程可由 threadIdx 和 blockIdx 索引，在实际应用中，可以有多维线程索引

Element-wise Add GPU 加速

两个向量 A 和 B，向量长度都为 1024,执行元素相加，并将结果存储在 C 中

@tvm.script.ir_module
class MyModuleVecAdd:
    @T.prim_func
    def main(A: T.Buffer[(1024,), "float32"],
            B: T.Buffer[(1024,), "float32"],
            C: T.Buffer[(1024,), "float32"]) -> None:
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i in T.grid(1024):
            with T.block("C"):
                vi = T.axis.remap("S", [i])
                C[vi] = A[vi] + B[vi]

首先将循环 i 拆分成两个循环:

sch = tvm.tir.Schedule(MyModuleVecAdd)
block_C = sch.get_block("C")
i, = sch.get_loops(block=block_C)
i0, i1 = sch.split(i, [None, 128])

将迭代器绑定到 GPU 线程块。每个线程由两个索引进行表示 threadIdx.x 和 blockIdx.x
1
2
3
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
sch.mod.show()

绑定后的代码：

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1024, "float32"], B: T.Buffer[1024, "float32"], C: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A[vi], B[vi])
                    T.writes(C[vi])
                    C[vi] = A[vi] + B[vi]

由于 Element-wise Add 不存在数据依赖，所以可以直接拆分到多个 block 中的多个 thread 中，一个 cycle 全部算完

窗口求和 GPU 加速

相邻三个窗口求和，输入向量 A 长度 1026，输出 B 长度 1024。（即无 padding 的权重为 [1, 1, 1] 的 conv1d）

@tvm.script.ir_module
class MyModuleWindowSum:
    @T.prim_func
    def main(A: T.Buffer[(1026,), "float32"],
            B: T.Buffer[(1024,), "float32"]) -> None:
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i in T.grid(1024):
            with T.block("C"):
                vi = T.axis.remap("S", [i])
                B[vi] = A[vi] + A[vi + 1] + A[vi + 2]

拆分循环并绑定到 block 和 thread

sch = tvm.tir.Schedule(MyModuleWindowSum)
nthread = 128
block_C = sch.get_block("C")
i,  = sch.get_loops(block=block_C)
i0, i1 = sch.split(i, [None, nthread])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
sch.mod.show()

拆分循环后 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1027, "float32"], B: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                # 启用 8 个 block 并发计算，每个 block 用 16 个 thread 并发
                # 因此每一个 thread 只需要计算 1 次乘加
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A[vi : vi + 3])
                    T.writes(B[vi])
                    B[vi] = A[vi] + A[vi + 1] + A[vi + 2]

提前缓存数据

1
2
3

A_shared = sch.cache_read(block_C, read_buffer_index=0, storage_scope="shared")
sch.compute_at(A_shared, i1)
sch.mod.show()

提前缓存数据后的 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1027, "float32"], B: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        A_shared = T.alloc_buffer([1027], dtype="float32", scope="shared")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                # 由上图 GPU 结构图可知
                # 不同 block 无法共享 share memory
                # 相同 block 的不同 thread 之间可以共享
                # 所以输出 128 个结果需要 130 个输入（本行 128 个加下一行 2 个）
                for ax0 in T.serial(130):
                    with T.block("A_shared"):
                        v0 = T.axis.spatial(1027, i_0 * 128 + ax0)
                        T.reads(A[v0])
                        T.writes(A_shared[v0])
                        A_shared[v0] = A[v0]
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A_shared[vi : vi + 3])
                    T.writes(B[vi])
                    B[vi] = A_shared[vi] + A_shared[vi + 1] + A_shared[vi + 2]

缓存数据可以使用多线程优化
- 因为内存是跨线程共享的，所以需要重新拆分循环并将获取过程的内部迭代器绑定到线程索引上，这种技术称为 cooperative fetching
1
2
3
4
ax = sch.get_loops(A_shared)[-1]
ax0, ax1 = sch.split(ax, [None, nthread])
sch.bind(ax1, "threadIdx.x")
sch.mod.show()

缓存数据优化后 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1026, "float32"], B: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        A_shared = T.alloc_buffer([1026], dtype="float32", scope="shared")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                for ax0_0 in T.serial(2):
                    for ax0_1 in T.thread_binding(128, thread="threadIdx.x"):
                        with T.block("A_shared"):
                            # 由上图 GPU 结构图可知
                            # 不同 block 无法共享 share memory
                            # 相同 block 的不同 thread 之间可以共享
                            # 所以输出 128 个结果需要 130 个输入（本行 128 个加下一行 2 个）
                            v0 = T.axis.spatial(1026, i_0 * 128 + (ax0_0 * 128 + ax0_1))
                            T.where(ax0_0 * 128 + ax0_1 < 130)
                            T.reads(A[v0])
                            T.writes(A_shared[v0])
                            A_shared[v0] = A[v0]
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A_shared[vi : vi + 3])
                    T.writes(B[vi])
                    B[vi] = A_shared[vi] + A_shared[vi + 1] + A_shared[vi + 2]

矩阵乘法 GPU 加速

IRModule 基础实现：

@tvm.script.ir_module
class MyModuleMatmul:
    @T.prim_func
    def main(A: T.Buffer[(1024, 1024), "float32"],
            B: T.Buffer[(1024, 1024), "float32"],
            C: T.Buffer[(1024, 1024), "float32"]) -> None:
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i, j, k in T.grid(1024, 1024, 1024):
            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]

绑定 block 和 thread + 本地存储分块优化

循环拆分，来增加整体内存复用，只需要从 A 和 B 加载一次条形数据（上图中的灰色部分），然后使用它们来计算矩阵乘法结果
下面代码中设置 V = 8

def blocking(sch,
            tile_local_y,
            tile_local_x,
            tile_block_y,
            tile_block_x,
            tile_k):
    block_C = sch.get_block("C")
    C_local = sch.cache_write(block_C, 0, "local")
    i, j, k = sch.get_loops(block=block_C)
    i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])
    j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])
    k0, k1 = sch.split(loop=k, factors=[None, tile_k])
    sch.unroll(k1)
    sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)
    sch.reverse_compute_at(C_local, j1)
    sch.bind(i0, "blockIdx.y")
    sch.bind(j0, "blockIdx.x")
    sch.bind(i1, "threadIdx.y")
    sch.bind(j1, "threadIdx.x")
    sch.decompose_reduction(block_C, k0)
    return sch
sch = tvm.tir.Schedule(MyModuleMatmul)
sch = blocking(sch, 8, 8, 8, 8, 4)
sch.mod.show()

输出优化后的 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        C_local = T.alloc_buffer([1024, 1024], dtype="float32", scope="local")
        for i_0 in T.thread_binding(16, thread="blockIdx.y"):
            for j_0 in T.thread_binding(16, thread="blockIdx.x"):
                for i_1 in T.thread_binding(8, thread="threadIdx.y"):
                    for j_1 in T.thread_binding(8, thread="threadIdx.x"):
                        # 一共使用 16 * 16 个 block 并发计算
                        # 每个 block 使用 8 * 8 个 thread 并发
                        # 所以每个 thread 只需计算输出为 8 * 8 的区域，因此只需要加载 A 中 8 行和 B 中 8 列数据
                        # 1. 初始化 8 * 8 的输出区域为 0
                        for i_2_init, j_2_init in T.grid(8, 8):
                            with T.block("C_init"):
                                vi = T.axis.spatial(1024, i_0 * 64 + i_1 * 8 + i_2_init)
                                vj = T.axis.spatial(1024, j_0 * 64 + j_1 * 8 + j_2_init)
                                T.reads()
                                T.writes(C_local[vi, vj])
                                C_local[vi, vj] = T.float32(0)
                        
                        # 2. 计算 8 * 8 输出区域的值，共计算 8 * 8 * 1024 次乘加
                        for k_0 in T.serial(256):
                            for k_1 in T.unroll(4):
                                for i_2, j_2 in T.grid(8, 8):
                                    with T.block("C_update"):
                                        vi = T.axis.spatial(1024, i_0 * 64 + i_1 * 8 + i_2)
                                        vj = T.axis.spatial(1024, j_0 * 64 + j_1 * 8 + j_2)
                                        vk = T.axis.reduce(1024, k_0 * 4 + k_1)
                                        T.reads(C_local[vi, vj], A[vi, vk], B[vk, vj])
                                        T.writes(C_local[vi, vj])
                                        C_local[vi, vj] = C_local[vi, vj] + A[vi, vk] * B[vk, vj]
                        
                        # 3. 把每个 thread 的 8 * 8 的输出区域拼成最后的 1024 * 1024 的输出
                        for ax0, ax1 in T.grid(8, 8):
                            with T.block("C_local"):
                                v0 = T.axis.spatial(1024, i_0 * 64 + i_1 * 8 + ax0)
                                v1 = T.axis.spatial(1024, j_0 * 64 + j_1 * 8 + ax1)
                                T.reads(C_local[v0, v1])
                                T.writes(C[v0, v1])
                                C[v0, v1] = C_local[v0, v1]

共享内存优化

与上图不同，图中矩阵 C 中 L * L 灰色区域表示一个 block 的计算输出
每个 L * L 灰色区域由多个 V * V 的小区域组成，表示一个 thread 的输出

同一个 block 中的多个 thread 可共享内存，因此可以重排同一个 block 中的 thread 数据，使得尽可能少的数据缓存到 shared memory 中
优化前：
- 每个 thread 需要计算输出矩阵中 8 * 8 的数据，需要从 local memory 中读取 8 * 8 * 1024 * 2 数据
- 每个 block 中的 thread 之间没有数据共享，所以需要从 local memory 中读取 $8 * 8 * 8 * 8 * 1024 * 2 = 2^{23}$ 个矩阵元素
优化后：
- 每个 block 计算输出矩阵的 64 * 64 的数据最少需要 $64 * 1024 * 2=2^{17}$ 的数据，可提前将这部分数据缓存到 shared memory
- 然后每个 thread 从 shared memory 读数据计算，需读取 $64 * 1024 * 2=2^{17}$ 个数据
内存优化前后每个 block 读取数据对比：
- 优化前：从 local memory 读取 $2^{23}$ 个矩阵元素
- 优化后：从 local memory 读取 $2^{17}$ 个矩阵元素到 shared memory，再从 shared memory 读取 $2^{17}$ 个数据计算
优化过程：

def cache_read_and_coop_fetch(sch, block, nthread, read_idx, read_loc):
    read_cache = sch.cache_read(block=block, read_buffer_index=read_idx, storage_scope="shared")
    sch.compute_at(block=read_cache, loop=read_loc)
    # vectorized cooperative fetch
    inner0, inner1 = sch.get_loops(block=read_cache)[-2:]
    inner = sch.fuse(inner0, inner1)
    _, tx, vec = sch.split(loop=inner, factors=[None, nthread, 4])
    sch.vectorize(vec)
    sch.bind(tx, "threadIdx.x")
def blocking_with_shared(
    sch,
    tile_local_y,
    tile_local_x,
    tile_block_y,
    tile_block_x,
    tile_k):
    block_C = sch.get_block("C")
    C_local = sch.cache_write(block_C, 0, "local")
    i, j, k = sch.get_loops(block=block_C)
    i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])
    j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])
    k0, k1 = sch.split(loop=k, factors=[None, tile_k])
    sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)
    sch.reverse_compute_at(C_local, j1)
    sch.bind(i0, "blockIdx.y")
    sch.bind(j0, "blockIdx.x")
    tx = sch.fuse(i1, j1)
    sch.bind(tx, "threadIdx.x")
    nthread = tile_block_y * tile_block_x
    cache_read_and_coop_fetch(sch, block_C, nthread, 0, k0)
    cache_read_and_coop_fetch(sch, block_C, nthread, 1, k0)
    sch.decompose_reduction(block_C, k0)
    return sch
sch = tvm.tir.Schedule(MyModuleMatmul)
sch = blocking_with_shared(sch, 8, 8, 8, 8, 8)
sch.mod.show()

优化后 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        C_local = T.alloc_buffer([1024, 1024], dtype="float32", scope="local")
        A_shared = T.alloc_buffer([1024, 1024], dtype="float32", scope="shared")
        B_shared = T.alloc_buffer([1024, 1024], dtype="float32", scope="shared")
        for i_0 in T.thread_binding(16, thread="blockIdx.y"):
            for j_0 in T.thread_binding(16, thread="blockIdx.x"):
                for i_1_j_1_fused in T.thread_binding(64, thread="threadIdx.x"):
                    for i_2_init, j_2_init in T.grid(8, 8):
                        with T.block("C_init"):
                            vi = T.axis.spatial(1024, i_0 * 64 + i_1_j_1_fused // 8 * 8 + i_2_init)
                            vj = T.axis.spatial(1024, j_0 * 64 + i_1_j_1_fused % 8 * 8 + j_2_init)
                            T.reads()
                            T.writes(C_local[vi, vj])
                            C_local[vi, vj] = T.float32(0)
                    for k_0 in T.serial(128):
                        for ax0_ax1_fused_0 in T.serial(2):
                            for ax0_ax1_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
                                for ax0_ax1_fused_2 in T.vectorized(4):
                                    with T.block("A_shared"):
                                        v0 = T.axis.spatial(1024, i_0 * 64 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 8)
                                        v1 = T.axis.spatial(1024, k_0 * 8 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 8)
                                        T.reads(A[v0, v1])
                                        T.writes(A_shared[v0, v1])
                                        A_shared[v0, v1] = A[v0, v1]
                        for ax0_ax1_fused_0 in T.serial(2):
                            for ax0_ax1_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
                                for ax0_ax1_fused_2 in T.vectorized(4):
                                    with T.block("B_shared"):
                                        v0 = T.axis.spatial(1024, k_0 * 8 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 64)
                                        v1 = T.axis.spatial(1024, j_0 * 64 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 64)
                                        T.reads(B[v0, v1])
                                        T.writes(B_shared[v0, v1])
                                        B_shared[v0, v1] = B[v0, v1]
                        for k_1, i_2, j_2 in T.grid(8, 8, 8):
                            with T.block("C_update"):
                                vi = T.axis.spatial(1024, i_0 * 64 + i_1_j_1_fused // 8 * 8 + i_2)
                                vj = T.axis.spatial(1024, j_0 * 64 + i_1_j_1_fused % 8 * 8 + j_2)
                                vk = T.axis.reduce(1024, k_0 * 8 + k_1)
                                T.reads(C_local[vi, vj], A_shared[vi, vk], B_shared[vk, vj])
                                T.writes(C_local[vi, vj])
                                C_local[vi, vj] = C_local[vi, vj] + A_shared[vi, vk] * B_shared[vk, vj]
                    for ax0, ax1 in T.grid(8, 8):
                        with T.block("C_local"):
                            v0 = T.axis.spatial(1024, i_0 * 64 + i_1_j_1_fused // 8 * 8 + ax0)
                            v1 = T.axis.spatial(1024, j_0 * 64 + i_1_j_1_fused % 8 * 8 + ax1)
                            T.reads(C_local[v0, v1])
                            T.writes(C[v0, v1])
                            C[v0, v1] = C_local[v0, v1]

程序自动变换

from tvm import meta_schedule as ms
sch_tuned = ms.tune_tir(
    mod=MyModuleMatmul,
    target="nvidia/tesla-p100",
    config=ms.TuneConfig(
      max_trials_global=64,
      num_trials_per_iter=64,
    ),
    work_dir="./tune_tmp",
    task_name="main"
)
sch_tuned.mod.show()

机器学习编译(5)——与机器学习框架的整合

Posted on 2022-10-30 Edited on 2025-03-21 In 机器学习编译 Valine:

URL

https://mlc.ai/zh/chapter_integration/index.html

使用 `Builder` 创建 `IRModule`

从张量表达式创建 `TensorIR`（主张量函数）

from tvm import te
# 定义 TensorIR 输入
A = te.placeholder((128, 128), name="A", dtype="float32")
B = te.placeholder((128, 128), name="B", dtype="float32")
type(A)
# tvm.te.tensor.Tensor
A.shape
# [128, 128]
# 由张量表达式自动生成 TensorIR
def te_matmul(A: te.Tensor, B: te.Tensor) -> te.Tensor:
    assert A.shape[1] == B.shape[0]
    n = A.shape[0]
    m = B.shape[1]
    k = te.reduce_axis((0, A.shape[1]), name="k")
    # 由张量表达式自动生成 TensorIR
    # 调用格式是：te.compute(output_shape, lambda, TensorIR_name)
    return te.compute(
        (n, m), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="matmul"
    )
C = te_matmul(A, B)
# 打印自动生成的 TensorIR，函数输入即为 [A, B, C]
te.create_prim_func([A, B, C]).show()

输出（自动生成的主张量函数）

# from tvm.script import tir as T
@T.prim_func
def func(
    A: T.Buffer[(128, 128), "float32"],
    B: T.Buffer[(128, 128), "float32"],
    matmul: T.Buffer[(128, 128), "float32"],
) -> None:
    # function attr dict
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    # body
    # with T.block("root")
    for i0, i1, i2 in T.grid(128, 128, 128):
        with T.block("matmul"):
            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
            T.reads(A[i, k], B[k, j])
            T.writes(matmul[i, j])
            with T.init():
                matmul[i, j] = T.float32(0)
            matmul[i, j] = matmul[i, j] + A[i, k] * B[k, j]

使用 `BlockBuilder` 构造 `IRModule`

自动生成的主张量函数还需要 计算图抽象 来将计算图拼起来

A = relax.Var("A", (128, 128), relax.DynTensorType(2, "float32"))
B = relax.Var("B", (128, 128), relax.DynTensorType(2, "float32"))
# 使用 BlockBuilder 将多个张量函数拼接成一个 IRModule
bb = relax.BlockBuilder()
with bb.function("main"):
    with bb.dataflow():
        C = bb.emit_te(te_matmul, A, B)
        D = bb.emit_te(te_relu, C)
        R = bb.emit_output(D)
    bb.emit_func_output(R, params=[A, B])
MyModule = bb.get()
MyModule.show()

输出（IRModule）

tvm.script.ir_module
class Module:
    @T.prim_func
    def te_matmul(rxplaceholder: T.Buffer[(128, 128), "float32"], rxplaceholder_1: T.Buffer[(128, 128), "float32"], matmul: T.Buffer[(128, 128), "float32"]) -> None:
        ...
    @T.prim_func
    def te_relu(rxplaceholder: T.Buffer[(128, 128), "float32"], relu: T.Buffer[(128, 128), "float32"]) -> None:
        ...
    @R.function
    def main(A: Tensor((128, 128), "float32"), B: Tensor((128, 128), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv = R.call_tir(te_matmul, (A, B), (128, 128), dtype="float32")
            lv1 = R.call_tir(te_relu, (lv,), (128, 128), dtype="float32")
            gv: Tensor((128, 128), "float32") = lv1
            R.output(gv)
        return gv

使用 BlockBuilder 创建 IRModule 与直接创建 IRMoudle 的对比
bb.emit_te 做了以下事情：
- 为 A 和 B 创建一个输入 te.placeholder
- 通过 te_matmul 函数运行它们
- 调用 te.create_prim_func 来创建一个 TensorIR 函数
- 通过 call_tir 生成对函数的调用

从 `Pytorch` 映射到 `IRModule`

`Pytorch` 模型

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(128, 128))
    def forward(self, x):
        x = torch.matmul(x, self.weight)
        x = torch.relu(x)
        return x
model = MyModel()
# 生成 Pytorch 计算图
fx_module = fx.symbolic_trace(model)

构造计算图之间的映射变换

# pytorch module parameter to IRModule parameter
def map_param(param: nn.Parameter):
    ndim = len(param.data.shape)
    return relax.const(
        param.data.cpu().numpy(), relax.DynTensorType(ndim, "float32")
    )
# pytorch module attribute to IRModule attribute
def fetch_attr(fx_mod, target: str):
    """Helper function to fetch an attr"""
    target_atoms = target.split('.')
    attr_itr = fx_mod
    for i, atom in enumerate(target_atoms):
        if not hasattr(attr_itr, atom):
            raise RuntimeError(f"Node referenced nonexistant target {'.'.join(target_atoms[:i])}")
        attr_itr = getattr(attr_itr, atom)
    return attr_itr
def from_fx(fx_mod, input_shapes, call_function_map, call_module_map):
    input_index = 0
    node_map = {}
    named_modules = dict(fx_mod.named_modules())
    bb = relax.BlockBuilder()
    fn_inputs = []
    fn_output = None
    with bb.function("main"):
        with bb.dataflow():
            for node in fx_mod.graph.nodes:
                if node.op == "placeholder":
                    # create input placeholder
                    shape = input_shapes[input_index]
                    input_index += 1
                    input_var = relax.Var(
                        node.target, shape, relax.DynTensorType(len(shape), "float32")
                    )
                    fn_inputs.append(input_var)
                    node_map[node] = input_var
                elif node.op == "get_attr":
                    node_map[node] = map_param(fetch_attr(fx_mod, node.target))
                elif node.op == "call_function":
                    node_map[node] = call_function_map[node.target](bb, node_map, node)
                elif node.op == "call_module":
                    named_module = named_modules[node.target]
                    node_map[node] = call_module_map[type(named_module)](bb, node_map, node, named_module)
                elif node.op == "output":
                    output = node_map[node.args[0]]
                    assert fn_output is None
                    fn_output = bb.emit_output(output)
        # output and finalize the function
        bb.emit_func_output(output, fn_inputs)
    return bb.get()

映射 `Pytorch Module` 到 `TensorIR`

# TensorIR 映射变换
def map_matmul(bb, node_map, node: fx.Node):
    A = node_map[node.args[0]]
    B = node_map[node.args[1]]
    return bb.emit_te(te_matmul, A, B)
# TensorIR 映射变换
def map_relu(bb, node_map, node: fx.Node):
    A = node_map[node.args[0]]
    return bb.emit_te(te_relu, A)
MyModule = from_fx(
    fx_module,
    input_shapes = [(1, 128)],
    call_function_map = {
      torch.matmul: map_matmul,
      torch.relu: map_relu,
    },
    call_module_map={},
)
MyModule.show()

映射后的 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def te_matmul(rxplaceholder: T.Buffer[(1, 128), "float32"], rxplaceholder_1: T.Buffer[(128, 128), "float32"], matmul: T.Buffer[(1, 128), "float32"]) -> None:
        ...
    @T.prim_func
    def te_relu(rxplaceholder: T.Buffer[(1, 128), "float32"], relu: T.Buffer[(1, 128), "float32"]) -> None:
        ...
    @R.function
    def main(x: Tensor((1, 128), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv = R.call_tir(te_matmul, (x, meta[relay.Constant][0]), (1, 128), dtype="float32")
            lv1 = R.call_tir(te_relu, (lv,), (1, 128), dtype="float32")
            gv: Tensor((1, 128), "float32") = lv1
            R.output(gv)
        return lv1

或映射到 `Pytorch Module` 到 `IRModule` 更高层的算子

def map_nn_relu_op(bb, node_map, node, nn_mod):
    A = node_map[node.args[0]]
    return bb.emit(relax.op.relu(A))
def map_nn_linear_op(bb, node_map, node, nn_mod):
    x = node_map[node.args[0]]
    w = map_param(nn_mod.weight)
    if nn_mod.bias is not None:
        b = map_param(nn_mod.bias)
    y = bb.emit(relax.op.dense(x, w))
    return bb.emit(relax.op.add(y, b))
MLPModuleHighLevel = from_fx(
    fx.symbolic_trace(mlp_model),
    input_shapes = [(1, 784)],
    call_function_map={
    },
    call_module_map={
        torch.nn.Linear: map_nn_linear_op,
        torch.nn.ReLU: map_nn_relu_op,
    },
)
MLPModuleHighLevel.show()

输出

@tvm.script.ir_module
class Module:
    @R.function
    def main(x: Tensor((1, 784), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv: Tensor((1, 128), "float32") = relax.nn.dense(x, meta[relay.Constant][0])
            lv1: Tensor((1, 128), "float32") = relax.add(lv, meta[relay.Constant][1])
            lv2: Tensor((1, 128), "float32") = relax.nn.relu(lv1)
            lv3: Tensor((1, 10), "float32") = relax.nn.dense(lv2, meta[relay.Constant][2])
            lv4: Tensor((1, 10), "float32") = relax.add(lv3, meta[relay.Constant][3])
            gv: Tensor((1, 10), "float32") = lv4
            R.output(gv)
        return lv4

总结

张量表达式 API 允许我们创建原始的 TensorIR 函数
BlockBuilder API 通过 emit_te 和其他函数创建 IRModule
通过将模型转换为 IRModule，实现与现有的机器学习框架的整合

机器学习编译(4)——自动程序优化

Posted on 2022-10-22 Edited on 2025-03-21 In 机器学习编译 Valine:

URL

https://mlc.ai/zh/chapter_auto_program_optimization/index.html

自动程序优化的原因

MLC 的本质是张量函数之间的转换，但我们不知道哪种转换是让模型运行更快的，所以需要使用自动程序优化，去自动搜索最有转换。

自动程序优化过程

end-to-end 构建模型

# IR_Module 使用自定义主张量函数和库张量函数
@tvm.script.ir_module
class MyModuleMixture:
    @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})
        ...
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        with R.dataflow():
            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir("env.relu", (lv0,), (1, 128), dtype="float32")
            out = R.call_tir("env.linear", (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out
# 注册库张量函数
@tvm.register_func("env.linear", override=True)
def torch_linear(
    x: tvm.nd.NDArray, w: tvm.nd.NDArray, b: tvm.nd.NDArray, out: tvm.nd.NDArray
):
    ...
# 注册库张量函数
@tvm.register_func("env.relu", override=True)
def lnumpy_relu(x: tvm.nd.NDArray, out: tvm.nd.NDArray):
    ...
# 绑定模型权重参数（nd_params 是模型权重），作用类似于 functools.partial()
MyModuleWithParams = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
# IR_Module -> 可执行程序 -> 虚拟机执行器
ex = relax.vm.build(MyModuleWithParams, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
# 执行
nd_res = vm["main"](data_nd)
# 测速
ftimer = vm.module.time_evaluator("main", tvm.cpu(), number=100)

自动优化 linear0 主张量函数

# 调优 API 只接受一个带有一个 main 函数的 IRModule，所以需要将原始 IRModule 中的 linear0 转成新 IRModule 的 main 函数
mod_linear = tvm.IRModule.from_expr(MyModuleMixture["linear0"].with_attr("global_symbol", "main"))
# 打印新IRModule
IPython.display.HTML(code2html(mod_linear.script()))
# 打印输出
@tvm.script.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"],
    ):
        # 函数中内容是 MyModuleMixture.linear0
# 自动调优 API，input 是需要调优的 IRModule，output 是调优后的 schedule，schedule.mod 是调优后的 IRModule
sch_tuned_linear = ms.tune_tir(
    mod=mod_linear,                 # 待调优 IRModule
    target="llvm --num-cores=1",    # 调优目标
    config=ms.TuneConfig(           # 自动调优配置
      max_trials_global=64,
      num_trials_per_iter=64,
    ),
    work_dir="./tune_tmp",
    task_name="main",
)
# 将返回的 IRModule 中的 main 函数更新到原 IRModule 的 linear0 中
# 绑定参数
MyModuleWithParams2 = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
# 获取调优后的 main 函数
new_func = sch_tuned_linear.mod["main"].with_attr("global_symbol", "linear0")
# 获取原 IRModule 的 linear0 张量函数
gv = MyModuleWithParams2.get_global_var("linear0")
# 更新调优后的 main 函数到原 IRModule 的 linear0 张量函数
MyModuleWithParams2.update_func(gv, new_func)
# 重新测速，速度变快

机器学习编译(3)——端到端模型执行

Posted on 2022-10-21 Edited on 2025-03-21 In 机器学习编译 Valine:

URL

https://mlc.ai/zh/chapter_end_to_end/index.html#

机器学习编译的本质与关注点

MLC 的本质：张量函数之间的转换
MLC 的关注点：
- 所以可能的张量函数抽象表达
- 所有可能的张量函数转换

构造 `IR_Module`

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def relu0(X: T.Buffer[(1, 128), "float32"], Y: T.Buffer[(1, 128), "float32"]):
        # function attr dict
        T.func_attr({"global_symbol": "relu0", "tir.noalias": True})
        for i, j in T.grid(1, 128):
            with T.block("Y"):
                vi, vj = T.axis.remap("SS", [i, j])
                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))
    @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})
        Y = T.alloc_buffer((1, 128), "float32")
        for i, j, k in T.grid(1, 128, 784):
            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] + 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])
                Z[vi, vj] = Y[vi, vj] + B[vj]
    @T.prim_func
    def linear1(
        X: T.Buffer[(1, 128), "float32"],
        W: T.Buffer[(10, 128), "float32"],
        B: T.Buffer[(10,), "float32"],
        Z: T.Buffer[(1, 10), "float32"],
    ):
        T.func_attr({"global_symbol": "linear1", "tir.noalias": True})
        Y = T.alloc_buffer((1, 10), "float32")
        for i, j, k in T.grid(1, 10, 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] + X[vi, vk] * W[vj, vk]
        for i, j in T.grid(1, 10):
            with T.block("Z"):
                vi, vj = T.axis.remap("SS", [i, j])
                Z[vi, vj] = Y[vi, vj] + B[vj]
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        with R.dataflow():
            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir(relu0, (lv0,), (1, 128), dtype="float32")
            out = R.call_tir(linear1, (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out

@tvm.script.ir_module 装饰 IR_Module
@T.prim_func 装饰 主张量函数
@R.function 装饰 高层神经网络执行的抽象，（将整个 IR_Module 中的主张量函数串起来组成一个计算图）
R.dataflow() 用于标记程序计算图区域
R.call_tir(prim_func, inputs, output_shape, dtype) 分配输出内存并和输入一起输入主张量函数

构建并运行模型

ex = relax.vm.build(MyModule, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
nd_res = vm["main"](
    data_nd, nd_params["w0"], nd_params["b0"], nd_params["w1"], nd_params["b1"]
)

可执行文件 = build(IR_Module)
虚拟机执行器 = 虚拟机(可执行文件)
运行结果 = 虚拟机执行器(模型输入)

使用现有库避免重复造轮子

@tvm.script.ir_module
class MyModuleWithExternCall:
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        # block 0
        with R.dataflow():
            lv0 = R.call_tir("env.linear", (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir("env.relu", (lv0,), (1, 128), dtype="float32")
            out = R.call_tir("env.linear", (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out

上图中的 env.linear 是库张量函数，同一个 IR_Module 中可使用库张量函数，也可使用自定义张量函数，也可以二者混用。

1
2
3

@tvm.register_func("env.linear", override=True)
def torch_linear(x: tvm.nd.NDArray, w: tvm.nd.NDArray, b: tvm.nd.NDArray, out: tvm.nd.NDArray):
    ...

@tvm.register_func，注册库张量函数

绑定参数到 IR_Module

以上的所有 IR_Module 都是在调用时才传入数据和 权重参数，但 权重参数 是不变的，所以可以将 权重参数 提前绑定到 IR_Module 中，调用时只传入输入数据。

1	MyModuleWithParams = relax.transform.BindParams("main", nd_params)(MyModuleMixture)

relax.transform.BindParams(计算图入口函数，模型参数)(IR_Module) 将模型参数绑定到 IR_Module 的计算图入口函数中，返回一个绑定好模型参数的 IR_Module。

总结

计算图抽象（一般表示 main 函数）有助于将元张量函数拼接在一起以进行端到端执行。
Relax 抽象的关键要素包括
- call_tir 构造，将目标传递规范的元函数嵌入到计算图中
- Dataflow block
计算图允许调用环境库函数和 TensorIR 函数。

相机内参与外参

Posted on 2022-10-16 Edited on 2025-03-21 In 3D Valine:

三种成像相关坐标系

像素坐标系：一种 2D 坐标，以图片左上角为原点，横轴（宽度方向）向右为 x 轴正方向，纵轴（高度方向）向下为 y 轴正方向，单位是像素
相机坐标系：一种 3D 坐标，以相机光心为原点，垂直相机平面远离相机方向为 z 轴正方向，垂直于 z 轴且平行于相机平面，水平向右为 x 轴正方向，竖直向下为 y 轴正方向，单位是米
世界坐标系：一种 3D 坐标，一种人为定义的，且x, y, z 轴两两垂直的坐标系，单位是米

相机内参

相机内参可以实现像素坐标系与相机坐标系之间相互转换，通常使用一个 3 * 3 矩阵表示
根据小孔成像和相似三角形原理，可以得出相机坐标系与成像坐标系点的对应关系： $\frac{Z}{f}=\frac{X}{x}=\frac{Y}{y}$
其中： $(X,Y,Z)$ 为相机坐标系下的点的坐标， $(x, y)$ 为投影到成像平面上的点的坐标， $f$ 表示焦距。
再根据成像坐标系到像素坐标系的对应关系：
$
\begin{matrix}
u=\alpha \cdot x + c_x \
v = \beta \cdot y + c_y
\end{matrix}
$
其中：
- $\alpha,\beta$ 分别表示 $x, y$ 方向上成像宽度到像素宽度的投影
- 由于成像坐标系原点为成像中心，像素坐标系原点为像素左上角，所以需要加上原点的偏移， $c_x,c_y$ 分别表示 $x, y$ 方向上原点的偏移。
所以，相机坐标系 $(X,Y,Z)$ 与像素坐标系 $(u, v)$ 可通过相机内参相互转换：
$
Z \begin{bmatrix} u \ v \ 1 \end{bmatrix} = \begin{bmatrix} f_x & 0 & c_x \ 0 & f_y & c_y \ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix}X\ Y\ Z \end{bmatrix}
$
其中： $f_x=\alpha \cdot f,\ f_y=\beta \cdot f,\ Z$ 表示相机坐标系下的深度
$\begin{bmatrix} f_x & 0 & c_x \\ 0 & f_y & c_y \\ 0 & 0 & 1 \end{bmatrix}$ 被称为相机内参 $K$

相机外参

相机外参可以实现相机坐标系与世界坐标系之间相互转换（刚体变换），通常用一个 3 * 3 的旋转矩阵 $R$ 和一个 3 * 1 的平移矩阵 $T$ 表示：
$
\begin{bmatrix} X_c\ Y_c \ Z_c \end{bmatrix}=\begin{bmatrix} R_{11} & R_{12} & R_{13} \ R_{21} & R_{22} & R_{23} \ R_{31} & R_{32} & R_{33} \end{bmatrix} \begin{bmatrix} X_w\ Y_w \ Z_w \end{bmatrix} + \begin{bmatrix} T_{1} \ T_{2} \ T_{3} \end{bmatrix}
$

其中： $(X_c,Y_c,Z_c)$ 表示相机坐标系下的点， $(X_w,Y_w,Z_w)$ 表示世界坐标系下的点
齐次化之后，得到一个 4 * 4 的矩阵：
$
\begin{bmatrix} X_c\ Y_c \ Z_c \ 1\end{bmatrix}=\begin{bmatrix} R_{11} & R_{12} & R_{13} & T_1 \ R_{21} & R_{22} & R_{23} & T_2 \ R_{31} & R_{32} & R_{33} & T_3 \ 0 & 0 & 0 & 1\end{bmatrix} \begin{bmatrix} X_w\ Y_w \ Z_w \ 1\end{bmatrix}
$

总结

$camera\ intrisic\ matrix \times camera\ coordinate\ system = depth\ in\ camera\ coordinate system \cdot image\ coordinate\ system$
$camera\ extrisic\ matrix \times world\ coordinate\ system = camera\ coordinate\ system$
$camera\ intrisic\ matrix \times camera\ extrisic\ matrix \times world\ coordinate\ system = depth\ in\ camera\ coordinate system \cdot image\ coordinate\ system$

Lift, Splat, Shoot: Encoding Images from Arbitrary Camera Rigs by Implicitly Unprojecting to 3D

Posted on 2022-09-30 Edited on 2025-03-21 In BEV Valine:

URL

paper: https://arxiv.org/pdf/2008.05711.pdf
code: https://github.com/real-zhangzhe/lift-splat-shoot (fork from official implementation, adding comments and predictions)

TL;DR

本文是 BEV (bird eye view) 的开山之作，通过隐式 2D 深度估计和像素坐标到世界坐标转换，将多张（6张）车周环视图拼接得到一张鸟瞰图。
具体实现请看代码，代码中有非常详细的注释。

Algorithm

在 inference 过程中，下半部分的 6 张环视图为输入，上半部分的鸟瞰图为输出（地图和本算法无关）。

Dataset

本文使用自动驾驶数据集 nuScense
- 官网：https://www.nuscenes.org/
- 数据集介绍论文：https://arxiv.org/pdf/1903.11027.pdf
输入的 6 张图来自上图的 6 个绿色 camera
世界坐标系如图 IMU 所示原点定为 车后轴中心，x 轴正方向为车辆前进方向， y 轴正方向为面向车辆前进方向的左手边，z 轴正方向为竖直向上。

算法细节

特别细节的看代码

1. 特征提取（Lift）

使用 EfficientNet 进行 2D 特征提取 + 隐式深度估计
- 输入：shape = [4, 6, 3, 128, 352]，分别表示：[batch, cameras, channel, height, width]
- 输出：shape = [24, 64, 41, 8, 22]，分别表示：[batch * cameras, features, depth, height, width]
  - 使用 64 维向量编码深度（不是直接预测深度，所以被称为隐式深度估计）
  - 深度从 4m ~ 45m，编码精度为 1m，所以有 41 种离散深度，相机坐标系下深度估计的目的是：从像素坐标转化为世界坐标
  - 长宽各下采样 16 倍，减小计算量

2. 像素坐标和相机坐标系下深度到世界坐标的映射（Splat）

使用如下参数将像素坐标和相机坐标系下深度映射到世界坐标
- 相机内参
- 相机外参
  - 旋转
  - 平移
- 像素坐标系内变换参数（缩放 + 裁剪（平移））
  - 原图(900, 1600) -> 模型输入图(128, 352) -> 模型预测图(8, 22)
体素池化：将属于同一个体素的深度估计向量求和
输入：
- 深度估计：shape = [24, 64, 41, 8, 22]
- 相机内外参和缩放参数
输出：shape = [4, 64, 200, 200]
- 200 * 200 个体素
  - X 方向上 [-50m, 50m) 0.5m 为一个 bin，200 个 bin
  - Z 方向上 [-50m, 50m) 0.5m 为一个 bin，200 个 bin
  - Y 方向不分 bin
- 每个体素用 64 维向量编码
本质是：
1. 构造一个 [24 * 41 * 8 * 22, 3] 的查找表，输入为 backbone 输出特征图的每一个 pixel，输出为这个 pixel 对应的世界坐标（这个查找表可由相机内外参和图像缩放系数计算得到）
2. 将离散的世界坐标点合并，合并规则是属于同一个体素的坐标点则合并

3. 体素编码降维（Shoot）

输入：shape = [4, 64, 200, 200]
输出：shape = [4, 1, 200, 200]，BEV 图

4. 训练 loss

~~非常简单粗暴~~

将 GT bbox3d 同样映射到 BEV 空间 [4, 1, 200, 200]，然后做 pixel-wise loss（分割 loss）

Thought

代码中 像素坐标和相机坐标系下深度到世界坐标的映射 部分比较难懂，需要较强的相机成像原理 / 图像 2D 转 3D 背景才能看懂

Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

Posted on 2022-09-29 Edited on 2025-03-21 In Transformer Valine:

URL

paper: https://arxiv.org/pdf/2103.14030.pdf
code: https://github.com/microsoft/Swin-Transformer

TL;DR

本文提出一种新型 Transformer 结构使用 Window Multi-head Self-attention (W-MSA) 和 Shifted Window Multi-head Self-attention (SW-MSA) 结构替代原始 Transformer 使用的 Multi-head Self-attention (MSA) 结构。
- 大大节省了原始 Transformer 的计算复杂度。
- 在视觉任务中吊打了一众 CNN 和 Transformer。

Algorithm

总体结构

swin-Transformer 总体按照类似 CNN 的层次化构建方式构建网络结构，分为 4 个 stage，每个 stage 都会将分辨率缩小一倍，channel 数扩大一倍 （like vgg） 。
Swin Tranformer Block 像大多数 Transformer Block 一样，不改变输入特征 shape，可以看做是一种比较高级（加入了 self-attention）的激活函数。
图中 Swin Tranformer Block 都是以连续偶数次出现，因为是一个 W-MSA Transformer Block + 一个 SW-MSA Transformer Block，如右边子图 b 所示。

Patch Partition

作用与 ViT 第一步很像：将图片切成若干等大小不重叠的 patch，patch_size = P，然后把每个 patch 从 (P, P, c) 拉成 1 维。
本实验中 patch_size = 4，所以一张图被裁切成了 $\frac{H}{4}\times\frac{W}{4}$ 个长度为 4 * 4 * 3 = 48 的向量。

Linear Embedding

与一个 shape = (48, C) 的矩阵乘将 48 维映射到 C 维。

与上一步结合可以变成一个 in_channel = 3, out_channel = C, kernel_size = stride = 4 的 Conv2d，官方实现中实际上也是这么做的（ class PatchEmbed ）。

Patch Merging

由两步组成，作用是 将分辨率缩小一倍，channel 扩大一倍。
- 将 $\frac{H}{4}\times\frac{W}{4}\times C$ 的输入使用 bayer2rggb (space2depth with block_size = 2) 变成 $\frac{H}{8}\times\frac{W}{8}\times 4C$ 。
- 再将 $\frac{H}{8}\times\frac{W}{8}\times 4C$ 与一个 $4C\times 2C$ 的矩阵相乘，输出一个 $\frac{H}{8}\times\frac{W}{8}\times 2C$ 的矩阵。

W-MSA

全称为 Windows Multi-head Self-attention，目的是为了减少 Self-attention 计算量。
- 原始的 MSA 会直接对完整输入计算其 self-attention 结果
- W-MSA 先将输入拆成大小为 $M\times M$ 且互不重叠的 windows，然后计算每个 window 的 self-attention 结果。
- 计算公式还是: $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$
MSA 和 W-MSA 计算量对比（假设输入特征图 $X \in \mathbb{R}^{H\times W\times C}$ $X \in R^{H \times W \times C}$ ，且 $W_q,W_k,W_v \in \mathbb{R}^{C\times C}$ $W_{q}, W_{k}, W_{v} \in R^{C \times C}$ ，且 Multi-head 的 head 数为 1)：
- MSA 计算量：
  - X -> Q / K / V 计算量： $\mathbb{R}^{H\times W\times C}\times \mathbb{R}^{C\times C}$ 计算量为 $3HWC^2$ 。
  - $QK^T$ 计算量： $\mathbb{R}^{HW\times C}\times \mathbb{R}^{C\times HW}$ 计算量为 $H^2W^2C$ 。
  - 不考虑 softmax 和 $\frac{..}{\sqrt{d_k}}$ 计算量。
  - softmax 结果 $\times V$ 计算量： $\mathbb{R}^{HW \times HW}\times \mathbb{R}^{HW \times C}$ 计算量为 $H^2W^2C$ 。
  - 因为需要 Multi-head，所以需要将 $\times V$ 之后的矩阵再 $\times W_o$ ，且 head 数为 1，计算量： $\mathbb{R}^{1*HW \times C}\times \mathbb{R}^{C \times C}$ 计算量为 $HWC^2$ 。
  - 总计算量： $4HWC^2 + 2H^2W^2C$ 。
- W-MSA 计算量：
  - 上图的 $H\rightarrow M,\ W\rightarrow M$ ，一共重复 $\frac{HW}{M^2}$ 次，所以总计算量为：
  $(4M^2C^2 + 2M^4C)\times \frac{HW}{M^2}=4HWC^2+2HWM^2C$
- 相比之下 W-MSA 会比 WSA 计算量少： $2HWC(HW-M^2)$ 。

SW-MSA

为了节省计算量 W-MSA 会将输入的完整特征图分 window，每个 window 独立去做 self-attention，这会导致 window 之间的关联性消失，这有悖于 self-attention 会在全图上建立长距离全局相关性依赖的特点。
所以，作者引入 SW-MSA 的方案通过滑窗解决这一问题。
滑窗会带来边角处的零碎区域（长或宽小于 window_size 的区域），由于 $H \% M = 0,\ W \% M = 0$ (W-MSA 要求)，所以零碎的区域可以通过调整位置拼成完整 window 。
完整的滑窗区域与 W-MSA 一样独立做 Self-attention。
由零碎区域拼成的滑窗区域中：
- 本属于同一零碎区域的位置可以 Self-attention 。
- 本不属于同一零碎区域的位置在原图上不相邻，不能做 Self-attention 。
- 具体做法是来自不同区域位置之间的 Self-attention 被 mask 掉 ，只保留来自相同区域的位置间的 Self-attention 。

W-MSA + SW-MSA

如网络结构图子图 b 所示，W-MSA 连接 SW-MSA 之后，数学表示：
$\hat{z}^l=W-MSA(LN(z^{l-1})) + z^{l-1}$
$z^l=MLP(LN(\hat z^l)) + \hat z^l$
$\hat z^{l+1} = SW-MSA(LN(z^l))+z^l$
$z^{l+1} = MLP(LN(\hat z^{l+1}))+\hat z^{l+1}$

Relative Position Bias

对 Self-attention 加上 相对位置偏置 bias 之后，精度提升巨大。
原始 Self-attention： $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$ 。
Relative Position Bias Self-attention： $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}} + Bias)V$ 。

Through

质疑了很多人都没有质疑的点，例如给 Self-attention 加上相对位置偏置。
本质是对 Tranformer 的一种很 work 的加速方法。
但 Shift window 逻辑比较复杂，没有经典模型该有的简约美，盲猜之后会被更 make sense 的结构替代。

Python 闭包与装饰器

Posted on 2022-09-24 Edited on 2025-03-21 In Python Valine:

Python 闭包 (closure)

闭包定义

闭包: 在一些语言中，在函数中可以（嵌套）定义另一个函数时，如果内部的函数引用了外部的函数的变量，则可能产生闭包。闭包可以用来在一个函数与一组“私有”变量之间创建关联关系。在给定函数被多次调用的过程中，这些私有变量能够保持其持久性。
支持将函数当成对象使用的编程语言，一般都支持闭包。比如 Python, JavaScript。

闭包的示例

代码

def function_1(arg_1):
    def function_2(arg_2):
        return arg_1 * arg_2
    return function_2
times_8 = function_1(8)
out = times_8(9)
print(f"times_8(9) = {out}")
# 闭包中的 cell
print(f"times_8.__closure__ = {times_8.__closure__}")
# 闭包中的 cell 对象的内容
print("times_8.__closure__.cell_contents:")
for i in times_8.__closure__:
    print(i.cell_contents)

输出

times_8(9) = 72
times_8.__closure__ = (<cell at 0x7ff39d4d2a30: int object at 0x5642a56c5e20>,)
times_8.__closure__.cell_contents:
8

闭包的用处

1. 可以读取函数内部的变量

如上面给出的例子

2. 让这些变量的值始终保持在内存中

例如一个棋盘游戏，棋子每次可以选择上下左右方向中的一个，在此方向上移动距离 step，使用闭包实现代码如下：

def create(pos=[0, 0]):
    
    def go(direction, step):
        new_x = pos[0] + direction[0] * step
        new_y = pos[1] + direction[1] * step
        
        pos[0] = new_x
        pos[1] = new_y
        
        return pos
    
    
    return go
player = create()
print(player([1, 0], 10))
print(player([0, 1], 20))
print(player([-1, 0], 10))

输出
1
2
3
[10, 0]
[10, 20]
[0, 20]

棋子每次更新后的位置都会存储在闭包中。

3. 用于装饰器

可以读取函数内部的变量 和 让这些变量的值始终保持在内存中 都可以使用 Python 的类实现，但 装饰器 是闭包的一个典型用处。

装饰器 (Decorators)

装饰器的定义

装饰器: 由闭包的概念引申而来，是一种 增加函数或类功能的方法，它可以快速地给不同的函数或类传入相同的功能。

函数装饰器的示例

代码

import time
def count_time(some_fun):
    def wrapper():
        t1 = time.time()
        some_fun()
        print(f"运行时间为: {round(time.time() - t1, 2)} s")
    return wrapper
# 装饰器语法糖
@count_time
def function_1():
    time.sleep(1)
    print("run function_1")
function_1()
# 不使用语法糖的方法
def function_2():
    time.sleep(1)
    print("run function_2")
new_function = count_time(function_2)
new_function()

输出

run function_1
运行时间为: 1.0 s
run function_2
运行时间为: 1.0 s

类别装饰器示例

代码

import time
class Timer:
    def __init__(self, func) -> None:
        self.func = func
    def __call__(self, *args, **kwargs):
        start = time.time()
        ret = self.func(*args, **kwargs)
        print(f"Time: {time.time()-start}")
        return ret
# 使用装饰器语法糖实现
@Timer
def add_1(a, b):
    time.sleep(1)
    print(f"{a} + {b} = {a+b}")
add_1(2, 3)
# 不使用装饰器语法糖实现
def add_2(a, b):
    time.sleep(1)
    print(f"{a} + {b} = {a+b}")
new_add_2 = Timer(add_2)
new_add_2(2, 3)

输出

2 + 3 = 5
Time: 1.0011768341064453
2 + 3 = 5
Time: 1.001098394393921

A White Paper on Neural Network Quantization

Posted on 2022-09-19 Edited on 2025-03-21 In Quantization Valine:

URL

https://arxiv.org/pdf/2106.08295.pdf

TL;DR

本文用比较简洁的方式给出了神经网络的通用量化方法，是量化领域的必读论文。

Algorithm

1. 量化基础知识

1.1 硬件背景

一个 $y=Wx+b$ 实际上是由 乘法器 和 累加器 组合而成的，实际的计算过程如下：

卷积实际上也是通过 image to column 操作变成 $y=Wx+b$ 操作

常见的 int8 量化会将上述过程变成如下过程：

weight 和 input 都被量化为 int8 ，同时保留各自的量化 scale，乘法操作是整形乘法器（更快），累加器是 int32 类型，最后再量化为 int8 放到 OCM 上

1.2 均匀仿射量化

均匀仿射量化也被称为 非对称量化，由三个量化参数定义：
- 比例因子 scale
- 零点 zero_point
- 比特宽度 bits
非对称量化：
- for unsigned integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil+z;0,2^b-1)$
- for signed integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil+z;-2^{b-1},2^{b-1}-1)$
- 这里的 $\lfloor\rceil$ 表示 round 运算
对称量化是非对称量化的简化版本，是将零点 zero_point 固定为 0
对称量化：
- for unsigned integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil;0,2^b-1)$
- for signed integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil;-2^{b-1},2^{b-1}-1)$
对称量化和非对称量化的含义：
2 的指数幂量化：
- 限制 $s=2^{-k}$
- 优势：scale 过程变成了硬件移位，对硬件更友好。
- 劣势：会使得 round 和 clip 误差的权衡变难。
量化颗粒度：
- per-tensor: 硬件更友好，但限制了量化的自由度。
- per-channel: 反之。

1.3 量化模拟

量化模拟是指在浮点计算设备上模拟定点计算设备的过程，通常用于训练。

左边是定点计算过程，右边是用浮点设备模型定点计算的过程

为了减少数据搬运和不必要的量化步骤，通常会做：
- batch norm 折叠：batch norm 在推理时是静态的，因此可以和前面的 conv 等层合并。
- 激活函数融合：在实际的硬件解决方案中，通常会在非线性操作（如 ReLU）之后直接进行量化，而不是先将激活写入内存然后再加载回计算核心。

1.4 实践考量

对称量化和非对称量化：
- 对称量化：zero-point == 0
- 非对称量化：zero-point != 0
为了方便计算，通常情况下，会将权重设置为对称量化（ $z_w=0$ ），将特征设置为非对称量化（ $z_x\ne 0$ ）
- 原因分析：
  - $W=S_w(W_{int} - Z_w)$
  - $X=S_x(X_{int} - Z_x)$
  - $WX=S_wS_x(W_{int} - Z_w)(X_{int} - Z_x)\\=S_wS_xW_{int}X_{int}-S_wS_xZ_wX_{int}-S_wS_xZ_xW_{int}+S_wS_xZ_wZ_x$
  - 在推理阶段： $S_w,\ S_x,\ Z_w,\ Z_x,\ W_{int}$ $S_{w}, S_{x}, Z_{w}, Z_{x}, W_{i n t}$ 已知，因此：
    - 等式的第三项和第四项可提前算出，无需推理耗时。
    - 第一项和第二项由于关联动态输入 $X_{int}$ ，因此需要额外耗时；但是如果设置 $Z_w=0$ ，则第二项恒等于0，可节省计算量。

2. 训练后量化（PTQ，post-training quantization）

训练后量化是指用 float32 精度训练的模型直接转成量化模型，无需任何数据和训练。

2.1 量化范围的设置

最大最小值法（min-max）： $q_{min}=minV,\ \ q_{max}=maxV$ ， $V$ 是待量化 tensor
均方差法（MSE）： $\argmin_{q_{min},q_{max}}||V-\hat{V}(q_{min}, q_{max})||^2_F$
交叉熵法（cross entropy）： $\argmin_{q_{min},q_{max}}=H(softmax(V),softmax(\hat{V}(q_{min},q_{max})))$ ，其中 $H$ 表示 cross entropy function
批量归一化法（BN based）： $q_{min}=min(\beta-\alpha\gamma),\ \ q_{max}=max(\beta+\alpha\gamma)$ ，其中 $\beta,\ \gamma$ 分布表示 batch norm 学到的 per channel 的 shift 和 scale， $\alpha>0$ 是超参数
组合法（comparsion）：以上方法的自由组合

使用不同量化方法分别量化 weight 和 activation 后的精度

2.2 跨层均衡（Cross-Layer Equalization）

这是一种 通过修改模型权重 来改善神经网络量化性能的技术，CLE 的目的是减少网络中不同 channel 之间由于量化引起的性能不平衡，这种问题在 depth-wise conv layer 中尤其容易出现。

mobilenetv2 第一个 depth-wise conv 层的 per output channel weight range

想要实现跨层均衡的模型，需要激活函数满足交换律，即： $f(sx)=sf(x)$ ，常见的 ReLU 和 PReLU 都满足。
CLE 原理：
- $y=f(W_2(W_1x+b_1)+b_2)\\=f(SW_2(S^{-1}W_1x+S^{-1}b_1)+b_2)\\=f(\hat{W_2}(\hat{W_1}x+\hat{b_1})+b_2)$
- 其中：
  - $\hat{W_1}=S^{-1}W_1$
  - $\hat{b_1}=S^{-1}b_1$
  - $\hat{W_2}=SW_2$
  - $S_i=\frac{\sqrt {r_i^1r_i^2}}{r_i^2}$ ，其中 $r_i^j$ 表示 j tensor 的 i channel
abosrbing high bias 是一种 解决模型中过大 bias 的技术，原理是：
- $y=W_2h+b_2\\=W_2(f(W_1x+b_1))+b_2\\=W_2(f(W_1x+b_1)+c-c)+b_2\\=W_2(f(W_1x+\hat{b_1})+c)+b_2\\=W_2(f(W_1x+\hat{b_1}))+\hat{b_2}\\=W_2\hat{h}+\hat{b_2}$
- 其中：
  - $\hat{b_2}=b_2+W_2c$
  - $\hat{h}=h-c$
  - $\hat{b_1}=b_1-c$
  - $c_i=max(0, min_x(W_{1i}x+b_{1i}))$

URL

背景知识

GPU 体系结构

Element-wise Add GPU 加速

窗口求和 GPU 加速

矩阵乘法 GPU 加速

程序自动变换

URL

使用 Builder 创建 IRModule

从张量表达式创建 TensorIR（主张量函数）

使用 BlockBuilder 构造 IRModule

从 Pytorch 映射到 IRModule

Pytorch 模型

构造计算图之间的映射变换

映射 Pytorch Module 到 TensorIR

或映射到 Pytorch Module 到 IRModule 更高层的算子

总结

URL

自动程序优化的原因

自动程序优化过程

end-to-end 构建模型

自动优化 linear0 主张量函数

URL

机器学习编译的本质与关注点

构造 IR_Module

构建并运行模型

使用现有库避免重复造轮子

绑定参数到 IR_Module

总结

三种成像相关坐标系

相机内参

相机外参

总结

URL

TL;DR

Algorithm

Dataset

相关知识

算法细节

1. 特征提取（Lift）

2. 像素坐标和相机坐标系下深度到世界坐标的映射（Splat）

3. 体素编码降维（Shoot）

4. 训练 loss

Thought

URL

TL;DR

Algorithm

总体结构

Patch Partition

Linear Embedding

Patch Merging

W-MSA

SW-MSA

W-MSA + SW-MSA

Relative Position Bias

Through

Python 闭包 (closure)

闭包定义

闭包的示例

闭包的用处

1. 可以读取函数内部的变量

2. 让这些变量的值始终保持在内存中

3. 用于装饰器

装饰器 (Decorators)

装饰器的定义

函数装饰器的示例

类别装饰器示例

URL

TL;DR

Algorithm

1. 量化基础知识

1.1 硬件背景

1.2 均匀仿射量化

1.3 量化模拟

1.4 实践考量

2. 训练后量化（PTQ，post-training quantization）

2.1 量化范围的设置

2.2 跨层均衡（Cross-Layer Equalization）

WIP

使用 `Builder` 创建 `IRModule`

从张量表达式创建 `TensorIR`（主张量函数）

使用 `BlockBuilder` 构造 `IRModule`

从 `Pytorch` 映射到 `IRModule`

`Pytorch` 模型

映射 `Pytorch Module` 到 `TensorIR`

或映射到 `Pytorch Module` 到 `IRModule` 更高层的算子

构造 `IR_Module`