Mojo 学习 —— 并行化

Mojo 学习 —— 并行化

Mojo 标准库 algorithm 模块中提供了很多并行化操作函数。包括

• functional 模块: 实现了一些高级函数
• memory 模块: 实现了并行内存拷贝

这些函数可以直接从 algorithm 中导入，例如

from algorithm import map

我们首先定义一个矩阵结构，用于后续并行化分析

from random import rand

alias type = DType.float32

struct Matrix[rows: Int, cols: Int]:
    # 使用 DTypePointer 类型的指针
    var data: DTypePointer[type]

    # 将数据初始化为 0
    fn __init__(inout self):
        self.data = DTypePointer[type].alloc(rows * cols)
        memset_zero(self.data, rows * cols)

    # 或使用指针数据初始化
    fn __init__(inout self, data: DTypePointer[type]):
        self.data = data

    # 将矩阵的值初始化为随机值
    @staticmethod
    fn rand() -> Self:
        var data = DTypePointer[type].alloc(rows * cols)
        rand(data, rows * cols)
        return Self(data)
    # 使用 [] 获取指定位置的值
    fn __getitem__(self, y: Int, x: Int) -> Scalar[type]:
        return self.load[1](y, x)
    # [] 设置值
    fn __setitem__(self, y: Int, x: Int, val: Scalar[type]):
        self.store[1](y, x, val)
    # 返回 x,y 后面的 nelts 个值
    fn load[nelts: Int](self, y: Int, x: Int) -> SIMD[type, nelts]:
        return self.data.load[width=nelts](y * self.cols + x)
    # 设置  x,y 后面的 nelts 个值
    fn store[nelts: Int](self, y: Int, x: Int, val: SIMD[type, nelts]):
        return self.data.store[width=nelts](y * self.cols + x, val)

然后定义一个基准测试函数（矩阵乘法），用于测试不同版本的函数执行效率的区别。

import benchmark

# 定义常量，矩阵的大小
alias M = 1024
alias N = 1024
alias K = 1024

@always_inline
fn bench[
    func: fn (Matrix, Matrix, Matrix) -> None](base_gflops: Float64):
    # 创建三个矩阵
    var C = Matrix[M, N]()
    var A = Matrix[M, K].rand()
    var B = Matrix[K, N].rand()

    # 调用矩阵乘法函数
    @always_inline
    @parameter
    fn test_fn():
        _ = func(C, A, B)
    # 基准测试
    var secs = benchmark.run[test_fn](max_runtime_secs=1).mean()
    # 释放内存
    A.data.free()
    B.data.free()
    C.data.free()
    # 打印测试结果
    var gflops = ((2 * M * N * K) / secs) / 1e9
    var speedup: Float64 = gflops / base_gflops

    print(gflops, "GFLOP/s, a", speedup, "x speedup over Python")

使用嵌套循环计算矩阵乘法

fn matmul_naive(C: Matrix, A: Matrix, B: Matrix):
    for m in range(C.rows):
        for k in range(A.cols):
            for n in range(C.cols):
                C[m, n] += A[m, k] * B[k, n]

以 Python 的计算结果为基准，测试 Mojo 运行结果

fn main():
    var python_gflops = 0.002
    bench[matmul_naive](python_gflops)
# 2.5135788794310976 GFLOP/s, a 1256.7894397155487 x speedup over Python

functional

functional 目前有 11 个高级函数，以及一列类型别名。我们主要介绍这些函数的使用

map

函数签名为：

map[func: fn(Int, /) capturing -> None](size: Int)

其功能为将 0-size 分别作为输入函数的参数，并运行，例如

fn map_demo():
    @parameter
    fn helper(num: Int):
        print(num, end=' ')

    map[helper](5)

map_demo()

注意，要使用闭包的形式，闭包的第一个参数为每次迭代获取的 0-size 的值，且闭包无返回值。输出结果为

0 1 2 3 4 

vectorize

函数签名为：

vectorize[func: fn[Int](Int, /) capturing -> None, simd_width: Int, /, *, unroll_factor: Int = 1](size: Int)

通过在从 0-size 的范围内映射函数来简化 SIMD 循环的优化，以 simd_width 为步长递增。剩余部分（size % simd_width）将在单独的迭代中运行。例如

from algorithm.functional import vectorize

# 要循环遍历的元素数量
alias size = 10
# 一个 SIMD 寄存器可以存储几个 Dtype.int32 元素（128 可以保存 8 个）
alias simd_width = simdwidthof[DType.int32]()

