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]