参数化：编译时元编程
Python 最令人惊叹的特性之一就是其可扩展的运行时元编程特性。这使得许多库得以开发，并提供了灵活可扩展的编程模型，让全球范围内的 Python 程序员受益。然而，这些特性也是有代价的：因为它们在运行时进行评估，直接影响底层代码的运行效率。由于 IDE 不了解这些特性，因此很难让代码补全等 IDE 功能理解它们，并利用它们改善开发人员的体验。
在 Python 生态系统之外，静态元编程也是开发的重要组成部分，它使得开发新的编程范式和高级库成为可能。在这个领域中有许多先前的艺术例子，具有不同的权衡，例如：
- 预处理器（如 C 预处理器、Lex/YACC 等）可能是最“笨重”的。它们是完全通用的，但在开发人员体验和工具集成方面是最差的。
- 一些语言（如 Lisp 和 Rust）支持（有时是“卫生”）的宏展开特性，可以实现语法扩展和减少样板代码，并具有较好的工具集成。
- 一些较旧的语言，如 C++，具有非常庞大和复杂的元编程语言（模板），它们是运行时语言的对应物。这些语言特性难以学习，并且编译时间和错误信息都很差。
- 一些语言（如 Swift）在核心语言中直接构建了许多特性，以提供良好的常见情况使用体验，但牺牲了通用性。
- 一些较新的语言，如 Zig，在编译过程中集成了一个语言解释器，并允许解释器在编译时对 AST 进行反射。这使得可以提供与宏系统类似的功能，同时具有更好的可扩展性和通用性。
对于 Modula.AI 在人工智能、高性能机器学习内核和加速器方面的工作，我们需要先进的元编程系统提供高抽象性能力。我们需要高层次的零成本抽象、表达力强的库以及多种算法变体的大规模集成。我们希望库开发人员能够像在 Python 中一样扩展系统，提供可扩展的开发平台。
然而，我们不愿牺牲开发人员体验（包括编译时间和错误信息），也对于构建一个难以教授的并行语言生态系统不感兴趣。我们可以借鉴这些先前的系统，但也有新的技术可以建立在其之上，包括 MLIR 和细粒度的语言集成缓存技术。
因此，Mojo 支持将编译时元编程作为编译器的一个独立编译阶段——在解析、语义分析和 IR 生成之后，但在降低到特定目标代码之前。它使用相同的宿主语言来作为运行时程序和元程序，并利用 MLIR 来可预测地表示和评估这些程序。
让我们来看几个简单的例子。
关于“参数”：Python 开发者通常可以交替使用“参数”和“变量”这两个词来表示“传递给函数的东西”。我们决定将“参数”和“参数表达式”重新定义为在 Mojo 中表示编译时值的概念，并继续使用“参数”和“表达式”来表示运行时值。这使我们能够围绕“参数化”和“参数化”的概念进行对齐。
定义参数化的类型和函数
您可以通过在方括号中指定参数名称和类型来对结构体和函数进行参数化（使用了PEP695语法的扩展版本）。与参数值不同，参数值在编译时已知，这使得额外的抽象层和代码重用成为可能，并且还可以实现编译器的优化，如自动调优。
例如，让我们来看一个SIMD类型，它代表硬件中的一个低级向量寄存器，可以容纳多个实例的标量数据类型。硬件加速器不断引入新的矢量数据类型，甚至CPU可能具有512位或更长的SIMD向量。为了在这些处理器上访问SIMD指令，数据必须被塑造成适当的SIMD宽度（数据类型）和长度（向量大小）。
然而，使用Mojo的内置类型定义所有不同的SIMD变量是不可行的。因此，Mojo的SIMD类型（定义为一个结构体）通过其方法公开了常见的SIMD操作，并使SIMD数据类型和大小的值成为参数化。这样就可以直接将数据映射到任何硬件上的SIMD向量。
下面是Mojo的SIMD类型定义的简化（非功能性）版本:
struct SIMD[type: DType, size: Int]:
var value: … # 这里有一些低级别的MLIR内容
fn __init__(inout self, *elems: SIMD[type, 1]): ...
@staticmethod
fn splat(x: SIMD[type, 1]) -> SIMD[type, size]: ...
fn cast[target: DType](self) -> SIMD[target, size]: ...
fn __add__(self, rhs: Self) -> Self: ...
对每个SIMD变体使用参数进行定义对于代码重用非常有用，因为SIMD类型可以静态地表达所有不同的矢量变体，而不需要语言预定义每个变体。
因为SIMD是一个参数化类型，在其函数中的self参数携带这些参数——完整的类型名称是SIMD[type, size]。尽管可以写出完整的类型名称（如splat()中的返回类型所示），但这样可能会冗长，所以我们建议使用Self类型（来自PEP673），就像__add__示例所做的那样。
参数重载
函数和方法可以根据其参数签名进行重载。重载解析逻辑根据以下规则进行候选过滤，按优先顺序：
- 具有最少隐式转换（在参数和参数之间）的候选者。
- 不带可变参数的候选者。
- 不带可变参数的参数的候选者。
- 具有最短参数签名的候选者。
如果应用这些规则后存在多个候选者，则重载解析失败。例如：
@register_passable("trivial")
struct MyInt:
"""可以隐式转换为`Int`的类型。"""
var value: Int
@always_inline("nodebug")
fn __init__(_a: Int) -> Self:
return Self {value: _a}
fn foo[x: MyInt, a: Int]():
print("foo[x: MyInt, a: Int]()")
fn foo[x: MyInt, y: MyInt]():
print("foo[x: MyInt, y: MyInt]()")
fn bar[a: Int](b: Int):
print("bar[a: Int](b: Int)")
fn bar[a: Int](*b: Int):
print("bar[a: Int](*b: Int)")
fn bar[*a: Int](b: Int):
print("bar[*a: Int](b: Int)")
fn parameter_overloads[a: Int, b: Int, x: MyInt]():
foo[x, a]()
bar[a](b)
bar[a, a, a](b)
parameter_overloads[1, 2, MyInt(3)]()
foo[x: MyInt, a: Int]()
bar[a: Int](b: Int)
bar[*a: Int](b: Int)
使用参数化类型和函数
您可以通过在方括号中传递值给参数来实例化参数化类型和函数。例如，对于上面的SIMD类型，type指定数据类型，size指定SIMD向量的长度（它必须是2的幂）：
let small_vec = SIMD[DType.float32, 4](1.0, 2.0, 3.0, 4.0)
let big_vec = SIMD[DType.float16, 32].splat(1.0)
let bigger_vec = (big_vec+big_vec).cast[DType.float32]()
let bigger_vec2 : SIMD[DType.float32, 32] = bigger_vec
print('small_vec type:', small_vec.element_type, 'length:', len(small_vec))
print('bigger_vec2 type:', bigger_vec2.element_type, 'length:', len(bigger_vec2))
small_vec type: float32 length: 4
bigger_vec2 type: float32 length: 32
请注意，cast()方法也需要一个参数来指定从转换中所需的类型（上面的方法定义期望一个target参数化值）。因此，就像SIMD结构体是一个泛型类型定义一样，cast()方法也是一个泛型方法定义，它在编译时而不是运行时实例化，基于参数值。
上面的代码展示了具体类型（即，使用已知类型值实例化SIMD），但是参数的主要优势来自于能够定义参数化算法和类型（使用参数值的代码）。例如，以下是如何定义使用SIMD的弱类型和宽度的参数化算法：
from math import sqrt
fn rsqrt[dt: DType, width: Int](x: SIMD[dt, width]) -> SIMD[dt, width]:
return 1 / sqrt(x)
print(rsqrt[DType.float16, 4](42))
[0.154296875, 0.154296875, 0.154296875, 0.154296875]
请注意，x参数实际上是基于函数参数的SIMD类型。运行时程序可以使用参数的值，因为参数在运行时程序需要之前的编译时解析（但是编译时参数表达式不能使用运行时值）。
Mojo编译器对于参数的类型推断也很聪明。请注意，上面的函数能够调用参数化的sqrt()函数而不需要指定参数——编译器将其参数推断为如果你显式写了sqrt[type, simd_width](x)。还要注意，rsqrt()选择定义其第一个参数名为width，即使SIMD类型将其命名为size，也没有问题。
参数表达式只是 Mojo 代码
参数表达式是在期望参数的地方出现的任意代码表达式（如 a+b）。参数表达式支持运算符和函数调用，就像运行时代码一样，并且所有的参数类型都使用与运行时程序相同的类型系统（如 Int 和 DType）。
由于参数表达式使用与运行时 Mojo 代码相同的语法和类型，您可以使用许多“依赖类型”功能。例如，您可以定义一个辅助函数来连接两个 SIMD 向量：
fn concat[ty: DType, len1: Int, len2: Int](
lhs: SIMD[ty, len1], rhs: SIMD[ty, len2]) -> SIMD[ty, len1+len2]:
var result = SIMD[ty, len1 + len2]()
for i in range(len1):
result[i] = SIMD[ty, 1](lhs[i])
for j in range(len2):
result[len1 + j] = SIMD[ty, 1](rhs[j])
return result
let a = SIMD[DType.float32, 2](1, 2)
let x = concat[DType.float32, 2, 2](a, a)
print('result type:', x.element_type, 'length:', len(x))
结果类型: float32 长度: 4
注意结果长度是输入向量长度的和，您可以用简单的 + 运算来表达。对于更复杂的例子，请看一下标准库中的 SIMD.shuffle()方法：它接受两个输入 SIMD 值，一个向量打乱掩码作为列表，并返回一个与打乱掩码长度相匹配的 SIMD。
强大的编译时编程
虽然简单的表达式很有用，但有时候您希望编写带有控制流的命令式编译时逻辑。例如，在 Mojo 的 Math 模块中，isclose() 函数对于整数使用精确相等性，但对于浮点数使用“接近”比较。您甚至可以在编译时进行递归。例如，这是一个示例的“树归约”算法，将一个向量中的所有元素递归求和为一个标量：
fn slice[ty: DType, new_size: Int, size: Int](
x: SIMD[ty, size], offset: Int) -> SIMD[ty, new_size]:
var result = SIMD[ty, new_size]()
for i in range(new_size):
result[i] = SIMD[ty, 1](x[i + offset])
return result
fn reduce_add[ty: DType, size: Int](x: SIMD[ty, size]) -> Int:
@parameter
if size == 1:
return x[0].to_int()
elif size == 2:
return x[0].to_int() + x[1].to_int()
alias half_size = size // 2
let lhs = slice[ty, half_size, size](x, 0)
let rhs = slice[ty, half_size, size](x, half_size)
return reduce_add[ty, half_size](lhs + rhs)
let x = SIMD[DType.index, 4](1, 2, 3, 4)
print(x)
print("元素求和:", reduce_add[DType.index, 4](x))
[1, 2, 3, 4]
元素求和: 10
这利用了 @parameter if 特性，它是一个在编译时运行的 if 语句。它要求它的条件是一个有效的参数表达式，并确保只有 if 语句的活动分支被编译到程序中。
Mojo类型只是参数表达式
虽然我们已经展示了如何在类型中使用参数表达式，类型注解本身也可以是任意表达式（就像在Python中一样）。Mojo中的类型具有特殊的元类型，允许定义类型参数化的算法和函数。
例如，我们可以创建一个简化的Array，该数组支持任意类型的元素（通过 AnyType 参数）：
struct Array[T: AnyType]:
var data: Pointer[T]
var size: Int
var cap: Int
fn __init__(inout self, size: Int, value: T):
self.cap = size * 2
self.size = size
self.data = Pointer[T].alloc(self.cap)
for i in range(self.size):
self.data.store(i, value)
fn __getitem__(self, i: Int) -> T:
return self.data.load(i)
fn __del__(owned self):
self.data.free()
var v = Array[Float32](4, 3.14)
print(v[0], v[1], v[2], v[3])
3.1400001049041748 3.1400001049041748 3.1400001049041748 3.1400001049041748
注意，T 参数被用作 value 参数的实际类型和 __getitem__ 函数的返回类型。参数允许 Array 类型根据不同的用例提供不同的API。
还有许多其他情况可以从参数的更高级使用中受益。例如，您可以并行地执行闭包N次，从上下文中提供一个值，如下所示：
fn parallelize[func: fn (Int) -> None](num_work_items: Int):
for i in range(num_work_items):
func(i)
另一个重要的例子是可变泛型，其中算法或数据结构可能需要定义在异构类型列表上，比如元组：
struct Tuple[*Ts: AnyType]:
var _storage : *Ts
虽然我们还没有足够的元类型辅助功能，但我们应该能够在将来编写类似于下面的代码（尽管重载仍然是处理此问题的更好方法）：
struct Array[T: AnyType]:
fn __getitem__[IndexType: AnyType](self, idx: IndexType)
-> (ArraySlice[T] if issubclass(IndexType, Range) else T):
...
alias：命名的参数表达式
经常希望对编译时值进行命名。而 var 定义了运行时值，let 定义了运行时常量，我们需要一种定义编译时临时值的方式。为此，Mojo使用 alias 声明。
例如，DType 结构使用别名来实现简单的枚举，如下所示（实际的 DType 实现细节有所不同）：
struct DType:
var value : UI8
alias invalid = DType(0)
alias bool = DType(1)
alias int8 = DType(2)
alias uint8 = DType(3)
alias int16 = DType(4)
alias int16 = DType(5)
...
alias float32 = DType(15)
这允许客户端将 DType.float32 作为参数表达式（也可以作为运行时值）自然地使用。请注意，这在编译时调用 DType 的运行时构造函数。
别名对于类型是常见的用法。因为类型是编译时表达式，因此像这样做非常方便：
alias Float16 = SIMD[DType.float16, 1]
alias UInt8 = SIMD[DType.uint8, 1]
var x : Float16 # F16就像一个“typedef”
与var和let一样，别名遵循作用域规则，您可以在函数内部使用本地别名，就像您预期的那样。
顺便说一句，None和AnyType都被定义为类型别名。
自动调优/自适应编译
Mojo参数表达式允许您写出像其他语言一样可移植的参数化算法，但是在编写高性能代码时，仍然需要选择具体的参数值来使用。例如，当编写高性能的数值算法时，您可能希望使用内存划分来加速算法，但使用的维度高度依赖于可用的硬件特性、缓存区的大小、融入内核的内容以及许多其他琐碎的细节。
甚至矢量长度也可能难以管理，因为典型机器的矢量长度取决于数据类型，并且某些数据类型（例如bfloat16）在所有实现中并不具备完全的支持。Mojo通过在标准库中提供“autotune”函数来提供帮助。例如，如果您想要向数据缓冲区写入一个与矢量长度无关的算法，则可以像这样编写：
from autotune import autotune, search
from benchmark import Benchmark
from memory.unsafe import DTypePointer
from algorithm import vectorize
fn buffer_elementwise_add_impl[
dt: DType
](lhs: DTypePointer[dt], rhs: DTypePointer[dt], result: DTypePointer[dt], N: Int):
"""对RHS和LHS中的N个元素执行逐元素加法，并将结果存储在RESULT中。"""
@parameter
fn add_simd[size: Int](idx: Int):
let lhs_simd = lhs.simd_load[size](idx)
let rhs_simd = rhs.simd_load[size](idx)
result.simd_store[size](idx, lhs_simd + rhs_simd)
alias vector_len = autotune(1, 4, 8, 16, 32)
vectorize[vector_len, add_simd](N)
fn elementwise_evaluator[dt: DType](
fns: Pointer[fn (DTypePointer[dt], DTypePointer[dt], DTypePointer[dt], Int) -> None],
num: Int,
) -> Int:
alias N = 64
let lhs = DTypePointer[dt].alloc(N)
let rhs = DTypePointer[dt].alloc(N)
let result = DTypePointer[dt].alloc(N)
for i in range(N):
lhs.store(i, 1)
rhs.store(i, 1)
var best_idx: Int = -1
var best_time: Int = -1
for i in range(num):
@parameter
fn wrapper():
fns.load(i)(lhs, rhs, result, N)
let cur_time = Benchmark(1).run[wrapper]()
if best_idx < 0 or cur_time < best_time:
best_idx = i
best_time = cur_time
print("time[", i, "] =", cur_time)
print("selected:", best_idx)
return best_idx
fn buffer_elementwise_add[
dt: DType
](lhs: DTypePointer[dt], rhs: DTypePointer[dt], result: DTypePointer[dt], N: Int):
alias best_impl: fn(DTypePointer[dt], DTypePointer[dt], DTypePointer[dt], Int) -> None
search[
fn(DTypePointer[dt], DTypePointer[dt], DTypePointer[dt], Int) -> None,
buffer_elementwise_add_impl[dt],
elementwise_evaluator[dt] -> best_impl
]()
best_impl(lhs, rhs, result, N)
现在，我们可以像通常一样调用我们的函数：
let N = 32
let a = DTypePointer[DType.float32].alloc(N)
let b = DTypePointer[DType.float32].alloc(N)
let res = DTypePointer[DType.float32].alloc(N)
for i in range(N):
a.store(i, 2.0)
b.store(i, 40.0)
res.store(i, -1)
buffer_elementwise_add[DType.float32](a, b, res, N)
print(a.load(10), b.load(10), res.load(10))
编译此代码的时候，Mojo会对此算法进行编译的分叉，并通过测量在目标硬件上实际效果最好的值来决定使用哪个值。它会评估vector_len表达式的不同取值，并根据用户定义的性能评估器选择最快的那个。因为它会逐个测量和评估每个选项，所以在float32和int8中可能会选择不同的向量长度。这个简单的特性非常强大，可以超越简单的整数常量，因为函数和类型也可以是参数表达式。
请注意，寻找最佳向量长度是由search()函数执行的。search()函数接受一个评估器和一个分叉函数，并将由评估器选择的最快实现作为参数结果返回。
自动调优是一种固有的指数级技术，它受益于Mojo编译器栈的内部实现细节（尤其是MLIR，集成缓存和编译分发）。这也是一个高级用户功能，需要不断的开发和迭代。