fn main():
    var p = DTypePointer[DType.int32].alloc(size)

    # 定义闭包来捕获指针
    @parameter
    fn closure[simd_width: Int](i: Int):
        print("storing", simd_width, "els at pos", i)
        p.store[width=simd_width](i, i)

    vectorize[closure, simd_width](size)
    print(p.load[width=size]())

在 SIMD 寄存器大小为 256 的机器上，每次迭代将设置 8 个 Int32 的值。10 % 8 的余数是 2，因此最后两个元素将在两次不同的迭代中设置：

storing 8 els at pos 0
storing 1 els at pos 8
storing 1 els at pos 9
[0, 0, 0, 0, 0, 0, 0, 0, 8, 9]

您还可以将循环展开，通过牺牲存储来提高性能：

vectorize[closure, width, unroll_factor=2](size)

类似于下面的伪代码，最后两次循环会重复运行

closure[8](0)
# 除非将 `size` 作为 `parameter` 传递，否则剩余循环不会展开
for i in range(8, 10):
    closure[1](i)
    closure[1](i)

如果我们在编译时可以确定 size 的值，可以将其作为 parameter，而不是 argument 从而减少余数的迭代次数。

只有当剩余循环为 2 的指数（2、4、8、16…）时，才会生效。如果余数不是 2 的指数，余下的循环仍会展开以提高性能。

重载函数的签名为：

vectorize[func: fn[Int](Int, /) capturing -> None, simd_width: Int, /, *, size: Int, unroll_factor: Int = 1]()

我们只需将上面倒数第二行代码改为

vectorize[closure, simd_width, size=size]()

其运行结果将会是

storing 8 els at pos 0
storing 2 els at pos 8
[0, 0, 0, 0, 0, 0, 0, 0, 8, 8]

我们使用 vectorize 来加速矩阵乘法的计算

# 每次向量化计算的长度
alias nelts = simdwidthof[DType.float32]() * 2

fn matmul_vectorized(C: Matrix, A: Matrix, B: Matrix):

    for m in range(C.rows):
        for k in range(A.cols):
            @parameter
            fn dot[nelts: Int](n: Int):
                C.store(m, n, C.load[nelts](m, n) + A[m, k] * B.load[nelts](k, n))
            vectorize[dot, nelts, size = C.cols]()
# 14.092552861535482 GFLOP/s, a 7046.2764307677407 x speedup over Python

提升差不多 6 倍

parallelize

parallelize 和 sync_parallelize 都可以执行并行任务，区别可能一个是同步一个是异步。使用方式差不多，parallelize 还有一个额外的参数

parallelize[func: fn(Int, /) capturing -> None](num_work_items: Int, num_workers: Int)

例如，可以使用并行化来赋值

from algorithm import sync_parallelize, parallelize, map

fn main():
    alias size = 10
    var p = Pointer[Int].alloc(size)
    @parameter
    fn fill_value(i: Int):
        p.store(i, i)
    print('store value to pointer')
    # sync_parallelize[fill_value](size)
    parallelize[fill_value](size, 3)

    @parameter
    fn print_value(i: Int):
        print(p.load(i), end=' ')
    print('get value from pointer')
    map[print_value](size)
    p.free()

store value to pointer
get value from pointer
0 1 2 3 4 5 6 7 8 9

我们可以把 parallelize 加入，进一步提升矩阵乘法的计算速度

fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix):
    @parameter
    fn calc_row(m: Int):
        for k in range(A.cols):
            @parameter
            fn dot[nelts : Int](n : Int):
                C.store[nelts](m,n, C.load[nelts](m,n) + A[m,k] * B.load[nelts](k,n))
            vectorize[dot, nelts, size = C.cols]()
    parallelize[calc_row](C.rows, C.rows)
# 116.56057998764479 GFLOP/s, a 58280.289993822393 x speedup over Python

woooooo~，这次比使用 vectorize 的方法提升了 8 倍多。速度杠杠的！

tile

tile 是一个生成器，可按照指定的任务大小列表，即将任务进行切片分组，并分组执行函数。

例如 func[3](5) 应启动对第 5、6、7 项的计算，在语义上应等同于 func[1](5)、func[1](6)、func[1](7)。

例如，tile[func, (3,2,1)](offset, upperbound) 会尝试从 offset 开始调用 func[3]，直到离 upperbound 还有三个任务，然后尝试调用 func[2]，再调用 func[1]。

