公開されたので早速 Mojo に入門する

Mojo が一般公開されたので、早速使ってみる。
M1 Mac では使えないそうで、GitHub codespaces を使う。
https://developer.modular.com/login
ここからログインして、始め方をみる。

https://github.com/modularml/mojo
をフォークして、codespacesを作成。
VSCodeの拡張機能を入れる。

このスクラップは多分に
https://docs.modular.com/
から引用しています。

ひとまず

mojo hoge.mojo

で実行できるのを確認。コードの雰囲気を掴んでいく。

Python のスーパーセットになっているはず。逆にPython以外の書き回しになっている箇所が高速化されている。

let: 不変var: 可変が高速な変数宣言の仕方。
fn が型安全な関数、defが従来のPythonの関数。

端的な書き方：

fn do_math():
    let x: Int = 1 # 型の明示的な宣言
    let y = 2 # 型推論される
    print(x + y)

fn add(x: Int, y: Int) -> Int: # 型アノテーション
    return x + y

fn main(): # main がエントリ
    var x: Int = 1
    x += 1
    print(x)

Taqqn

Swift作者だけあって、所有権周りはこだわりを感じる。SwiftはGCの代わりにARC(Automatic Reference Counting) を導入してる。
Rust ほど、たくさん書くことはなく、Pythonの書き味を残している印象。

fn add(borrowed x: Int, borrowed y: Int) -> Int: # borrowed はfnのデフォルトで、省略可能。中では変更不可能。
    return x + y

fn add_inout(inout x: Int, inout y: Int) -> Int:
    x += 1
    y += 1
    return x + y

var a = 1
var b = 2
c = add_inout(a, b)
print(a)
print(b)
print(c)

fn set_fire(owned text: String) -> String:
    text += "🔥"
    return text

fn mojo():
    let a: String = "mojo"
    let b = set_fire(a)
    print(a)
    print(b)

mojo()

transfer 演算子(C++, Rust でいうmove) が^となっていて使いやすい。Rustのようにデフォルトでmove出ないのは正解に思える。C++のstd::moveより簡潔。

let b = set_fire(a^)

一旦記事にまとめた。

Taqqn

Python モジュールのimport。便利そう。

from python import Python

let np = Python.import_module("numpy")

ar = np.arange(15).reshape(3, 5)
print(ar)
print(ar.shape)

Taqqn

mojo はコンパイラ設計用の拡張中間表現MLIR にアクセスできる。MLIRには、math方言やamdgpu方言などが存在するが、相互運用可能らしい。
この辺りは後で勉強。

struct OurBool:
    var value: __mlir_type.i1 # MLIRのbooleanプリミティブ

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

Taqqn

move の仕様と__moveinit__を使えば、例えばpointerを触ってる箇所を唯一的にできる。
そうでなくても、moveの時だとメモリのallocationなどを効率化できたりしてよさそう。

# This is not really a unique pointer, we just model its behavior here
# to serve the examples below.
struct UniquePointer:
    var ptr: Int
    
    fn __init__(inout self, ptr: Int):
        self.ptr = ptr
    
    fn __moveinit__(inout self, owned existing: Self):
        self.ptr = existing.ptr
        
    fn __del__(owned self):
        self.ptr = 0

わかりやすい、MyIntの実装。

struct MyInt:
    var value: Int
    
    fn __init__(inout self, v: Int):
        self.value = v

    fn __copyinit__(inout self, other: MyInt):
        self.value = other.value
        
    # self and rhs are both immutable in __add__.
    fn __add__(self, rhs: MyInt) -> MyInt:
        return MyInt(self.value + rhs.value)
        
    # ... now this works:
    fn __iadd__(inout self, rhs: MyInt):
        self = self + rhs

Taqqn

`@register_passable`

@register_passable をstruct につけると参照渡しではなく値わたしとなり、小さいメモリのデータに対しては効率化可能。

@register_passable
struct Pair:
    var a: UInt32
    var b: UInt32

    fn __init__(a: UInt32, b: UInt32) -> Self:
        return Self{a: a, b: b}
fn main():
    let x = Pair(5, 10)
    print(x.a)

@register_passable な時は、

