Mojo的特征与参数(参数化部分)详解

许多语言都具有元编程功能：即编写生成或修改代码的代码。Python 具有动态元编程功能：装饰器、元类等功能。这些功能使 Python 非常灵活且高效，但由于它们是动态的，因此会产生运行时开销。其他语言具有静态或编译时元编程功能，如 C 预处理器宏和 C++ 模板。这些功能可能受到限制且难以使用。

为了支持 Modular 在 AI 领域的工作，Mojo 旨在提供功能强大、易于使用且运行时成本为零的元编程。这种编译时元编程使用与运行时程序相同的语言，因此您无需学习新语言 - 只需学习一些新功能即可。

主要新功能是参数。您可以将参数视为编译时变量，该变量将成为运行时常量。这种“参数”用法可能与您在其他语言中习惯的用法不同，在其他语言中，“参数”和“参数”通常可以互换使用。在 Mojo 中，“参数”和“参数表达式”指的是编译时值，而“参数”和“表达式”指的是运行时值。

在 Mojo 中，您可以向结构或函数添加参数。您还可以定义命名参数表达式（别名），将其用作运行时常量。

参数化函数

---

要定义参数化函数，请在参数列表前面的方括号中添加参数。每个参数的格式与参数一样：参数名称，后跟冒号和类型（必填）。在以下示例中，该函数有一个参数count类型为Int。

fn repeat[count: Int](msg: String):
    @parameter
    for i in range(count):
        print(msg)
1
2
3
4

@parameter此处显示的指令导致在for编译时评估循环。仅当循环限制是编译时常量时，该指令才有效。由于count是参数， 因此range(count)可以在编译时计算。

调用参数化函数时，您可以为参数提供值，就像函数参数一样：

repeat[3]("Hello")
1

输出：

Hello
Hello
Hello
1
2
3

编译器在编译期间解析参数值，并为每个唯一参数值创建函数的具体版本repeat。解析参数值并展开循环后，该repeat3 函数大致相当于：

fn repeat_3(msg: String):
    print(msg)
    print(msg)
    print(msg)
1
2
3
4

这并不代表编译器生成的实际代码。在解析参数时，Mojo 代码已经转换为MLIR中的中间表示。

如果编译器无法将所有参数值解析为常量值，则编译失败。

参数和泛型

---

“泛型”是指可以作用于多种类型值的函数，或可以容纳多种类型值的容器。例如， List可以容纳不同类型的值，因此您可以拥有一个值列表Int，或一个String值列表。

在 Mojo 中，泛型使用参数来指定类型。例如，List 采用类型参数，因此整数向量写为List[Int]。因此，所有泛型都使用参数，但并非所有使用参数的东西都是泛型。

例如，repeat上一节中的函数包含类型为Int的形参和类型为String的实参。它是参数化的，但不是泛型的。泛型函数或结构体在类型上是参数化的。例如，我们可以重写repeat以采用符合以下Stringable特征的任何类型的参数：

fn repeat[MsgType: Stringable, count: Int](msg: MsgType):
    @parameter
    for i in range(count):
        print(msg)

# Must use keyword parameter for `count`
repeat[count=2](42)
1
2
3
4
5
6
7

输出：

42
42
1
2

此更新的函数接受任何Stringable类型，因此您可以向其传递Int、 String或Bool值。
另外一个示例：

fn repeat[MsgType: Stringable, count: Int](msg: MsgType):
    @parameter
    for i in range(count):
        print(msg)
def main():
    # Must use keyword parameter for `count`
    repeat[count=2]("zhangsan")
1
2
3
4
5
6
7

输出：

zhangsan
zhangsan
1
2

如果不指定 ，则不能将count作为位置关键字传递MsgType。您可以将其放在//后面MsgType以指定它始终由参数推断。现在您可以按count位置传递以下参数：

