Mojo 中的低级别 IR

Mojo 中的低级别 IR

Mojo 是一种高级编程语言，具有丰富的现代化特性。但是，Mojo 也为您提供了访问所有低级原语的权限，这些原语是编写强大且零成本抽象所必需的。

这些原语在 MLIR 中实现，MLIR是一种可扩展的中间表示（IR）格式，用于编译器设计。许多不同的编程语言和编译器将其源代码程序转换为 MLIR，由于 Mojo 提供了对MLIR功能的直接访问权限，这意味着 Mojo 程序可以享受这些工具所带来的好处。

更进一步，Mojo 的独特之处在于其零成本抽象与 MLIR 互操作性的结合，这意味着 Mojo程序可以充分利用任何与 MLIR 接口的东西。虽然这可能不是 Mojo 程序员通常需要做的事情，但是当扩展系统以与新的数据类型或新颖的加速器功能接口时，这是一个极其强大的功能。

为了说明这些想法，我们将在下面使用 MLIR 实现一个名为 OurBool 的布尔类型。让我们开始一个简单的介绍。

什么是 MLIR

MLIR 是一种程序的中间表示形式，类似于汇编语言，其中指令序列操作内存中的值。

更重要的是，MLIR 是模块化和可扩展的。MLIR 由不断增长的“方言”组成。每个方言定义了操作和优化：例如，“数学”方言提供了正弦和余弦等数学运算，“amdgpu” 方言提供了特定于 AMD 处理器的操作等。

MLIR 的每个方言都可以与其他方言互操作。这就是为什么 MLIR 被称为解锁异构计算的原因：随着更先进、更快的处理器和架构的发展，新的MLIR方言被实现为这些环境生成最优代码。任何新的 MLIR 方言都可以无缝地转换为其他方言，因此随着更多的方言被添加，所有现有的 MLIR 都变得更加强大。

这意味着我们自己的自定义类型，如下面的 OurBool 类型，可以用来为程序员提供高级的 Python-like 接口。但是“在幕后”，Mojo 和 MLIR 将为我们方便的高级类型优化未来的新处理器。

关于为什么 MLIR 是这样一种革命性的技术还有很多要写的内容，但让我们回到 Mojo 并定义 OurBool 类型。在途中会有机会了解更多关于 MLIR 的知识。

定义 OurBool 类型

我们可以使用 Mojo 的 struct 关键字定义一个新的类型 OurBool。

struct OurBool:
    var value: __mlir_type.i1

一个布尔值可以代表 0 或 1，“true” 或 “false”。为了存储这个信息，OurBool 有一个名为 value 的成员。它的类型直接在 MLIR 中表示，使用 MLIR 内置类型 i1。实际上，您可以在 Mojo 中使用任何 MLIR 类型，只需将类型名称前加上 __mlir_type 即可。

正如我们下面将要看到的，使用 i1 来表示我们的布尔值将使我们能够利用与 i1 类型接口的所有MLIR操作和优化 - 并且它们有很多！

定义了 OurBool 之后，我们现在可以声明一个这种类型的变量：

let a: OurBool

利用 MLIR

自然地，我们可能接下来尝试创建一个 OurBool 的实例。然而，此时尝试这样做会导致错误：

let a = OurBool() # error: 'OurBool' does not implement an '__init__' method

与 Python 一样，__init__ 是一个特殊方法，可以定义来自定义类型的行为。我们可以实现一个不接受任何参数的 __init__ 方法，返回一个值为 “false” 的 OurBool。

struct OurBool:
    var value: __mlir_type.i1
    fn __init__(inout self):
        self.value = __mlir_op.`index.bool.constant`[
            value : __mlir_attr.`false`,
        ]()

为了初始化底层的 i1 值，我们使用来自其 “index” 方言的 MLIR 操作，称为index.bool.constant 。

MLIR的“index”方言为我们提供了操作内置MLIR类型的操作，例如我们用于存储OurBool值的i1。index.bool.constant 操作接受一个 true 或 false 的编译时常量作为输入，并产生具有给定值的类型为 i1 的运行时输出。

因此，如图所示，除了任何 MLIR 类型之外，Mojo 还通过 __mlir_op 前缀直接访问任何 MLIR 操作，并通过 __mlir_attr 前缀直接访问任何属性。MLIR 属性用于表示编译时常量。

如您在上面看到的，与MLIR交互的语法并不总是很漂亮：MLIR属性在方括号[…]之间传递，并且操作通过后缀（…）执行，可以采用运行时参数值。然而，大多数Mojo程序员不需要直接访问MLIR，对于少数需要访问MLIR的程序员来说，这种“丑陋”的语法赋予了他们超能力：他们可以定义易于使用的高级类型，但内部连接到MLIR及其强大的方言系统。

我们认为这非常令人兴奋，但让我们回到现实中来：定义了 __init__ 方法后，我们现在可以创建 OurBool 类型的实例：

let b = OurBool()

Mojo 中的 Value 语义

现在我们可以使用 OurBool 了，但是使用它却是另一回事：

let a = OurBool()
let b = a # error: 'OurBool' does not implement the '__copyinit__' method

Mojo 默认使用“值语义”，这意味着在分配给 b 时，它期望创建一个 a 的副本。但是，Mojo 不对 OurBool 或其底层 i1 值的复制方式做任何假设。错误表明我们应该实现一个 __copyinit__ 方法，该方法将实现复制逻辑。

然而，对于我们的情况来说，OurBool 是一个非常简单的类型，只有一个“可平凡地复制”的成员。我们可以使用装饰器告诉 Mojo 编译器，省去了我们自己定义 __copyinit__ 样板代码的麻烦。可平凡地复制的类型必须实现一个 __init__ 方法，该方法返回自身的实例，因此我们必须稍微重写我们的初始化器。

@register_passable("trivial")
struct OurBool:
    var value: __mlir_type.i1
    fn __init__() -> Self:
        return Self {
            value: __mlir_op.`index.bool.constant`[
                value : __mlir_attr.`false`,
            ]()
        }

现在我们可以根据需要复制 OurBool：

let c = OurBool()
let d = c

编译时常量

拥有只能表示“假”的布尔类型并不是很有用。让我们定义代表 OurBool 值的 true 和 false 编译时常量。

首先，让我们为 OurBool 定义另一个 __init__ 构造函数，它将其 i1 值作为参数：

@register_passable("trivial")
struct OurBool:
    var value: __mlir_type.i1
    # ...
    fn __init__(value: __mlir_type.i1) -> Self:
        return Self {value: value}