MLIR 新编程语言MOJO

发布时间 2023-05-06 20:20:37作者: myrj

什么是MLIR?

MLIR是程序的中间表示,与汇编语言没有什么不同,在汇编语言中,一组连续的指令对内存中的值进行操作。

更重要的是,MLIR是模块化和可扩展的。MLIR由越来越多的“方言”组成每种方言都定义了操作和优化:例如“数学”方言提供数学运算,如正弦和余弦运算阿姆德普方言提供特定于AMD处理器的操作,等等。

MLIR的每一种方言都可以互通。这就是为什么说MLIR开启了异构计算:随着更新、更快的处理器和架构的开发,新的MLIR方言被实施来为这些环境生成最佳代码。任何新的MLIR方言都可以无缝地翻译成其他方言,因此随着更多方言的加入,所有现有的MLIR都会变得更加强大。

这意味着我们自己的自定义类型,如OurBool类型,可以用来为程序员提供一个高级的、类似Python的接口。但是“在幕后”,Mojo和MLIR将为未来出现的每一款新处理器优化我们方便的高级类型。

关于为什么MLIR是如此革命性的技术,还有很多要写,但是让我们回到魔咒和定义OurBool类型。一路上会有机会更多地了解MLIR。

定义OurBool类型

我们可以用魔咒struct关键字来定义新类型OurBool:

struct OurBool:
    var value: __mlir_type.i1

布尔值可以表示0或1,“真”或“假”为了存储这些信息,OurBool有一个名为的成员value。它的类型被表示直接在MLIR,使用MLIR内置类型i1。事实上,您可以在Mojo中使用任何MLIR类型,只需在类型名前面加上__mlir_type.

正如我们将在下面看到的,用i1将允许我们利用与接口的所有MLIR操作和优化i1类型——而且有很多类型!

定义了OurBool,我们现在可以声明这种类型的变量:

var a: OurBool

利用MLIR

自然,我们接下来可能会尝试创建OurBool。但是,此时尝试这样做将导致错误:

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

和Python一样,__init__是一个特殊方法可以对其进行定义以自定义类型的行为。我们可以实现一个__init__方法,该方法不采用任何参数,并返回OurBool具有“假”值。

struct OurBool:
    var value: __mlir_type.i1

    fn __init__(self&):
        self.value = __mlir_op.`index.bool.constant`[
            value : __mlir_attr.`false`,
        ]()

初始化基础i1值,我们使用MLIR操作从它的“索引”方言,名为index.bool.constant.

MLIR的“索引”方言为我们提供了操作内置MLIR类型的操作,例如i1我们用来存储的值OurBool。这index.bool.constant操作需要true或者false编译时常数作为输入,并生成i1用给定的值。

因此,如上所示,除了任何MLIR类型,Mojo还通过__mlir_op前缀,并通过__mlir_attr前缀。MLIR属性用于表示编译时常数。

正如您在上面看到的,与MLIR交互的语法并不总是很好:MLIR属性在方括号之间传递[...],操作通过括号后缀来执行(...),它可以接受运行时参数值。然而,大多数Mojo程序员不需要直接访问MLIR,对于少数这样做的人来说,这种“丑陋”的语法给了他们超能力:他们可以定义易于使用的高级类型,但可以在内部插入MLIR及其强大的方言系统。

我们认为这非常令人兴奋,但是让我们回到现实:定义了一个__init__方法,我们现在可以创建我们的OurBool类型:

let b = OurBool()

Mojo中的值语义

我们现在可以实例化OurBool,但使用它是另一回事:

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

默认情况下,Mojo使用“值语义”,这意味着它期望创建a分配给时b。然而,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

编译时常数

拥有一个只能表示“false”的布尔类型不是很有用让我们定义表示真和假的编译时常数OurBool价值观。

首先,让我们定义另一个__init__的构造函数OurBool那需要时间i1作为参数的值:

@register_passable("trivial")
struct OurBool:
    var value: __mlir_type.i1
    # ...

    fn __init__(value: __mlir_type.i1) -> Self:
        return Self {value: value}

这允许我们定义编译时常数OurBool值,使用alias关键词。首先,我们来定义一下OurTrue:

alias OurTrue = OurBool(__mlir_attr.`true`)

这里我们传入一个MLIR编译时常量值true,它具有i1键入我们的新__init__构造函数需要。我们可以使用稍微不同的语法OurFalse:

alias OurFalse: OurBool = __mlir_attr.`false`

OurFalse被声明为类型OurBool,然后分配一个i1类型–在这种情况下OurBool我们添加的构造函数被隐式调用。

有了真常数和假常数,我们也可以简化我们原来的__init__的构造函数OurBool。我们可以简单地返回我们的,而不是构造MLIR值OurFalse常数:

alias OurTrue = OurBool(__mlir_attr.`true`)
alias OurFalse: OurBool = __mlir_attr.`false`

@register_passable("trivial")
struct OurBool:
    # ...
    fn __init__() -> Self:
        return OurFalse

还要注意,我们可以定义OurTrue在我们定义之前OurBool。Mojo编译器足够聪明来解决这个问题。

有了这些常量,我们现在可以用真值和假值来定义变量OurBool:

let e = OurTrue
let f = OurFalse

