mojo实现高阶函数(algorithm)
functional
实现高阶函数。
您可以从 algorithm 包导入这些 API。例如：
from algorithm import map
别名:
-
Static1DTileUnitFunc = fn[Int](Int, /) capturing -> None: Signature of a 1d tiled function that performs some work with a static tile size and an offset. i.e. func<tile_size: Int> (offset: Int) -
Dynamic1DTileUnitFunc = fn(Int, Int, /) capturing -> None: Signature of a 1d tiled function that performs some work with a dynamic tile size and an offset. i.e. func(offset: Int, tile_size: Int) -
BinaryTile1DTileUnitFunc = fn[Int](Int, Int, /) capturing -> None: Signature of a tiled function that performs some work with a dynamic tile size and a secondary static tile size. -
Static2DTileUnitFunc = fn[Int, Int](Int, Int, /) capturing -> None: Signature of a 2d tiled function that performs some work with a static tile size and an offset. i.e. func<tile_size_x: Int, tile_size_y: Int> (offset_x: Int, offset_y: Int) -
SwitchedFunction = fn[Bool]() capturing -> None -
SwitchedFunction2 = fn[Bool, Bool]() capturing -> None -
Static1DTileUnswitchUnitFunc = fn[Int, Bool](Int, Int, /) capturing -> None: Signature of a tiled function that performs some work with a static tile size and an offset. i.e. func<tile_size: Int> (offset: Int) -
Static1DTileUnitFuncWithFlag = fn[Int, Bool](Int, /) capturing -> None -
Dynamic1DTileUnswitchUnitFunc = fn[Bool](Int, Int, Int, /) capturing -> None -
Static1DTileUnitFuncWithFlags = fn[Int, Bool, Bool](Int, /) capturing -> None
map
map[func: fn(Int, /) capturing -> None](size: Int)
将函数映射到从0到size的范围上。
Parameters:
- func (
fn(Int, /) capturing -> None)：要映射的函数。
Args:
- size (
Int)：元素的数量。
vectorize
vectorize[func: fn[Int](Int, /) capturing -> None, simd_width: Int, unroll_factor: Int](size: Int)
通过在 0 到 的范围内映射函数并在每一步size递增 来简化 SIMD 优化循环simd_width。其余部分size % simd_width将在单独的迭代中运行。
下面的示例演示了如何通过使用计算机上的 SIMD 寄存器同时设置多个值来提高循环的性能：
from algorithm.functional import vectorize# The amount of elements to loop through
alias size = 10
# How many Dtype.int32 elements fit into the SIMD register (4 on 128bit)
alias simd_width = simdwidthof[DType.int32]()fn main():var p = DTypePointer[DType.int32].alloc(size)# @parameter allows the closure to capture the `p` pointer@parameterfn 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 寄存器大小为 128 的机器上，这将在每次迭代时设置 4xInt32 值。 10 % 4 的余数为 2，因此最后两个元素将在两次单独的迭代中设置：
storing 4 els at pos 0
storing 4 els at pos 4
storing 1 els at pos 8
storing 1 els at pos 9
[0, 0, 0, 0, 4, 4, 4, 4, 8, 9]
您还可以展开循环以潜在地提高性能，但代价是二进制大小：
vectorize[closure, width, unroll_factor=2](size)
在生成的程序集中，函数调用将被重复，从而导致算术、比较和条件跳转操作减少。程序集在伪代码中看起来像这样：
closure[4](0)
closure[4](4)
# Remainder loop won't unroll unless `size` is passed as a parameter
for i in range(8, 10):closure[1](i)closure[1](i)
size如果已知编译时间可以减少剩余部分的迭代，则可以将其作为参数传递。仅当余数是 2 的指数 (2, 4, 8, 16, …) 时才会发生这种情况。如果不是 2 的指数，余数循环仍会展开以提高性能。
Parameters:
- func (
fn[Int](Int, /) capturing -> None)：将在循环体中调用的函数。 - simd _ width (
Int)：SIMD向量宽度。 - unroll_factor (
Int)：主循环的展开因子（默认 1 ）****。
Args:
- size (
Int): 循环的上限。
vectorize[func: fn[Int](Int, /) capturing -> None, simd_width: Int, size: Int, unroll_factor: Int]()
通过在 0 到 的范围内映射函数并在每一步size递增 来简化 SIMD 优化循环simd_width。size % simd_width如果它是 2 的指数，则其余部分将在单次迭代中运行。
下面的示例演示了如何通过使用计算机上的 SIMD 寄存器同时设置多个值来提高循环的性能：
from algorithm.functional import vectorize# The amount of elements to