例如，将 9 个任务分为 3 组，每组 3 个任务

fn tile_test():
    @parameter
    fn func[tile_x: Int](x: Int):
        print(tile_x, x)
    tile[func, tile_size_list=3](0, 9)
# 3 0
# 3 3
# 3 6

类似于

fn tile_test():
    @parameter
    fn func(x: Int, y: Int):
        print(x, y)
    tile[func](0, 9, 3)
# 0 3
# 3 3
# 6 3

我们将 tile 也加入矩阵乘法的计算中

from algorithm import tile

fn matmul_tiled_parallelized(C: Matrix, A: Matrix, B: Matrix):
    @parameter
    fn calc_row(m: Int):
        @parameter
        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):
            for k in range(y, y + tile_y):
                @parameter
                fn dot[nelts: Int](n: Int):
                    C.store(m, n + x, C.load[nelts](m, n + x) + A[m, k] * B.load[nelts](k, n + x))
                vectorize[dot, nelts, size = tile_x]()

        # We hardcode the tile factor to be 4.
        alias tile_size = 4
        tile[calc_tile, nelts * tile_size, tile_size](0, 0, A.cols, C.cols)

    parallelize[calc_row](C.rows, C.rows)
# 176.02634276604189 GFLOP/s, a 88013.171383020948 x speedup over Python

elementwise

函数签名为

elementwise[func: fn[Int, Int](StaticIntTuple[$1], /) capturing -> None, simd_width: Int, rank: Int](shape: StaticIntTuple[rank])

将函数闭包 func[width, rank](indices) 作为子任务运行，子任务的数量会根据指定的宽度和形状进行拆分并覆盖所有数据，当所有子任务完成时返回。

例如，我们可以为一个矩阵类型的数据指针，按行赋值

from algorithm.functional import elementwise

alias nrow = 3
alias ncol = 4

fn main():
    var p = DTypePointer[DType.int32].alloc(nrow * ncol)

    @parameter
    fn wise[i: Int, j: Int](x: StaticIntTuple[j]):
        var pos = x[1] + x[0] * i
        print("storing", i, "elements at position", pos)
        p.store[width=i](pos, x[0])

    elementwise[wise, ncol, 2]((nrow, ncol))
    print(p.load[width=nrow * ncol]())

输出结果

storing 4 elements at position 0
storing 4 elements at position 4
storing 4 elements at position 8
[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]

unswitch

unswitch 是一种简单的模式，类似于循环内的判断条件外提。该模式有助于进行代码转换，从而减少生成代码中的分支数量。

简单来说就是，将下面这种循环模式

for i in range(...)
    if i < xxx:
        ...

转变为

if i < ...
    for i in range(...)
        ...
else
    for i in range(...)
        if i < xxx:
            ...

减少循环中的判断次数

from algorithm import unswitch

fn tile_test():
    @parameter
    fn switch[sign: Bool]():
        if sign:
            for i in range(10):
                print(i, end=' ')
        else:
            for i in range(10):
                print(i * -1, end=' ')

    unswitch[switch](False)
    unswitch[switch](True)
# 0 -1 -2 -3 -4 -5 -6 -7 -8 -9 0 1 2 3 4 5 6 7 8 9

memory

该模块中目前只有一个函数 parallel_memcpy，其重要功能是并行将内存缓冲区 src 中的 count 个元素复制到 dest

先导入进来

from algorithm import parallel_memcpy

例如，使用简单的并行拷贝

from algorithm import parallel_memcpy

fn main():
    var src = DTypePointer[DType.int8].alloc(8)
    src.store(0, SIMD[DType.int8, 8](1, 2, 3, 4, 5, 6, 7, 8))

    var dest = DTypePointer[DType.int8].alloc(4)
    parallel_memcpy(dest, src, 4)
    
    print(dest.load[width=4]())
    
    src.free()
    dest.free()
# [1, 2, 3, 4]

还可以指针任务数量以及每个任务复制元素的个数。例如

from algorithm import parallel_memcpy

fn main():
    var src = DTypePointer[DType.int8].alloc(8)
    src.store(0, SIMD[DType.int8, 8](1, 2, 3, 4, 5, 6, 7, 8))

    var dest = DTypePointer[DType.int8].alloc(8)
    parallel_memcpy(dest, src, 8, count_per_task=4, num_tasks=2)

    print(dest.load[width=8]())

    src.free()
    dest.free()
# [1, 2, 3, 4, 5, 6, 7, 8]