执行__bool__

当然,布尔在编程中无处不在的原因是它们可以用于程序控制流。然而,如果我们试图使用OurBool这样,我们得到一个错误:

let a = OurTrue
if a: print("It's true!") # error: 'OurBool' does not implement the '__bool__' method

当Mojo试图执行我们的程序时,它需要能够决定是否打印“这是真的!”或者不是。它还不知道OurBool表示一个布尔值——Mojo只看到一个大小为1位的结构。然而,Mojo也提供了传递布尔特性的接口,这与Mojo的标准库类型所使用的接口相同,比如Bool。实际上,这意味着Mojo给了你完全的控制权:任何与语言的标准库打包在一起的类型都是你可以定义自己版本的类型。

在我们的错误消息中,Mojo告诉我们实现一个__bool__方法打开OurBool将表明它具有布尔性质。

谢天谢地,__bool__实现起来很简单:Mojo的标准库和内置类型都是在MLIR之上实现的,所以内置Bool类型还定义了一个采用i1,就像OurBool:

@register_passable("trivial")
struct OurBool:
    var value: __mlir_type.i1
    # ...

    fn __bool__(self) -> Bool:
        return Bool(self.value)

现在我们可以使用OurBool任何我们可以使用内置的地方Bool类型:

let g = OurTrue
if g: print("It's true!")
It's true!

使用避免类型转换__mlir_i1__

我们的OurBool类型看起来很棒,通过提供到Bool,它可以在内置的任何地方使用Bool类型可以。但是在上一节中,我们向您承诺了“完全控制”,即定义内置于Mojo或其标准库中的任何类型的您自己的版本的能力。无疑Bool不实现__bool__把自己变成Bool?

事实上并不是这样:当Mojo计算一个条件表达式时,它实际上试图将其转换成一个MLIRi1值,通过搜索特殊的接口方法__mlir_i1__。(自动转换为Bool发生原因是Bool已知实现了__mlir_i1__方法。)

同样,Mojo被设计成可扩展和模块化的。通过实现所有的特殊方法Bool我们可以创造一种类型来完全取代它。让我们通过实现__mlir_i1__OurBool:

@register_passable("trivial")
struct OurBool:
    var value: __mlir_type.i1
    # ...

    fn __mlir_i1__(self) -> __mlir_type.i1:
        return self.value

我们仍然可以使用OurBool就像我们之前做的那样:

let h = OurTrue
if h: print("No more Bool conversion!")
No more Bool conversion!

但是这一次,没有转换到Bool发生。你可以尝试添加print致大会的声明__bool____mlir_i1__方法,甚至移除__bool__方法,自己去看。

使用MLIR添加功能

我们还有很多方法可以改进OurBool。其中许多都涉及到实现特殊的方法,有些您可能在Python中见过,有些是特定于Mojo的。例如,我们可以实现OurBool通过添加一个__invert__方法。我们还可以添加一个__eq__方法,该方法允许两个OurBool要与==接线员。

让Mojo与众不同的是,我们可以使用MLIR来实现这些功能。实施__eq__例如,我们使用index.casts铸造我们的操作i1MLIR索引方言的值index键入,然后index.cmp比较它们是否相等的操作:

@register_passable("trivial")
struct OurBool:
    var value: __mlir_type.i1
    # ...

    fn __eq__(self, rhs: OurBool) -> Self:
        let lhsIndex = __mlir_op.`index.casts`[_type : __mlir_type.index](
            self.value
        )
        let rhsIndex = __mlir_op.`index.casts`[_type : __mlir_type.index](
            rhs.value
        )
        return Self(
            __mlir_op.`index.cmp`[
                pred : __mlir_attr.`#index<cmp_predicate eq>`
            ](lhsIndex, rhsIndex)
        )

然后我们可以实现__invert__根据__eq__:

@register_passable("trivial")
struct OurBool:
    # ...
    fn __invert__(self) -> Self:
        return OurFalse if self == OurTrue else OurTrue

这允许我们使用~运算符withOurBool:

let i = OurFalse
if ~i: print("It's false!")
It's false!

这种可扩展的设计甚至允许“内置”的Mojo类型,如BoolInt,甚至Tuple(!!)将根据MLIR在Mojo标准库中实现,而不是硬编码到Mojo语言中。这也意味着这些类型几乎没有什么是用户定义的类型所不能实现的。

推而广之,这意味着Mojo为机器学习工作流带来的令人难以置信的性能并不是由于幕后执行的某种魔法——你可以定义自己的高级类型,在实现中使用低级MLIR来实现前所未有的速度和控制。

模块化的承诺

正如我们所见,Mojo与MLIR的集成允许Mojo程序员实现与Mojo自己的内置和标准库类型同等的零成本抽象。

MLIR是开源的和可扩展的:新的方言一直在增加,然后这些方言就可以在Mojo中使用了。与此同时,Mojo代码变得更加强大,并针对新硬件进行了优化——Mojo程序员无需额外工作。

这意味着您自己的自定义类型,无论是OurBool或者OurTensor,可以用来给程序员提供一个易于使用且不变的界面。但在幕后,MLIR将为明天的计算环境优化那些方便的高级类型。

换句话说:Mojo不是魔术,它是模块化的。