fn repeat[MsgType: Stringable, //, count: Int](msg: MsgType):
    @parameter
    for i in range(count):
        print(msg)

# MsgType is always inferred, so first positional keyword `2` is passed to `count`
repeat[2](42)
1
2
3
4
5
6
7

Mojo 对泛型的支持还处于早期阶段。您可以使用特征和参数编写这样的泛型函数。您还可以编写像 List和这样的泛型集合。

参数化结构

---

您还可以向结构体添加参数。您可以使用参数化结构体来构建通用容器。例如，通用数组类型可能包含如下代码：

from memory.unsafe_pointer import UnsafePointer, initialize_pointee_copy, destroy_pointee

struct GenericArray[ElementType: CollectionElement]:
    var data: UnsafePointer[ElementType]
    var size: Int

    fn __init__(inout self, *elements: ElementType):
        self.size = len(elements)
        self.data = UnsafePointer[ElementType].alloc(self.size)
        for i in range(self.size):
            initialize_pointee_move(self.data.offset(i), elements[i])

    fn __del__(owned self):
        for i in range(self.size):
            destroy_pointee(self.data.offset(i))
        self.data.free()

    fn __getitem__(self, i: Int) raises -> ref [__lifetime_of(self)] ElementType:
        if (i < self.size):
            return self.data[i]
        else:
            raise Error("Out of bounds")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

该结构体有一个参数，ElementType它是您想要存储在数组中的数据类型的占位符，有时称为类型参数。ElementType其类型为 CollectionElement，它是一种 表示可以复制和移动的任何类型的特征。

与参数化函数一样，使用参数化结构体时需要传入参数值。在这种情况下，创建的实例时 GenericArray，需要指定要存储的类型，如Int、 或 Float64。（这有点令人困惑，因为在这种情况下传递的参数值是类型。没关系：Mojo 类型是有效的编译时值。）

您将看到它ElementType在整个结构体中被使用，而您通常会看到类型名称。例如，作为elements构造函数中的正式类型，以及方法的返回类型__getitem__()。

以下是使用的示例GenericArray：

var array = GenericArray[Int](1, 2, 3, 4)
for i in range(array.size):
    print(array[i], end=" ")
1
2
3

输出：

1 2 3 4
1

参数化结构体可以使用类型Self来表示结构体的具体实例（即指定了所有参数）。例如，您可以添加一个静态工厂方法，GenericArray其签名如下：

struct GenericArray[ElementType: CollectionElement]:
    ...

    @staticmethod
    fn splat(count: Int, value: ElementType) -> Self:
        # Create a new array with count instances of the given value
1
2
3
4
5
6

这里，Self相当于写成GenericArray[ElementType]。也就是说，你可以像splat()这样调用该方法：

GenericArray[Float64].splat(8, 0)
1

该方法返回 的一个实例GenericArray[Float64]。

案例研究：SIMD类型

---

为了得到参数化类型的真实示例，让我们看一下Mojo中的SIMD标准库中的类型。

单指令多数据 (SIMD)是许多现代 CPU、GPU 和自定义加速器内置的并行处理技术。SIMD 允许您同时对多个数据执行单个操作。例如，如果您想对数组中每个元素求平方根，则可以使用 SIMD 来并行化工作。

处理器使用硬件中的低级向量寄存器来实现 SIMD，这些寄存器可保存标量数据类型的多个实例。为了在这些处理器上使用 SIMD 指令，必须将数据调整为适当的 SIMD 宽度（数据类型）和长度（向量大小）。处理器可能支持 512 位或更长的 SIMD 向量，并支持从 8 位整数到 64 位浮点数的多种数据类型，因此定义所有可能的 SIMD 变体并不实际。

Mojo 的SIMD类型（定义为结构）通过其方法公开常见的 SIMD 操作，并使 SIMD 数据类型和大小值参数化。这允许您将数据直接映射到任何硬件上的 SIMD 向量。

以下是 Mojo 类型定义的精简版（非功能性版本）SIMD：

struct SIMD[type: DType, size: Int]:
    var value: … # Some low-level MLIR stuff here

    # Create a new SIMD from a number of scalars
    fn __init__(inout self, *elems: SIMD[type, 1]):  ...

    # Fill a SIMD with a duplicated scalar value.
    @staticmethod
    fn splat(x: SIMD[type, 1]) -> SIMD[type, size]: ...

    # Cast the elements of the SIMD to a different elt type.
    fn cast[target: DType](self) -> SIMD[target, size]: ...

    # Many standard operators are supported.
    fn __add__(self, rhs: Self) -> Self: ...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

因此，您可以像这样创建和使用 SIMD 向量：

var vector = SIMD[DType.int16, 4](1, 2, 3, 4