fn __init__(inout self) の最初のself がなくなる。
Self{a: a, b: b} のようなinitializer syntaxが使用可能になる。
コピー操作が前提となるので、__moveinit__が実装できなくなる。

基本的なコンストラクタを自動で実装する@value を用いて、

@value
@register_passable
struct Pair:
    var a: Int
    var b: Int

としてもよい。

`@register_passable("trivial")`

さらに、
@register_passable("trivial") とすると

__init__
__copyinit__
__moveinit__
__del__
が定義できなくなる。自明なコピーしか許さない、本当にシンプルな構造体となる。

Taqqn

@always_inline("nodebug")
よくあるinline化してデバッグ情報を落とすやつ。

Taqqn

`alias`

コンパイル時定数として使える。
Mojoでは型もコンパイル時定数。したがって、型のエイリアスにもaliasを使う。

Taqqn

[size: Int] がジェネリクスのパラメータ (Intも使える。型ならDType) となっていて@parameter if 担っている場所をコンパイル時に解決する。次元による場合分けなどによさそう

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)

Taqqn

コンパイル時に決定される、関数内の関数

fn buffer_elementwise_add_impl[
    dt: DType
](lhs: DTypePointer[dt], rhs: DTypePointer[dt], result: DTypePointer[dt], N: Int):
    """Perform elementwise addition of N elements in RHS and LHS and store
    the result in 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)

Taqqn

関数のオーバーロードの解決

@register_passable("trivial")
struct MyInt:
    """A type that is implicitly convertible to `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: MyInt, a: Int]()` is called because it requires no implicit
    # conversions, whereas `foo[x: MyInt, y: MyInt]()` requires one.
    foo[x, a]()

    # `bar[a: Int](b: Int)` is called because it does not have variadic
    # arguments or parameters.
    bar[a](b)

    # `bar[*a: Int](b: Int)` is called because it has variadic parameters.
    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)

Taqqn

object は TypeScript の any 的なやつかな？
例外を投げうるときは raises

fn matrix_getitem(self: object, i: object) raises -> object:
    return self.value[i]


fn matrix_setitem(self: object, i: object, value: object) raises -> object:
    self.value[i] = value
    return None


fn matrix_append(self: object, value: object) raises -> object:
    self.value.append(value)
    return None

Taqqn

さっきからちょくちょく出てくる SIMD (Single Instruction, Multiple Data)

以下のように全体の要素に対する掛け算が可能。ハードウェアレベルで最適化される。

from DType import DType

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

ハードウェアでサポートされてる(小さい)ベクトル要素。
複数のデータに対して同じ命令を同時に実行することができるコンピュータアーキテクチャの一種。データの並列処理を効率化可能。

alias は以下:

Int8 = SIMD[si8, 1]: 8 ビットの符号付きスカラー整数
UInt8 = SIMD[ui8, 1]: 8 ビットの符号なしスカラー整数
Int16 = SIMD[si16, 1]: 16 ビットの符号付きスカラー整数
UInt16 = SIMD[ui16, 1]: 16 ビットの符号なしスカラー整数
Int32 = SIMD[si32, 1]: 32 ビットの符号付きスカラー整数
UInt32 = SIMD[ui32, 1]: 32 ビットの符号なしスカラー整数
Int64 = SIMD[si64, 1]: 64 ビットの符号付きスカラー整数
UInt64 = SIMD[ui64, 1]: 64 ビットの符号なしスカラー整数
Float16 = SIMD[f16, 1]: 16 ビット浮動小数点数
Float32 = SIMD[f32, 1]: 32 ビット浮動小数点数
Float64 = SIMD[f64, 1]: 64 ビット浮動小数点数

Taqqn

Tensor は別に存在。多次元配列が標準ライブラリになっている。しかし、加減乗除は定義されていない。

参考までに以下のような実装が存在。

Taqqn

DType.float64 などと型を参照する、DType は class であると書かれている。

Taqqn

https://mojodojo.dev/guides/
とてもよくまとまっている。
基本的な型から、copy move の話とか

Taqqn

Mojo には3種類のstructがある。

ムーブ・コピー不可
ムーブ可能
- __moveinit__ の実装が条件
コピー・ムーブ可能
- __copyinit__ の実装が条件
- ムーブの時には__moveinit__ があればそちらがよばれるし、なければ__copyinit__が呼ばれる

