Mojo 学习笔记（四）

Posted on April 6, 2024

原本打算按照Mojo的 《Get Started》文档学下来的，因为（三）里面说的类型问题，我想先看看 Mojo 的类型系统。Mojo 里面的类型包括了和Python兼容的 PythonObject 和 Mojo自己的 native object.

PythonObject

一旦Mojo使用了 Python的解释器，那么会得到 PythonObject：

x = Python.evaluate("5 + 10")
print(x)

var py = Python.import_module("builtins")
print(py.type(x))
print(py.id(x))

将输出：

15
<class 'int'>
140553709109936

x 在内存中的表示和我们在python中运行它是完全一样，因此它被定义为一个 PythonObject。实际上在Python中会为每个对象分配内存，而 id() 打印的就是它的内存地址。

【TIP：使用 import “builtins” 的方式，就可以使用python中的所有keyword】

由于Python中的对象可以动态改变其类型，因此实际上在分配的对象内存中就保存着此对象的类型和具体的一些值信息。在Mojo中使用这类PythonObject对象时，实际上和Python的行为是完全一致的，这带来了很多便利，但同时也为这些灵活性带来了代价，因为每个对象中都有许多字段需要额外的内存，也需要CPU 来分配、获取值以及做垃圾收集等。

Mojo object

因此，在Mojo中，在可能的情况下，就不会使用 PythonObject，而是其 native object 类型。

x = 5 + 10
print(x)

这时候，同样会打印 15。但这时候的 x​ 将不再是PythonObject（也即不再是一个指向分配内存的指针），而直接是一个64bit的“整数”。这样就可以带来一系列的优化操作。特别的，Mojo 专门为了提升 Python 的性能而生，而且特别关注发挥现代 CPU 的性能（他们认为目前的业界对CPU的重视程度不够），其实说到底就是想更好地利用现代CPU 的SIMD 特性。

SIMD

SIMD，是指一条指令，多个数据（Single Instruction, Multiple Data)。从早期的Intel x86指令集里面的 MMX，SSE 到现在的AVX, AVX2，包括 Arm 指令集里面的 neon，现在的CPU硬件都有特殊的寄存器和指令集让你通过一条指令同时计算多条数据。早期主要是为了多媒体处理而生的，但实际上目前对于AI 计算也是特别的重要。因此 Mojo 也专门为 SIMD 设计了数据类型。

x = SIMD[DType.uint8, 4](1, 2, 3, 4)
y = SIMD[DType.uint8, 4](10, 20, 30, 40)

a = x+y
print(a)

b = x*y
print(b)

var c=x*10+y 
print(c)

x 和 y 就分别定义了一个 4 个值的向量，每个值是一个uint8 (1 byte)，总共是 32bit， 4bytes。而在Mojo语言中，就会直接支持一些向量的运算。因此上面的代码会输出：

[11, 22, 33, 44]
[10, 40, 90, 160]
[20, 40, 60, 80]

【这里 c 前面必须加 var，否则会报 error: expression must be mutable in assignment，但为啥 a, b 不需要，还不是很明白】

【插入】：后来发现了 c 前面必须加 var 的原因，是因为我测试用的mojo脚本的文件名就是 c.mojo，因此 c 就被当成是代表此脚本的变量了。要是改成 d=x*10+y，实际上也是不需要 var的。闹了一个大乌龙 。如果 c.mojo 如下：

def greet(): 
    print("hello")
def main():
    c.greet()

是会打印出 hello 来的。不知道是bug还是设计如此。

你也可以用一个值来初始化一个向量：

x = SIMD[DType.uint8, 4](18)
print(x)

输出：

[18, 18, 18, 18]

标量类型（Scalar）

有趣的是，Mojo 里面内置的标量类型，其实都是用 SIMD 定义的，也就是1个值的向量而已：

Scalar = SIMD[?, 1]: Represents a scalar dtype.
Int8 = SIMD[si8, 1]: Represents an 8-bit signed scalar integer.
UInt8 = SIMD[ui8, 1]: Represents an 8-bit unsigned scalar integer.
Int16 = SIMD[si16, 1]: Represents a 16-bit signed scalar integer.
UInt16 = SIMD[ui16, 1]: Represents a 16-bit unsigned scalar integer.
Int32 = SIMD[si32, 1]: Represents a 32-bit signed scalar integer.
UInt32 = SIMD[ui32, 1]: Represents a 32-bit unsigned scalar integer.
Int64 = SIMD[si64, 1]: Represents a 64-bit signed scalar integer.
UInt64 = SIMD[ui64, 1]: Represents a 64-bit unsigned scalar integer.
BFloat16 = SIMD[bf16, 1]: Represents a 16-bit brain floating point value.
Float16 = SIMD[f16, 1]: Represents a 16-bit floating point value.
Float32 = SIMD[f32, 1]: Represents a 32-bit floating point value.
Float64 = SIMD[f64, 1]: Represents a 64-bit floating point value.

你可以这样使用：

var x = Int8(1)
y = Int16(1)
var z : Int32 = 1

一个小问题

那么问题来了，最短的话1个数值就是标量，那我们能定义多长的SIMD 变量呢？实验下来似乎没有长度限制，但必须是2的N次方。否则会报错：

var aa = SIMD[DType.int8, 10](1)
print(aa)

会输出：

error: no viable expansion found
note: call expansion failed - no concrete specializations
note: constraint failed: simd width must be power of 2

而下面的语句就没有问题：

var aa = SIMD[DType.int8, 1024](1)
标签: 技术 Mojo