Taqqn

毎回 __init__ __moveinit__, __copyinit__ を実装するのは大変なので、@value デコレータが存在する。Kotlin の data class 的なやつ。

@value
struct Pair:
    var x: Int
    var y: Int

fn main():
  let pair = Pair(5, 10)

  # Move object
  var pair2 = pair^
  # Copy object
  let pair3 = pair2
  # Edit original
  pair2.x = 20

  # Print both the original and copy
  print(pair2.x)
  print(pair3.x)

Taqqn

__moveinit__ が思ったより呼ばれないと思ったら、__del__が実装されてないからだった。以下にまとめた。

Taqqn

ListLiteral と Tuple

型パラメータで型安全性が意識されている。

let list: ListLiteral[Int, FloatLiteral, StringLiteral] = [1, 5.0, "Mojo🔥"]
print(list.get[2, StringLiteral]()) # Mojo🔥

let tup: Tuple[Int, StringLiteral, FloatLiteral] = (1, "Mojo", 3.0)
print(tup.get[0, Int]()) # 1

Taqqn

いつかの Cython 記事との比較のために、実装。
参考:

from utils.vector import DynamicVector

fn vec():
  var sum: Int = 0
  var l = DynamicVector[Int]()
  for i in range(1000000):
    l.push_back(i*2)
  for i in range(1000000):
    sum += l[i]
  print(sum)

fn main():
  vec()

コアダンプした。1000000を10にしても落ちた。

$ mojo vec.mojo
mojo: /__w/modular/modular/third-party/llvm-project/llvm/include/llvm/ADT/ArrayRef.h:443: T &llvm::MutableArrayRef<mlir::OpOperand>::operator[](size_t) const [T = mlir::OpOperand]: Assertion `Index < this->size() && "Invalid index!"' failed.
[5282:5282:20230912,113418.139653:ERROR file_io_posix.cc:144] open /sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq: No such file or directory (2)
[5282:5282:20230912,113418.139715:ERROR file_io_posix.cc:144] open /sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq: No such file or directory (2)
Please submit a bug report to https://github.com/modularml/mojo/issues and include the crash backtrace along with all the relevant source codes.
 #0 0x0000560045738957 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x5bb957)
 #1 0x000056004573652e (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x5b952e)
 #2 0x000056004573902f (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x5bc02f)
 #3 0x00007facd1b73420 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x14420)
 #4 0x00007facd15fc00b raise /build/glibc-SzIz7B/glibc-2.31/signal/../sysdeps/unix/sysv/linux/raise.c:51:1
 #5 0x00007facd15db859 abort /build/glibc-SzIz7B/glibc-2.31/stdlib/abort.c:81:7
 #6 0x00007facd15db729 get_sysdep_segment_value /build/glibc-SzIz7B/glibc-2.31/intl/loadmsgcat.c:509:8
 #7 0x00007facd15db729 _nl_load_domain /build/glibc-SzIz7B/glibc-2.31/intl/loadmsgcat.c:970:34
 #8 0x00007facd15ecfd6 (/lib/x86_64-linux-gnu/libc.so.6+0x33fd6)
 #9 0x0000560046d9ee90 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x1c21e90)
#10 0x00005600471fbe86 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x207ee86)
#11 0x00005600471fc511 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x207f511)
#12 0x0000560047201b52 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x2084b52)
#13 0x0000560047201c4f (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x2084c4f)
#14 0x0000560045cd3232 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0xb56232)
#15 0x0000560047edb376 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x2d5e376)
#16 0x0000560047edc0b1 (/home/codespace/.modular/pkg/packages.modular.com_mojo/bin/mojo+0x2d5f0b1)
#17 0x00007facd1b67609 start_thread /build/glibc-SzIz7B/glibc-2.31/nptl/pthread_create.c:478:7
#18 0x00007facd16d8133 clone /build/glibc-SzIz7B/glibc-2.31/misc/../sysdeps/unix/sysv/linux/x86_64/clone.S:97:0
Aborted (core dumped)

どうやら DynamicVector にバグがあるらしい。様子を見る。