見出し画像

Pythonを過去のものにするMojo🔥が爆誕

2023/10/28更新
アップルシリコンが Python 素の状態の9万倍のベンチマーク結果が出ています。


2010年代前半、AI業界全体がpythonを基本インフラとして活動してきたこともあり、ChatGPT時代に入った今日現在もpythonが主流です。しかし、元々pythonは処理速度よりも生産性を優先させたインタープリタ型のscriptに過ぎず、CPUやGPUの資源を効率的に発揮させてるとは言い難い状況です。
ChatGPTなどもpythonで実装されてますが、今後、Mojoに乗り換えることで桁違いの効率性とChat出力速度が期待されます。


なぜ Mojo🔥 なのか?


言語の構文やエコシステムは革新する必要はない

既存のPythonのエコシステムを採用。Pythonのライブラリなども使用可。
既にPython知ってるエンジニアであれば、新しく覚えることは少なく、大半のエンジニアの習得コストも僅かです。

ChatGPTなど応用的なAIも一つの言語でカバー

PythonはCPUのマルチコア、GPGPUなどのH/Wアクセラレーション連携を最初から網羅されて設計されていません。(※そもそも設計当初、GPGPUは存在しなかった)Mojo🔥は、これらH/Wアクセラレーションを取り込んで再設計。

Pythonの構文で、C++やRustのような高性能も併せ持つ

Mojoは、Pythonのように使いやすく、C++やRustのような性能を持つプログラミング言語です。さらに、MojoはPythonのライブラリエコシステム全体を活用する機能を提供します。
Pythonの動的機能を維持しながら、システムプログラミングのための新しいプリミティブを追加することで、Pythonのスーパーセットとなるように設計されています。これらの新しいシステムプログラミングのプリミティブにより、Mojoの開発者は、現在C、C++、Rust、CUDA、その他のアクセラレータシステムを必要とする高性能なライブラリを構築することができます。動的言語とシステム言語の長所を組み合わせることで、抽象度を超えて機能し、初心者プログラマーに優しく、加速器からアプリケーションプログラミングやスクリプトまで、多くのユースケースでスケールする統一プログラミングモデルを提供します。


Pythonエンジニア向けのMojo🔥 / hello world

そもそも覚えることは少ないので、例文提示を列挙します。

拡張子は *.mojo です。しかも絵文字(emoji)対応で
*.🔥 でも可です。ただ、絵文字だと grepやfindとの相性も悪いので、*.mojoあたりが無難かと個人的には思います。

$ mojo hello.mojo

いわゆる hello world

def main():
    print("hello world")
    for x in range(9, 0, -3):
        print(x)


let宣言とvar宣言

Mojoのdefの内部では、名前に値を代入することができ、Pythonのように暗黙のうちに関数スコープの変数が作成されます。これは非常にダイナミックでセレモニー性の低いコードの書き方を提供しますが、2つの理由で難点があります。

システムプログラマは、型安全性やパフォーマンスのために、値が不変であることを宣言したいことがよくあります。
システムプログラマは、型安全性とパフォーマンスのために、値が不変であることを宣言したいと思うことがよくあります。
これをサポートするために、Mojoはスコープ付きの実行時値宣言を提供します。letはimmutableで、varはmutableです。これらの値はレキシカルスコープを使用し、名前のシャドーイングをサポートします。

def your_function(a, b):
    let c = a
    c = b  # error: c is immutable

    if c != b:
        var c = b
        stuff()
def your_function():
    let x: Int8 = 42
    let y: Int64 = 17

    let z: Int8
    if x != 0:
        z = 1
    else:
        z = foo()
    use(z)

letとvarはdef宣言の中では完全にオプトインであることに注意してください。

構造体

Mojoは、多くのプログラミング言語で使用されている最先端のコンパイラとコード生成システムを提供するMLIRとLLVMをベースにしています。これにより、データ構成やデータフィールドへの直接アクセスなど、性能を向上させるための制御をより適切に行うことができます。現代のシステムプログラミング言語の重要な特徴は、このような複雑な低レベルの操作の上に、性能を落とすことなく高レベルで安全な抽象化を構築することができることです。Mojo では、これは struct 型によって提供されます。

Mojo の struct は Python のクラスに似ています。どちらもメソッド、フィールド、演算子のオーバーロード、メタプログラミングのためのデコレータなどをサポートしています。

class vs struct :

Python クラスは動的です。動的なディスパッチ、モンキーパッチ (または「スウィズリング」)、実行時にインスタンスプロパティを動的に結合することができます。

Mojoの構造体は静的です。コンパイル時に結合されます(実行時にメソッドを追加することはできません)。構造体を使用すると、安全で使いやすい一方で、柔軟性とパフォーマンスを両立させることができます。

以下は、構造体の簡単な定義です:

@value
struct MyPair:
    var first: Int
    var second: Int
    def __lt__(self, rhs: MyPair) -> Bool:
        return self.first < rhs.first or
              (self.first == rhs.first and
               self.second < rhs.second)

構文上、Pythonのクラスと比較した場合の最大の違いは、「構造体」のインスタンスプロパティはすべてvarまたはlet宣言で明示的に宣言する必要があることです。

Mojoでは、「構造体」の構造や内容はあらかじめ設定されており、プログラムの実行中に変更することはできません。オブジェクトの属性をその場で追加、削除、変更できるPythonとは異なり、Mojoでは構造体にはそれができません。つまり、プログラムを実行している途中でdelを使ってメソッドを削除したり、値を変更したりすることができないのです。

しかし、struct の静的な性質には大きな利点があります。Mojoがコードをより速く実行できるようになります。プログラムは、構造体の情報がどこにあり、どのように使用するかを正確に把握しているので、余分な手順や遅延が発生しません。

Mojoの構造体は、演算子のオーバーロード(+や-などの数学記号が自分のデータでどのように動作するかを変更できる)など、Pythonですでにご存知の機能とも非常にうまく連携しています。さらに、すべての「標準型」(Int、Bool、String、Tupleなど)は、構造体を使用して作られています。これは、構造体が言語そのものに組み込まれているのではなく、使用できるツールの標準セットの一部であることを意味します。これにより、コードを書く際に、より柔軟でコントロールしやすくなります。

Int vs int

Mojoでは、Pythonのint(小文字の "i")とは異なるInt(大文字の "I")となり、この違いは意図的なものです。

Pythonでは、int型は非常に大きな数値を扱うことができ、2つの数値が同じオブジェクトであるかどうかをチェックするなどの特別な機能を備えていますが、その反面、非効率でした。MojoのIntは違います。シンプルで高速、そしてコンピュータのハードウェアが素早く処理できるように調整された設計になっています。

ハードウェア資源と密接に連携する必要があるプログラマ(システム・プログラマ)に、ハードウェアとやりとりするための透明で信頼できる方法を提供します。

Mojoが問題を起こすことなくPythonとの連携を維持するために、異なる名前(intの代わりにInt)を使用することで、Pythonのintの動作方法を変更することで両立を維持しています。

さらに、Intは、Mojoで作成する他のカスタムデータ型と同じ命名スタイルに従います。IntはMojoの標準ツールセット(+や-などの数学記号が独自のデータでどのように機能するか)に含まれている構造体です。すべての「標準型」(Int、Bool、String、Tupleなど)は、構造体を使用して作られています。これは、構造体が言語そのものに組み込まれているのではなく、使用できるツールの標準セットの一部であることを意味します。

マンデルブロー集合をサンプルにして比較

まずは典型的なマンデルブロー描画でpythonと比較

Mojoは高性能なコードを書くのに適しているだけでなく、ライブラリやツールの巨大なPythonエコシステムを活用することも可能です。Pythonのシームレスな相互運用性により、Mojoはパフォーマンスを犠牲にすることなく、Pythonを得意とするもの、特にGUIに使用することができます。古典的なマンデルブロ集合アルゴリズムを取り上げ、Mojoで比較してみましょう。

複素数型を下記のように定義してみます。

@register_passable("trivial")
struct Complex:
    var real: F32
    var imag: F32

    fn __init__(real: F32, imag: F32) -> Self:
        return Self {real: real, imag: imag}

    fn __add__(lhs, rhs: Self) -> Self:
        return Self(lhs.real + rhs.real, lhs.imag + rhs.imag)

    fn __mul__(lhs, rhs: Self) -> Self:
        return Self(
            lhs.real * rhs.real - lhs.imag * rhs.imag,
            lhs.real * rhs.imag + lhs.imag * rhs.real,
        )

    fn norm(self) -> F32:
        return self.real * self.real + self.imag * self.imag

半径2の複素円から「脱出」するまで、各画素に対して反復的に複素関数を計算し、脱出するまでの反復回数を数えるという、マンデルブロの中核となるアルゴリズムを書いてみます。

alias xmin: F32 = -2
alias xmax: F32 = 0.6
alias xn = 960
alias ymin: F32 = -1.5
alias ymax: F32 = 1.5
alias yn = 768

# Compute the number of steps to escape.
def mandelbrot_kernel(c: Complex) -> Int:
    max_iter = 200
    z = c
    for i in range(max_iter):
        z = z * z + c
        if z.norm() > 4:
            return i
    return max_iter


def compute_mandelbrot() -> Matrix:
    # create a matrix. Each element of the matrix corresponds to a pixel
    result = Matrix(xn, yn)

    dx = (xmax - xmin) / xn
    dy = (ymax - ymin) / yn

    y = ymin
    for j in range(yn):
        x = xmin
        for i in range(xn):
            result[i, j] = mandelbrot_kernel(Complex(x, y))
            x += dx
        y += dy
    return result

脱出するための反復回数を色でプロットすると、典型的なマンデルブロ集合のプロットが得られます。これを描画するために、Mojoから直接Pythonのmatplotlibを利用することができます。

def make_plot(m: Matrix):
    np = Python.import_module("numpy")
    plt = Python.import_module("matplotlib.pyplot")
    colors = Python.import_module("matplotlib.colors")
    dpi = 64
    width = 10
    height = 10 * yn // xn

    fig = plt.figure(1, [width, height], dpi)
    ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], False, 1)

    light = colors.LightSource(315, 10, 0, 1, 1, 0)

    image = light.shade(m.to_numpy(), plt.cm.hot, colors.PowerNorm(0.3), "hsv", 0, 0, 1.5)
    plt.imshow(image)
    plt.axis("off")
    plt.show()

make_plot(compute_mandelbrot())
print("finished")
引用:https://docs.modular.com/mojo/notebooks/Mandelbrot.html

マンデルブローのベクトル化

マンデルブロアルゴリズムの素の実装をしましたが、高速化するためにできることが2つあります。ピクセルのエスケープが判明したときにループの反復を早期に停止することと、ループをベクトル化して複数のピクセルを同時に計算することにより、Mojoのハードウェアへのアクセスを活用することができます。そのために、vectorize高次ジェネレータを使用します。

まず、メインの反復ループをベクトル化する方法で定義します。

alias MAX_ITERS = 200
fn mandelbrot_kernel_simd[simd_width:Int](c: ComplexGenericSIMD[DType.f32, simd_width]) -> SIMD[DType.f32, simd_width]:
    var z = c
    var nv = SIMD[DType.f32, simd_width](0)
    var escape_mask = SIMD[DType.bool, simd_width](0)

    for i in range(MAX_ITERS):
        if escape_mask: # All the elements have escaped, so exit.
            break
        z = z*z + c
        # Only update elements that haven't escaped yet
        escape_mask = escape_mask.select(escape_mask, z.norm() > 4)
        nv = escape_mask.select(nv, nv + 1)
    
    return nv

上記の関数は、simd_widthをパラメータとして、simd_widthのピクセルを処理します。ベクトルレーン内のすべてのピクセルを処理し終わると、エスケープされます。上記と同じ反復ループを使用することができますが、今回は代わりに各行内でベクトル化します。vectorizeジェネレータを使用し、単純な関数呼び出しにします。

from Functional import vectorize
from Math import iota
from TargetInfo import dtype_simd_width


def compute_mandelbrot_simd() -> Matrix:
    # create a matrix. Each element of the matrix corresponds to a pixel
    var result = Matrix(xn, yn)

    let dx = (xmax - xmin) / xn
    let dy = (ymax - ymin) / yn

    var y = ymin
    alias simd_width = dtype_simd_width[DType.f32]()

    for row in range(yn):
        var x = xmin
        @parameter
        fn _process_simd_element[simd_width:Int](col: Int):
            let c = ComplexGenericSIMD[DType.f32, simd_width](dx*iota[simd_width, DType.f32]() + x, y)
            result.store[simd_width](col, row, mandelbrot_kernel_simd[simd_width](c))
            x += simd_width*dx

        vectorize[simd_width, _process_simd_element](xn)
        y += dy
    return result


make_plot(compute_mandelbrot_simd())
print("finished")

マンデルブローの並列化

上記のベクトル化された実装は効率的ですが、行を並列化することでより良い性能を得ることができます。これもMojoのparallelize higher order functionを使えば簡単です。呼び出しを実行する関数だけを変更する必要があります。

from Functional import parallelize 

def compute_mandelbrot_simd_parallel() -> Matrix:
    # create a matrix. Each element of the matrix corresponds to a pixel
    var result = Matrix(xn, yn)

    let dx = (xmax - xmin) / xn
    let dy = (ymax - ymin) / yn

    alias simd_width = dtype_simd_width[DType.f32]()

    @parameter
    fn _process_row(row:Int):
        var y = ymin + dy*row
        var x = xmin
        @parameter
        fn _process_simd_element[simd_width:Int](col: Int):
            let c = ComplexGenericSIMD[DType.f32, simd_width](dx*iota[simd_width, DType.f32]() + x, y)
            result.store[simd_width](col, row, mandelbrot_kernel_simd[simd_width](c))
            x += simd_width*dx
            
        vectorize[simd_width, _process_simd_element](xn)

    parallelize[_process_row](yn)
    return result


make_plot(compute_mandelbrot_simd_parallel())
print("finished")


行列の乗算をMojoで実装

Mojoで行列の乗算(matmul)アルゴリズムを記述する方法を説明します。純粋なPythonの実装から始め、次に型を追加し、実装をベクトル化、タイリング、並列化することによって最適化を続けます。

Pythonでの実装

%%python
def matmul_python(C, A, B):
    for m in range(C.rows):
        for n in range(C.cols):
            for k in range(A.cols):
                C[m, n] += A[m, k] * B[k, n]

128×128の正方形行列を使った実装をベンチマークし、達成されたGFLopsを計測してみましょう。

%%python
import numpy as np
from timeit import timeit

class Matrix:
    def __init__(self, value, rows, cols):
        self.value = value
        self.rows = rows
        self.cols = cols
        
    def __getitem__(self, idxs):
        return self.value[idxs[0]][idxs[1]]
    
    def __setitem__(self, idxs, value):
        self.value[idxs[0]][idxs[1]] = value

def benchmark_matmul_python(M, N, K):
    A = Matrix(list(np.random.rand(M, K)), M, K)
    B = Matrix(list(np.random.rand(K, N)), K, N)
    C = Matrix(list(np.zeros((M, N))), M, N)
    secs = timeit(lambda: matmul_python(C, A, B), number=2)/2
    gflops = ((2*M*N*K)/secs) / 1e9
    print(gflops, "GFLOP/s")
    return gflops
python_gflops = benchmark_matmul_python(128, 128, 128).to_f64()

0.0016717199881536883 GFLOP/s

python の素の状態だとほとんど性能は出ません。

Pythonの実装をMojoにインポート

Mojoの使い方はPythonと同じで簡単です。まず、Mojoのstdlibから、これから使うそのモジュールをインクルードしましょう:

ユーティリティのインポートとマトリックスの定義 (クリックで表示/非表示)
そして、Pythonのコードをコピー&ペーストします。MojoはPythonのスーパーセットなので、同じPythonのコードがMojoのコードとして実行されます。

# This exactly the same Python implementation, 
# but is infact Mojo code!
def matmul_untyped(C, A, B):
    for m in range(C.rows):
        for n in range(C.cols):
            for k in range(A.cols):
                C[m, n] += A[m, k] * B[k, n]

128x128行列乗算をベンチマークをしてみると、

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


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


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


def matrix_init(rows: Int, cols: Int) -> object:
    value = object([])
    return object(
        Attr("value", value), Attr("__getitem__", matrix_getitem), Attr("__setitem__", matrix_setitem), 
        Attr("rows", rows), Attr("cols", cols), Attr("append", matrix_append),
    )

def benchmark_matmul_untyped(M: Int, N: Int, K: Int, python_gflops: F64):
    C = matrix_init(M, N)
    A = matrix_init(M, K)
    B = matrix_init(K, N)
    for i in range(M):
        c_row = object([])
        b_row = object([])
        a_row = object([])
        for j in range(N):
            c_row.append(0.0)
            b_row.append(random_f64(-5, 5))
            a_row.append(random_f64(-5, 5))
        C.append(c_row)
        B.append(b_row)
        A.append(a_row)

    @parameter
    fn test_fn():
        try:
            _ = matmul_untyped(C, A, B)
        except:
            pass

    let secs = F64(Benchmark().run[test_fn]()) / 1_000_000_000
    let gflops = ((2*M*N*K)/secs) / 1e9
    let speedup : F64 = gflops / python_gflops
    print(gflops, "GFLOP/s, a", speedup.value, "x speedup over Python")

0.029258 GFLOP/s, a 17.501798 x speedup over Python

まだまだですが、Pythonの素の状態に比較して約17.5倍です。
素のPythonよりも優れたパフォーマンスを達成しているものの、Mojoから得られる恩恵をまだ受けてないです。Mojoに入力の型を教えれば、コードの多くを最適化してディスパッチコストを削減できます(型を型チェックにしか使わないPythonとは異なり、Mojoは型情報を利用してパフォーマンスの最適化も行います)。

その最初の一歩として、まずMatrix構造体を定義します。Matrix構造体には、データポインタとサイズフィールドが含まれています。Matrix構造体は任意のデータ型にパラメータを設定できますが、ここでは簡潔にするためにデータ型をf32に設定しています。

struct Matrix:
    var data: DTypePointer[DType.f32]
    var rows: Int
    var cols: Int

    fn __init__(inout self, rows: Int, cols: Int):
        self.data = DTypePointer[DType.f32].alloc(rows * cols)
        rand(self.data, rows*cols)
        self.rows = rows
        self.cols = cols

    fn __del__(owned self):
        self.data.free()

    fn zero(inout self):
        memset_zero(self.data, self.rows * self.cols)

    @always_inline
    fn __getitem__(self, y: Int, x: Int) -> F32:
        return self.load[1](y, x)

    @always_inline
    fn load[nelts:Int](self, y: Int, x: Int) -> SIMD[DType.f32, nelts]:
        return self.data.simd_load[nelts](y * self.cols + x)

    @always_inline
    fn __setitem__(self, y: Int, x: Int, val: F32):
        return self.store[1](y, x, val)

    @always_inline
    fn store[nelts:Int](self, y: Int, x: Int, val: SIMD[DType.f32, nelts]):
        self.data.simd_store[nelts](y * self.cols + x, val)

上記のMatrix型は,Pythonの実装をコピー&ペーストして,型のアノテーションを追加しただけのものです.

# Note that C, A, and B have types.
fn matmul_naive(C: Matrix, A: Matrix, B: Matrix):
    for m in range(C.rows):
        for k in range(A.cols):
            for n in range(C.cols):
                C[m, n] += A[m, k] * B[k, n]

ではベンチマークしてみましょう。

@always_inline
def benchmark[func : fn(Matrix, Matrix, Matrix) -> None]
    (M : Int, N : Int, K : Int, python_gflops: F64):
    var C = Matrix(M, N)
    C.zero()
    var A = Matrix(M, K)
    var B = Matrix(K, N)

    @always_inline
    @parameter
    fn test_fn():
        _ = func(C, A, B)

    let secs = F64(Benchmark().run[test_fn]()) / 1_000_000_000
    # Prevent matrices from being destroyed before we finished benchmarking them.
    _ = A.data
    _ = B.data
    _ = C.data

    let gflops = ((2*M*N*K)/secs) / 1e9
    let speedup : F64 = gflops / python_gflops
    print(gflops, "GFLOP/s, a", speedup.value, "x speedup over Python")

128x128行列ではなく敢えて512x512行列にしてみました。

3.120776 GFLOP/s, a 1866.805513 x speedup over Python
約1866倍も高速化しました。

最内周ループのベクトル化
ベクター命令を利用することで、さらに最適化を進めてみます。ベクター幅を仮定するのではなく、dtype_simd_widthを使用して指定されたdtypeのsimd幅を問い合わせるのです。これにより、他のハードウェアに移行する際にも、コードの移植性が高まります。SIMD命令を活用するのは、次のように簡単です:

# Mojo has SIMD vector types, we can vectorize the Matmul code as follows.
alias nelts = dtype_simd_width[DType.f32]() # The SIMD vector width.
fn matmul_vectorized_0(C: Matrix, A: Matrix, B: Matrix):
    for m in range(C.rows):
        for k in range(A.cols):
            for nv in range(0, C.cols, nelts):
                C.store[nelts](m,nv, C.load[nelts](m,nv) + A[m,k] * B.load[nelts](k,nv))
        
            # Handle remaining elements with scalars.
            for n in range(nelts*(C.cols//nelts), C.cols):
                C[m,n] += A[m,k] * B[k,n]

上記の実装をベンチマークすることができます。多くのコンパイラは、ナイーブループを検出して最適化を実行できることに注意してください。しかし、Mojoでは、明示的に、どのような最適化が適用されるかを正確に制御することができます。

14.318266GFLOP/s、Python比で8564.990560倍も高速化

ベクトル化は一般的な最適化であり、Mojoはベクトル化を実行する高次の関数を提供しています。vectorize関数は、ベクトル幅と、ベクトル幅をパラメータとし、ベクトル化された方法で評価されることになる関数を受け取ります。

# Simplify the code by using the builtin vectorize function
from Functional import vectorize
fn matmul_vectorized_1(C: Matrix, A: Matrix, B: Matrix):
    for m in range(C.rows):
        for k in range(A.cols):
            @parameter
            fn dot[nelts : Int](n : Int):
                C.store[nelts](m,n, C.load[nelts](m,n) + A[m,k] * B.load[nelts](k,n))
            vectorize[nelts, dot](C.cols)

この両者とではほとんど性能差はありません。
13.978696 GFLOP/s、Python比で8361.864443倍

Matmulの並列化
最新のプロセッサから最高のパフォーマンスを得るには、それらが持つ複数のコアを利用する必要があります。Mojoでは、parallelize関数で簡単に実現することができます。matmulの実装を変更し、マルチスレッド化してみましょう(簡単のため、M次元でのみ並列化します):

# Parallelize the code by using the builtin parallelize function
from Functional import parallelize
fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix):
    @parameter
    fn calc_row(m: Int):
        for k in range(A.cols):
            @parameter
            fn dot[nelts : Int](n : Int):
                C.store[nelts](m,n, C.load[nelts](m,n) + A[m,k] * B.load[nelts](k,n))
            vectorize[nelts, dot](C.cols)
        
    parallelize[calc_row](C.rows)

24.091962 GFLOP/s、Python比で14411.481509倍高速化

タイリング Matmul
タイリングは、キャッシュの局所性を高めるためにmatmulで行われる最適化です。このアイデアは、サブマトリックスをキャッシュに常駐させ、再利用を増加させることです。タイル関数自体は、Mojoで次のように書くことができる:

from Functional import Static2DTileUnitFunc as Tile2DFunc
# Perform 2D tiling on the iteration space defined by end_x and end_y.
fn tile[tiled_fn: Tile2DFunc, tile_x: Int, tile_y: Int](end_x: Int, end_y: Int):
    # Note: this assumes that ends are multiples of the tiles.
    for y in range(0, end_y, tile_y):
        for x in range(0, end_x, tile_x):
            tiled_fn[tile_x, tile_y](x, y)

2次元の反復空間に対して2次元のタイリングを行います。
このように定義すると、matmulカーネル内で使用することができます。簡単のために、タイルの高さを4とし、ベクトル化もしたいので、タイルの幅を4 * neltsとします(列に対してベクトル化するので)。

# Use the above tile function to perform tiled matmul.
fn matmul_tiled_parallelized(C: Matrix, A: Matrix, B: Matrix):
    @parameter
    fn calc_row(m: Int):
        @parameter
        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):
            for k in range(y, y + tile_y):
                @parameter
                fn dot[nelts : Int,](n : Int):
                    C.store[nelts](m,n + x, C.load[nelts](m,n+x) + A[m,k] * B.load[nelts](k,n+x))
                vectorize[nelts, dot](tile_x)

        # We hardcode the tile factor to be 4.
        alias tile_size = 4
        tile[calc_tile, nelts * tile_size, tile_size](A.cols, C.cols)

    parallelize[calc_row](C.rows)

23.391380 GFLOP/s、Python比で13992.403260倍高速化。

上記の実装におけるオーバーヘッドの原因の一つは、dot関数のvectorizeによって導入されたループを展開していないからです。これはMojoのvectorize_unroll高次関数によって行うことができます:

# Unroll the vectorized loop by a constant factor.
from Functional import vectorize_unroll
fn matmul_tiled_unrolled_parallelized(C: Matrix, A: Matrix, B: Matrix):
    @parameter
    fn calc_row(m: Int):
        @parameter
        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):
            for k in range(y, y + tile_y):
                @parameter
                fn dot[nelts : Int,](n : Int):
                    C.store[nelts](m,n+x, C.load[nelts](m,n+x) + A[m,k] * B.load[nelts](k,n+x))

                # Vectorize by nelts and unroll by tile_x/nelts
                # Here unroll factor is 4
                vectorize_unroll[nelts, tile_x//nelts, dot](tile_x)

        alias tile_size = 4
        tile[calc_tile, nelts*tile_size, tile_size](A.cols, C.cols)
      
    parallelize[calc_row](C.rows)

24.263229 GFLOP/s、Python比で14513.931176倍の高速化。

ループ展開で若干高速化が進みました。

from Autotune import autotune, search
from Time import now
from Pointer import Pointer

alias matmul_fn_sig_type = fn(Matrix, Matrix, Matrix) -> None

タイルファクターの選択は行列乗算のパフォーマンスに大きく影響しますが、最適なタイルファクターはハードウェアに大きく依存し、キャッシュ構成やその他のモデル化しにくい効果に影響されます。私たちは、ハードウェアについてすべてを知ることなく、ポータブルなコードを書きたいので、Mojoに、オートチューニングを使用して自動的に最適なタイルファクターを選択するように依頼することができます。

# Autotune the tile size used in the matmul.
@adaptive
fn matmul_autotune_impl(C: Matrix, A: Matrix, B: Matrix):
    @parameter
    fn calc_row(m: Int):
        @parameter
        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):
            for k in range(y, y + tile_y):
                @parameter
                fn dot[nelts : Int,](n : Int):
                    C.store[nelts](m,n+x, C.load[nelts](m,n+x) + A[m,k] * B.load[nelts](k,n+x))
                vectorize_unroll[nelts, tile_x // nelts, dot](tile_x)

        # Instead of hardcoding to tile_size = 4, search for the fastest 
        # tile size by evaluting this function as tile size varies.
        alias tile_size = autotune(1, 2, 4, 8, 16, 32)
        tile[calc_tile, nelts * tile_size, tile_size](A.cols, C.cols)
      
    parallelize[calc_row](C.rows)

これにより、matmul関数の候補が複数生成されます。Mojoに最適なタイルファクターを見つける方法を教えるために、Mojoが各候補を評価するために使用できる評価関数を提供します。

fn matmul_evaluator(funcs: Pointer[matmul_fn_sig_type], size: Int) -> Int:
    print("matmul_evaluator, number of candidates: ", size)

    let eval_begin: Int = now()

    # This size is picked at random, in real code we could use a real size
    # distribution here.
    let M = 512
    let N = 512
    let K = 512
    print("Optimizing for size:", M, "x", N, "x", K)

    var best_idx: Int = -1
    var best_time: Int = -1

    alias eval_iterations = 10
    alias eval_samples = 10

    var C = Matrix(M, N)
    var A = Matrix(M, K)
    var B = Matrix(K, N)
    let Cptr = Pointer[Matrix].address_of(C).address
    let Aptr = Pointer[Matrix].address_of(A).address
    let Bptr = Pointer[Matrix].address_of(B).address

    # Find the function that's the fastest on the size we're optimizing for
    for f_idx in range(size):
        let func = funcs.load(f_idx)

        @always_inline
        @parameter
        fn wrapper():
            func(C, A, B)
        let cur_time = Benchmark(1, 100_000, 500_000_000, 1000_000_000).run[wrapper]()

        if best_idx < 0:
            best_idx = f_idx
            best_time = cur_time
        if best_time > cur_time:
            best_idx = f_idx
            best_time = cur_time

    let eval_end: Int = now()
    # Prevent matrices from being destroyed before we finished benchmarking them.
    _ = A.data
    _ = B.data
    _ = C.data
    print("Time spent in matmul_evaluator, ms:", (eval_end - eval_begin) // 1000000)
    print("Best candidate idx:", best_idx)
    return best_idx

最後に、単に最良の候補を呼び出すだけのエントリ関数を定義する必要があります。

fn matmul_autotune(C: Matrix, A: Matrix, B: Matrix):
    alias best_impl: matmul_fn_sig_type
    search[
        matmul_fn_sig_type,
        VariadicList(matmul_autotune_impl.__adaptive_set),
        matmul_evaluator -> best_impl
    ]()
    # Run the best candidate
    return best_impl(C, A, B)

結果的にbenchmarkを取ると

matmul_evaluator, number of candidates: 6 Optimizing for size: 512 x 512 x 512 Time spent in matmul_evaluator, ms: 8650 Best candidate idx: 2 23.488668 GFLOP/s, a 14050.599768 x speedup over Python

この測定は、実行したH/W条件に依存するのでベンチマーク結果は注目対象ではありません。重要なのは、nVidiaのCUDAからスマホのエッジまで、実行時に最適なH/W構成を検索して実行することを mojo がサポートしている点です。


強力な型チェック

Pythonのように柔軟な型を使用することができても、Mojoでは厳格な型チェックを使用することができます。型チェックを行うことで、コードの予測性、管理性、安全性を高めることができます。

強力な型チェックを採用する主な方法の 1 つに、Mojo の struct 型があります。Mojo の struct 定義では、コンパイル時決定の名前を定義し、型コンテキストでその名前を参照すると、定義されている値に対する強力な仕様として扱われます。たとえば、上記のMyPair構造体を使用する次のコードを考えてみましょう。異なる名前(intの代わりにInt)を使用することで、Pythonのintの動作を変更せずに、両方の型をMojoで維持することができます。

さらに、Intは、Mojoで作成する他のカスタムデータ型と同じ命名スタイルに従います。IntはMojoの標準ツールセット(+や-などの数学記号が独自のデータでどのように機能するか)に含まれている構造体です。すべての「標準型」(Int、Bool、String、Tupleなど)は、構造体を使用して作られています。これは、構造体が言語そのものに組み込まれているのではなく、使用できるツールの標準セットの一部であることを意味します。これにより、コードを書く際に、より柔軟でコントロールしやすくなります。

def pairTest() -> Bool:
    let p = MyPair(1, 2)
    return p < 4 # gives a compile-time error

このコードを実行すると、MyPair.__lt__のRHSが要求する「4」をMyPairに変換できない、というコンパイル時エラーが発生します。

これは、システム・プログラミング言語を扱う際にはお馴染みの経験ですが、Pythonの仕組みとは異なります。PythonにはMyPy型アノテーションと構文的に同じ機能がありますが、コンパイラによって強制されることはなく、代わりに静的解析に役立つヒントとなります。型を特定の宣言に結びつけることで、Mojoは互換性を壊すことなく、古典的な型注釈のヒントと強力な型指定の両方を扱うことができます。

強力な型の使用例は型チェックだけではありません。型が正確であることが分かっているので、その型に基づいてコードを最適化し、レジスタで値を渡し、引数渡しやその他の低レベルの詳細についてCと同じように効率化することができます。これは、Mojoがシステムプログラマーに提供する安全性と予測可能性の保証の基礎です。つまり、言語自体にハードワイヤされるのではなく、使用できるツールの標準セットの一部であるということです。これにより、コードを書く際に、より柔軟でコントロールしやすくなります。

C/C++ のような静的型付け言語を、そのまま言語仕様として導入するのではなく、今後、Mojoとして記述するコードに対して静的型付けの恩恵を得られる得られるという設計方針です。

オーバーロードされた関数とメソッド

Pythonのように、Mojoでは引数のデータ型を指定せずに関数を定義することができ、Mojoはそれらを動的に処理します。これは、任意の入力を受け入れ、動的なディスパッチにデータの処理方法を決定させるだけで動作する表現力豊かなAPIが欲しい場合に便利です。しかし、前述のように型安全性を確保したい場合、Mojoはオーバーロードされた関数やメソッドを完全にサポートしています。

これにより、同じ名前で異なる引数を持つ複数の関数を定義することができます。これは、C++、Java、Swiftなど、多くの言語で見られる一般的な機能です。

関数呼び出しを解決する場合、Mojoは各候補を試し、動作するものを使用するか(1つしか動作しない場合)、最も近いものを選ぶか(近いものを判断できる場合)、どれを選ぶべきか判断できない場合は呼び出しが曖昧であると報告します。後者の場合、呼び出し側で明示的なキャストを追加することで、曖昧さを解消することができます。

struct Array[T: AnyType]:
    fn __getitem__(self, idx: Int) -> T: ...
    fn __getitem__(self, idx: Range) -> ArraySlice: ...

構造体やクラス内のメソッドをオーバーロードしたり、モジュールレベルの関数をオーバーロードしたりできます。

Mojoは、結果型のみによるオーバーロードをサポートせず、型推論に結果型や文脈上の型情報を使用しないため、シンプルで高速、かつ予測可能な状態を維持できます。Mojoの型チェッカーは定義上シンプルで高速なので、「式が複雑すぎる」というエラーは決して発生しません。

型定義をせずに引数名を残すと、関数は動的な型を持つPythonと同じように振る舞われます。単一の引数型を定義するとすぐに、Mojoはオーバーロード候補を探し、上記のように関数呼び出しを解決します。

fn 定義

上記の拡張機能は、低レベルプログラミングを提供する礎であり、抽象化機能を提供しますが、多くのシステムプログラマは、Mojoのdefが提供するものよりも、より多くの制御と予測可能性を好みます。defは非常に動的で柔軟性があり、一般的にPythonと互換性があるように必然的に定義されています。引数は変更可能で、ローカル変数は最初の使用時に暗黙的に宣言され、スコープが強制されるわけではありません。これは、高レベルのプログラミングやスクリプトには最適ですが、システムプログラミングには必ずしも最適ではありません。これを補完するために、Mojoはdefの「厳密モード」のようなfn宣言を提供します。

代替案:fnのような新しいキーワードを使う代わりに、@strict defのような修飾子やデコレータを追加することができます。 しかし、いずれにしても新しいキーワードを取る必要があり、そのためのコストはほとんどかかりません。また、システムプログラミングの領域では、実際にはfnは常に使用されるので、ファーストクラスにすることはおそらく理にかなっています。

呼び出し側としては、fnとdefは互換性があります。defが提供できてfnが提供できないものはありません(逆もまた然り)。違いは、fnがより制限され、そのボディの内側で制御されることです(代わりに厳密です)。具体的には、fnはdef関数と比較して、いくつかの制限があります:

引数の値は、(varのように)変更可能ではなく、(letのように)関数本体で不変であることがデフォルトです。これにより、意図しない変異を防ぐことができ、コピー不可能な型を引数として使用することができるようになります。

引数の値には型指定が必要であり(メソッド内のselfを除く)、誤って型指定が漏れてしまうのを防ぐことができる。同様に、戻り値の型指定がない場合は、未知の戻り値の代わりにNoneを返すと解釈されます。なお、どちらも明示的にobjectを返すと宣言することができ、必要に応じてdefの動作にオプトインすることができる。

ローカル変数の暗黙的な宣言は無効で、すべてのローカル変数を宣言する必要があります。これにより、名前の誤記を防ぎ、letやvarが提供するスコープと連携することができます。

どちらも例外の発生をサポートしていますが、これはraisesキーワードを使ってfnで明示的に宣言する必要があります。

プログラミングのパターンはチームによって大きく異なるので、このレベルの厳しさは誰にでも当てはまるものではないでしょう。C++に慣れていて、PythonでMyPyスタイルの型注釈をすでに使っている人はfnsの使用を好みますが、より高度なプログラマやML研究者はdefを使い続けると予想しています。 Mojoでは、あるメソッドは一方、他のメソッドは他方で実装するなど、defとfn宣言を自由に混在させて、それぞれのチームやプログラマの使用ケースに最適なものを判断することができます。

copyinit__および__moveinit__特殊メソッド

Mojo は、C++ や Swift などの言語で見られる「値セマンティクス」を完全にサポートしており、@value デコレーター を使用して、フィールドの単純な集約を非常に簡単に定義することができます。

高度な使用例として、Mojoではカスタムコンストラクタ(Pythonの既存の__init__特殊メソッドを使用)、カスタムデストラクタ(既存の__del__特殊メソッドを使用)、新しい copyinit および moveinit 特殊メソッドを使用してカスタムコピーおよび移動コンストラクタを定義することができます。

これらの低レベルのカスタマイズフックは、低レベルのシステムプログラミングを行う際、例えば、手動でメモリ管理を行う場合に便利です。例えば、動的な文字列型があり、構築時に文字列データのメモリを確保し、値が破棄されるときにメモリを破棄する必要があるとします。

struct MyString:
    var data: Pointer[Int8]

    # StringRef has a data + length field
    def __init__(inout self, input: StringRef):
        let data = Pointer[Int8].alloc(input.length+1)
        data.memcpy(input.data, input.length)
        data[input.length] = 0
        self.data = Pointer[Int8](data)

    def __del__(owned self):
        self.data.free()

このMyString型は、この動作の簡単な例を示すために、低レベルの関数を使って実装されています。より現実的な実装では、短い文字列の最適化などを使用します。

fn useStrings():
    var a: MyString = "hello"
    print(a)   # Should print "hello"
    var b = a  # ERROR: MyString doesn't implement __copyinit__

    a = "Goodbye"
    print(b)   # Should print "hello"
    print(a)   # Should print "Goodbye"

Mojoコンパイラは、aからbへの文字列をコピーする方法を知らないので、コピーすることができません。MyString には Pointer (低レベルの C ポインタに相当) のインスタンスが含まれており、Mojo はそれがどのようなデータを指しているのか、どのようにコピーすればよいのかを知りません。より一般的には、いくつかの型(原子番号など)は、クラスのインスタンスと同様にアドレスがアイデンティティを提供するため、コピーしたり移動したりできません。

この場合、文字列はコピー可能であることが望まれます。これを可能にするために、__copyinit__という特殊メソッドを実装します:

struct MyString:
    ...
    def __copyinit__(inout self, existing: Self):
        self.data = Pointer(strdup(self.data.address))

この実装では、上記のコードは正しく動作し、b = aのコピーは、独自の寿命とデータを持つ文字列の論理的に異なるインスタンスを生成します。コピーは、上記のコード行で指示されているように、Cスタイルのstrdup()関数で作成されます。Mojoは__moveinit__メソッドもサポートしており、Rustスタイルのムーブ(ライフタイムの終了時に値を取得する)とC++スタイルのムーブ(値の内容は削除されるが、デストラクタは実行される)の両方を可能にし、カスタムムーブロジックを定義することができます。値のライフタイムに関する詳細は後述します。

Mojoは、型をコピー可能、移動のみ、移動不可にする機能を含む、値の寿命に関する完全な制御を提供します。これは、SwiftやRustのような言語が提供する、値が少なくとも移動可能であることを必要とするよりも多彩です。それ自体がコピーを作成することなく、__copyinit__メソッドに既存のものを渡すことができる方法に興味がある場合は、後述の借用引数に関するセクションを御覧ください。


引数渡し制御とメモリ所有権

PythonとMojoの両方において、言語の多くは関数呼び出しを中心に展開されます:多くの組み込みの動作は、標準ライブラリでメソッドで実装されています。これらの関数内部では、引数渡しによって多くのメモリ所有権が決定されます。

PythonとMojoが引数を渡す方法:

Pythonのdef関数に渡される値はすべて、参照セマンティクスを使用します。つまり、関数は渡された変更可能なオブジェクトを変更することができ、その変更は関数の外側で見ることができます。しかし、引数が指し示すオブジェクトを変更しても、その変更は関数の外では見えないという問題を抱えています。

Mojoのdef関数に渡されるすべての値は、デフォルトで値セマンティクスを使用します。Python と比較すると、これは重要な違いです: Mojoのdef関数はすべての引数のコピーを受け取り、関数内で引数を変更することができますが、変更は関数の外では見えません。

モジョのfn関数に渡されるすべての値は、デフォルトで不変の参照です。つまり、関数は元のオブジェクトを読むことはできますが(コピーではありません)、オブジェクトを一切変更することはできません。

Mojo fn で不変の引数を渡すためのこの規約は、"borrowing" と呼ばれます。以下のセクションでは、def関数とfn関数の両方について、Mojoで引数渡しの動作を変更する方法を説明します。

不変の引数(借用)

借用オブジェクトとは、オブジェクトのコピーを受け取る代わりに、関数が受け取るオブジェクトへの不変の参照のことです。したがって、呼び出し側の関数はオブジェクトに対する完全な読み取りと実行のアクセス権を持ちますが、それを変更することはできません(呼び出し側は依然としてオブジェクトの排他的な「所有権」を持っています)。

これはfn引数のデフォルトの動作ですが、必要に応じてbuiltedキーワードで明示的に定義することができます(def引数にもbuiltedを適用することができます):

fn use_something_big(borrowed a: SomethingBig, b: SomethingBig):
    """a "と "b "は同じように渡されます。"借用"がデフォルトだからです。"""。
    a.print_id()
    b.print_id()

このデフォルトは、メソッドの self 引数を含む、すべての引数に一律に適用されます。これは、大きな値を渡すときや、参照カウントされたポインタ(Python/Mojoクラスのデフォルトです)のような高価な値を渡すときに、引数を渡すときにコピーコンストラクタとデストラクタを呼び出す必要がないため、より効率的です。

上のコードをベースにした応用例:

# copyコストが高い大きい構造体なので __copyinit__ メソッドがありません
struct SomethingBig:
    var id_number: Int
    var huge: InlinedArray[Int, 100000]
    fn __init__(inout self): …

    fn set_id(inout self, number: Int):
        self.id_number = number

    fn print_id(self):  # Same as: fn print_id(borrowed self):
        print(self.id_number)

fn try_something_big():
    let big = SomethingBig()
    big.print_id()
    use_something_big(big, big)

fn 関数のデフォルトの引数規約を借用しているため、Mojo はデフォルトで正しいことを行うシンプルで論理的なコードになっています。たとえば、print_id()メソッドを呼び出すときや、use_something_big()を呼び出すときに、SomethingBigのすべてをコピーしたり移動したりしたくありません。

この借用引数規約は、C++のconst&による引数渡しと似ているところがあり、値のコピーを回避し、呼び出し側の変異性を無効にすることができます。しかし、この借用規約は、C++のconst&とは2つの重要な点で異なっています。

Mojoコンパイラは、(Rustに似た)借用チェッカーを実装し、不変の参照が存在する場合にコードが動的に値への可変参照を形成するのを防ぎます。また、同じ値への複数の可変参照を防止します。上記のuse_something_bigの呼び出しのように複数のborrowを持つことは許されていますが、ミュータブル参照で何かを渡すこととborrowを同時に行うことはできません。

Int、Float、SIMDのような小さな値は、余分なインダイレクトを通さず、レジスタで直接渡されます(これは、@register_passableデコレーターで宣言されているためです)。これは、C++やRustなどの言語と比較した場合、大幅な性能向上となり、この最適化をすべての呼び出しサイトから型に対する宣言的なものに移行します。

Rustと同様に、Mojoの借用チェッカーは、不変量の排他性を強制的に実現します。RustとMojoの大きな違いは、Mojoはborrowによる受け渡しに呼び出し側の紋章を必要としないことです。また、Mojoは小さな値を渡すときに効率的であり、Rustはborrowで渡す代わりに値を移動させることをデフォルトとしています。これらの方針と構文の決定により、Mojoはより使いやすいプログラミングモデルを提供することができます。

ミュータブル引数(inout)

一方、fn関数を定義する際に、引数をミュータブルにしたい場合(関数内部での引数の変更が関数外部でも見えるようにしたい場合)、inoutキーワードで引数をミュータブルと宣言しなければなりません。

次の例では、__iadd__関数がself.を変更しようとしています:

struct Int:
    # self and rhs are both immutable in __add__.
    fn __add__(self, rhs: Int) -> Int: ...

    # ... but this cannot work for __iadd__
    fn __iadd__(self, rhs: Int):
        self = self + rhs  # ERROR: cannot assign to self!

ここでの問題は、これはMojoのfn関数なので、selfは不変であり、引数の内部状態を変更することはできないことです。解決策としては、selfの引数名にinoutキーワードを追加することで、引数がミュータブルであることを宣言します。

struct Int:
    # ...
    fn __iadd__(inout self, rhs: Int):
        self = self + rhs    # OK

これで、self引数は関数内でミュータブルとなり、呼び出し側で変更が可視化されます。たとえ、呼び出し側が配列の添え字のような自明でない計算でアクセスする場合も可視化されます。

fn show_mutation():
    var x = 42
    x += 1
    print(x)    # prints 43 of course

    var a = InlinedFixedVector[16, Int](...)
    a[4] = 7
    a[4] += 1    # Mutate an element within the InlinedFixedVector
    print(a[4])  # Prints 8

    let y = x
    y += 1       # ERROR: Cannot mutate 'let' value

Mojoは、上記のInlinedFixedVector要素のインプレース変異を、一時バッファへの__getitem__の呼び出しと、呼び出し後の__setitem__によるストアを発することで実装しています。letの値の変異は、不変の値への変異可能な参照を形成することができないため、失敗します。同様に、コンパイラは、__getitem__は実装しているが__setitem__は実装していない場合、inout引数で添え字を使おうとする拒否されます。

もちろん、複数のinout引数を宣言することも可能です。例えば、次のようなswap関数を定義して使用することができます:

fn swap(inout lhs: Int, inout rhs: Int):
    let tmp = lhs
    lhs = rhs
    rhs = tmp

fn show_swap():
    var x = 42
    var y = 12
    swap(x, y)
    print(x)  # Prints 12
    print(y)  # Prints 42



引数を転送する(ownedと^)

Mojoがサポートする最後の引数規約は、owned引数規約です。この規約は、値を排他的に所有したい関数に使用され、多くの場合、固定後の ^ 演算子とともに使用されます。

例えば、ユニークポインタのような移動専用の型を扱う場合を想像してみてください。borrowの規約により、式なしでユニークポインタを扱うのは簡単ですが、ある時点で他の関数に所有権を移したい場合があります。これが ^ "転送" 演算子の役割です。

fn usePointer():
    let ptr = SomeUniquePtr(...)
    use(ptr)        # Perfectly fine to pass to borrowing function.
    take_ptr(ptr^)  # Pass ownership of the `ptr` value to another function.

    use(ptr) # ERROR: ptr is no longer valid here!

移動可能な型の場合、^演算子は値結合の寿命を終了させ、値の所有権を他のもの(この場合、take_ptr()関数)に移します。これをサポートするために、関数を所有する引数を取るものとして定義することができます。例えば、take_ptr()を次のように定義します:

fn take_ptr(owned p: SomeUniquePtr):
    use(p)

own宣言されているので、take_ptr()関数は、その値へのユニークなアクセス権を持っていることを知ることができます。これは、一意なポインタのようなものにとって非常に重要であり、コピー禁止させたい時に便利です。

例えば、デストラクタやムーブコンストラクタを使用する際に、ownの規約を目にすることが多いでしょう。例えば、先ほどのMyString型は次のように定義することができます:

struct MyString:
    var data: Pointer[Int8]

    # StringRef has a data + length field
    def __init__(inout self, input: StringRef): ...
    def __copyinit__(inout self, existing: Self): ...

    def __moveinit__(inout self, owned existing: Self):
        self.data = existing.data

    def __del__(owned self):
        self.data.free()

値を破棄するためには値を所有する必要があるため、__del__関数でownを指定することは重要です。

defとfnの引数渡しを比較する

Mojoのdef関数は、本質的にfn関数のシンタックシュガーにに過ぎません。明示的な型注釈のないdef引数は、デフォルトでObjectになります。(inoutやownedなどの)規約キーワードのないdef引数は、引数と同じ名前のmutable varに暗黙のコピーで渡されます。(これには、型に__copyinit__メソッドがあることが必要です。)

例えば、この2つの関数は同じ動作をします:

def example(inout a: Int, b: Int, c):
    # b and c use value semantics so they're mutable in the function
    ...

fn example(inout a: Int, b_in: Int, c_in: Object):
    # b_in and c_in are immutable references, so we make mutable shadow copies
    var b = b_in
    var c = c_in
    ...

Objectのような小さな型の参照はコピーするのが簡単なので、シャドウコピーは通常オーバーヘッドを追加しません。高価なのは参照カウントの調整ですが、これは移動の最適化で解消されます。

Pythonとの統合

MojoでおなじみのPythonモジュールを簡単に使用することができます。任意のPythonモジュールをMojoプログラムにインポートし、Mojo型からPython型を作成することができます。

Pythonモジュールのインポート

MojoでPythonモジュールをインポートするには、モジュール名を指定してPython.import_module()を呼び出すだけです:

from PythonInterface import Python

# This is equivalent to Python's `import numpy as np`
let np = Python.import_module("numpy")

# Now use numpy as if writing in Python
a = np.array([1, 2, 3])

Python NumPyをインポートして、他のPythonモジュールをインポートします。ただし、Pythonオブジェクトを操作しているため、print()のような一部のMojo関数が動作しないことを忘れないでください。MojoでPythonの型を表示したい場合は、Pythonの組み込みprint()関数を使用する必要があります:

a = np.array([1, 2, 3])
builtins = Python.import_module("builtins")
builtins.print(a)

現在のところ、個々のメンバー(単一のPythonクラスや関数など)をインポートすることはできず、Pythonモジュール全体をインポートして、モジュール名を通してメンバーにアクセスする必要があります。

ローカルPythonモジュールのインポート

Mojoで使用したいローカルPythonコードがある場合、そのディレクトリをPythonパスに追加してから、モジュールをインポートしてください。

例えば、次のような Python ファイルがあるとします:

import numpy as np

def my_algorithm(a, b):
    array_a = np.random.rand(a, a)
    return array_a + b

それをインポートしてMojoで使用するには

from PythonInterface import Python

Python.add_to_path("path/to/module")
let mypython = Python.import_module("mypython")
let builtins = Python.import_module("builtins")

let c = mypython.my_algorithm(2, 3)
builtins.print(c)

MojoでPythonを使用する際に、メモリ管理について心配する必要はありません。

PythonでのMojoの型

Mojoのプリミティブ型は暗黙のうちにPythonのオブジェクトに変換されます。現在、リスト、タプル、整数、浮動小数点数、ブーリアン、文字列をサポートしています。

例えば、Pythonの型を表示するPython関数があるとします:

def type_printer(my_list, my_tuple, my_int, my_string, my_float):
    print(type(my_list))
    print(type(my_tuple))
    print(type(my_int))
    print(type(my_string))
    print(type(my_float))

インポートしてMojo型に渡しても問題ないです:

from PythonInterface import Python

Python.add_to_path("/path/to/module")
let mypython2 = Python.import_module("mypython2")
mypython2.type_printer([0, 3], (False, True), 4, "orange", 3.4)

Python型に暗黙のうちに変換した後の型が出力されます。

Mojoはまだ標準的なDictionaryを持っていないので、Mojoの辞書からPythonの辞書を作成することはまだ不可能です。しかし、MojoでPythonの辞書を扱うことはできます。

パラメータ化:コンパイル時のメタプログラミング

Pythonの最も素晴らしい機能の1つは、ランタイムで拡張可能なメタプログラミング機能です。これにより、さまざまなライブラリが可能になり、柔軟で拡張性のあるプログラミングモデルが提供され、Pythonプログラマー全体が恩恵を受けています。ただし、これらの機能にはコストがかかります。ランタイムで評価されるため、基礎となるコードの実行効率に直接影響を与えます。また、これらの機能はIDEには知られていないため、コード補完などのIDE機能が理解し、開発者の体験を向上させるために使用することが困難です。

Pythonエコシステムの外では、静的メタプログラミングも開発の重要な部分であり、新しいプログラミングパラダイムや高度なライブラリの開発を可能にします。この領域にはさまざまなトレードオフを伴う事例が数多く存在します。例えば:

プリプロセッサ(Cプリプロセッサ、Lex/YACCなど)は、もっとも重い手法かもしれません。完全に汎用的ですが、開発者の体験やツールの統合においては最も劣っています。

LispやRustなどの一部の言語はマクロ展開機能をサポートしており、構文の拡張や冗長なコードの削減を可能にしています。ツールの統合においてはやや優れた機能を持っています。

C++などの一部の古い言語は、ランタイム言語と対応する非常に大規模かつ複雑なメタプログラミング言語(テンプレート)を持っています。これらは学習が難しく、コンパイル時間やエラーメッセージが劣っていることが特筆されます。

Swiftのような一部の言語は、一般性を犠牲にする代わりに、主要なケースに対して良好なエルゴノミクスを提供するために、コア言語に多くの機能を組み込んでいます。

Zigのような最新の言語では、言語インタプリタをコンパイルフローに組み込み、ASTをコンパイル中に反映させることができます。これにより、マクロシステムと同様の機能を提供し、拡張性と汎用性が向上します。

ModularのAI、高性能な機械学習カーネル、およびアクセラレータの開発において、高度なメタプログラミングシステムによる高い抽象化能力が必要です。高レベルでゼロコストの抽象化、表現力のあるライブラリ、複数のアルゴリズムバリアントの大規模な統合が求められます。Pythonと同様にライブラリ開発者がシステムを拡張可能です。

ただし、Mojoはコンパイル時間や分かり難いエラーメッセージなどを犠牲にすることは望んでいません。また、教育が困難な並行の言語エコシステムを構築することにも興味はありません。Mojoはこれまでのシステムから学びつつ、MLIRや細粒度な言語統合キャッシング技術などの新しい技術を利用してシステムを構築することができます。

そのため、Mojoコンパイラに組み込まれたコンパイル時メタプログラミングをサポートしています。このメタプログラミングは、パース、意味解析、IR生成の後、ターゲット固有のコードに変換される前の別のコンパイル段階として実行されます。実行時プログラムと同じホスト言語を使用し、これらのプログラムを予測可能な方法で表現および評価するためにMLIRを活用しています。

パラメータ化された型と関数を定義する

Mojoの構造体や関数はそれぞれパラメータ化されることがありますが、例を挙げると、なぜパラメータ化するのかがよくわかります。SIMD型は、スカラーデータ型の複数のインスタンスを保持するハードウェアの低レベルベクトルレジスタを表します。最近のハードウェアアクセラレータは512ビット以上のSIMDベクターを持つCPUを扱うことも珍しくなくなりました。ハードウェアには多様性がありますが(SSE、AVX-512、NEON、SVE、RVVなど)、多くの演算は数値演算やMLカーネル開発者がよく使うもので、SIMD型はそれらをMojoプログラマに公開します。

以下は、Mojoの標準ライブラリにあるSIMD APIの(カットされた)バージョンです:

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: ...

Mojo のパラメータは、PEP695 構文の拡張版を使用して角括弧内で宣言されます。パラメータは、Mojo プログラムの通常の値のように名前が付けられ、型を持ちますが、ターゲットプログラムによって実行時ではなく、コンパイル時に評価されます。ランタイムプログラムはパラメータの値を使用することができます。なぜなら、パラメータはランタイムプログラムが必要とする前にコンパイル時に解決されるからです。しかし、コンパイル時のパラメータ式はランタイムの値を使用することはできません。

SIMD構造体は、typeパラメータとsizeパラメータによってパラメータ化されています。キャストメソッドはさらにターゲットパラメータでパラメータ化されています。SIMDはパラメータ化された型なので、自己の引数の型がパラメータを担います(完全な型名はSIMD[type, size])。これを書き出すことは常に有効ですが(splat()の戻り値の型に示されているように)、これは冗長になりかねないので、__add__の例のように(PEP673の)Self型を使用することが推奨されます。


パラメータ化された型と関数の使用

SIMD型の場合、sizeはSIMDベクトルの要素数を指定し、typeは要素タイプを指定します。例えば、小さな浮動小数点ベクトルを表すために「4xFloat」、機械学習型「bfloat16」を持つAVX-512システムでは「32xbfloat16」を使用するとよいでしょう:

fn funWithSIMD():
    # Make a vector of 4 floats.
    let small_vec = SIMD[DType.f32, 4](1.0, 2.0, 3.0, 4.0)

    # Make a big vector containing 1.0 in bfloat16 format.
    let big_vec = SIMD[DType.bf16, 32].splat(1.0)

    # Do some math and convert the elements to float32.
    let bigger_vec = (big_vec+big_vec).cast[DType.f32]()

    # You can write types out explicitly if you want of course.
    let bigger_vec2 : SIMD[DType.f32, 32] = bigger_vec

cast()メソッドには、どの型にキャストするかを示す追加のパラメータが必要であることに注意してください。これは、cast()の呼び出しをパラメータ化することで処理されます。上の例では具体的な型を示しましたが、パラメータの大きな力は、パラメトリックなアルゴリズムや型を定義する能力から生まれます。例えば、長さやDTypeを問わないようなパラメトリックなアルゴリズムを定義することは非常に簡単です。

fn rsqrt[width: Int, dt: DType](x: SIMD[dt, width]) -> SIMD[dt, width]:
    return 1 / sqrt(x)

Mojoコンパイラは、パラメータを使った型推論は賢いです。この関数は、パラメータを指定せずにパラメトリックなsqrt()関数を呼び出すことができ、コンパイラは、sqrtwidth,typeと明示的に書いたようにそのパラメータを推測することに注意してください。また、rsqrt()は最初のパラメータをwidthという名前で定義することにしましたが、SIMD型は挑戦せずにsizeという名前にしていることに注意してください。

パラメータ式は単なるMojoコード

すべてのパラメータとパラメータ式は、ランタイムプログラムと同じ型システムを使って型付けされます: Mojoの標準ライブラリでは、IntとDTypeは構造体として実装されています。パラメータは非常に強力で、演算子による式の使用や、コンパイル時の関数呼び出しなど、ランタイムプログラムと同じようにサポートします。これにより、多くの「依存型」機能を使用することができます。例えば、2つのSIMDベクトルを連結するヘルパー関数を定義したい場合、次のようになります:

fn concat[ty: DType, len1: Int, len2: Int](
    lhs: SIMD[ty, len1], rhs: SIMD[ty, len2]) -> SIMD[ty, len1+len2]:
      ...

fn use_vectors(a: SIMD[DType.f32, 4], b: SIMD[DType.f16, 8]):
    let x = concat(a, a)  # Length = 8
    let y = concat(b, b)  # Length = 16

結果の長さは入力ベクトルの長さの合計であり、単純な + 演算でそれを表現できることに注意してください。より複雑な例として、標準ライブラリのSIMD.shuffle()メソッドを見てみましょう。これは2つの入力SIMD値、リストとしてのベクトルシャッフルマスクを受け取り、シャッフルマスクの長さと一致するSIMDを返します。

強力なコンパイルタイムプログラミング

単純な式も便利ですが、時には制御フローを伴う命令的なコンパイルタイムロジックを書きたいこともあります。例えば、Mojo Mathモジュールのisclose()関数は、整数には厳密な等式を使用しますが、浮動小数点には「近い」比較を使用します。コンパイル時に再帰を行うことも可能です。例えば、ベクトルの全要素を再帰的に合計してスカラーにする「ツリーリダクション」アルゴリズムの例を以下に示します:

struct SIMD[type: DType, size: Int]:
    ...
    fn reduce_add(self) -> SIMD[type, 1]:
        @parameter
        if size == 1:
            return self[0]
        elif size == 2:
            return self[0] + self[1]

        # Extract the top/bottom halves, add them, sum the elements.
        let lhs = self.slice[size // 2](0)
        let rhs = self.slice[size // 2](size // 2)
        return (lhs + rhs).reduce_add()

これは、コンパイル時に実行されるif文である@parameter ifの機能を利用したものです。これは、その条件が有効なパラメータ式であることを要求し、if文のライブブランチのみがプログラムにコンパイルされることを保証するものである。

Mojo型は単なるパラメータ式

型内でパラメータ式を使用する方法を紹介しましたが、PythonとMojoの両方において、型の注釈はそれ自体が任意の式になり得ます。Mojoの型には特別なメタタイプ型があり、型パラメトリックなアルゴリズムや関数を定義することができます。例えば、C++のstd::vectorクラスのようなアルゴリズムを次のように定義することができます:

struct DynamicVector[type: AnyType]:
    ...
    fn reserve(inout self, new_capacity: Int): ...
    fn push_back(inout self, value: type): ...
    fn pop_back(inout self): ...
    fn __getitem__(self, i: Int) -> type: ...
    fn __setitem__(inout self, i: Int, value: type): ...

fn use_vector():
    var v = DynamicVector[Int]()
    v.push_back(17)
    v.push_back(42)
    v[0] = 123
    print(v[1])      # Prints 42
    print(v[0])      # Prints 123

getitem__関数の値引数および戻り値の形式的な型として、型パラメータが使用されていることに注意してください。パラメータによって、DynamicVector型はさまざまなユースケースに基づいた異なるAPIを提供することができます。他にも、より高度なユースケースの恩恵を受けるケースはたくさんあります。例えば、並列処理ライブラリではparallelForEachNアルゴリズムが定義されており、クロージャをN回並列に実行し、コンテキストから値を供給します。その値はどのような型でも受け付けます:

fn parallelize[
    arg_type: AnyType,
    func: fn(Int, arg_type) -> None,
](rt: Runtime, num_work_items: Int, arg: arg_type):
    # Not actually parallel: see Functional.mojo for real impl.
    for i in range(num_work_items):
        func(i, arg)

これは、funcパラメータが先のarg_typeパラメータを参照することができ、その型が順次洗練されていくためです。

このことが重要なもう一つの例は、アルゴリズムやデータ構造が異種型のリストに対して定義される必要があるような、可変型ジェネリックの場合です:

struct Tuple[*ElementTys: AnyType]:
    var _storage : *ElementTys
struct Array[T: AnyType]:
    fn __getitem__[IndexType: AnyType](self, idx: IndexType)
       -> (ArraySlice[T] if issubclass(IndexType, Range) else T):
       ...

別名:名前付きパラメータ式

コンパイル時の値に名前を付けたいことはよくあることです。var は実行時の値を定義し、let は実行時の定数を定義しますが、コンパイル時の一時的な値を定義する方法が必要です。このために、Mojoはエイリアス宣言を使用します。たとえば、DType構造体は、次のように列挙子のエイリアスを使用して単純なenumを実装しています:

struct DType:
    var value : Int8
    alias invalid = DType(0)
    alias bool = DType(1)
    alias si8 = DType(2)
    alias ui8 = DType(3)
    alias si16 = DType(4)
    alias ui16 = DType(5)
    ...
    alias f32 = DType(15)

これにより、クライアントはDType.f32をパラメータ式として(ランタイム値としても機能する)自然に使用できるようになります。なお、これはDTypeのランタイムコンストラクタをコンパイル時に呼び出していることになります。

型はコンパイル時の表現であるため、このようなことができると便利です:

alias F32 = SIMD[DType.f32, 1]
alias UI8 = SIMD[DType.ui8, 1]

var x : F32   # F32 works like a "typedef"

varやletと同様、エイリアスはスコープに従うので、期待通りに関数内でローカルエイリアスを使用することができます。


オートチューニング/適応型コンパイル

Mojoのパラメータ式では、他の言語でできるような移植性の高いパラメータアルゴリズムを書くことができますが、高性能なコードを書くときには、パラメータに使用する具体的な値を選択する必要があります。例えば、高性能な数値計算アルゴリズムを書く場合、アルゴリズムを高速化するためにメモリタイリングを使用したいと思うかもしれませんが、使用する次元は、使用可能なハードウェア機能、キャッシュのサイズ、カーネルに融合するもの、その他多くの厄介な詳細によって大きく異なります。

一般的なマシンのベクトル長はデータ型に依存し、bfloat16のようにすべての実装で完全にサポートされていないデータ型もあるため、ベクトル長の管理も困難な場合があります。Mojoは、標準ライブラリにautotune関数を提供することでこれを支援します。例えば、データのバッファに対してベクトル長にとらわれないアルゴリズムを書きたい場合、次のように書くことができます:

from Autotune import autotune

def exp_buffer_impl[dt: DType](data: ArraySlice[dt]):
    # Pick vector length for this dtype and hardware
    alias vector_len = autotune(1, 4, 8, 16, 32)

    # Use it as the vectorization length
    vectorize[exp[dt, vector_len]](data)

このコードのインスタンス化をコンパイルする際、Mojoはこのアルゴリズムのコンパイルをフォークし、ターゲットハードウェアで実際に最もよく機能するものを測定して、どの値を使用するかを決定します。vector_len式のさまざまな値を評価し、ユーザー定義の性能評価器に従って最速のものを選びます。各オプションを個別に測定・評価するため、例えばF32とSI8では異なるベクトル長を選択することもあります。このシンプルな機能は、関数や型がパラメータ式でもあるため、単純な整数定数を超えて非常に強力です。

searchは評価者とフォークされた関数を受け取り、評価者が選んだ最速の実装をパラメータ結果として返します。

from Autotune import search

fn exp_buffer[dt: DType](data: ArraySlice[dt]):
    # Forward declare the result parameter.
    alias best_impl: fn(ArraySlice[dt]) -> None

    # Perform search!
    search[
      fn(ArraySlice[dt]) -> None,
      exp_buffer_impl[dt],
      exp_evaluator[dt] -> best_impl
    ]()

    # Call the selected implementation
    best_impl(data)

この例では、exp_evaluatorを性能評価関数として検索関数に与えています。性能評価関数は、候補となる関数のリストで呼び出され、最適な関数のインデックスを返す必要があります。Mojoの標準ライブラリには、関数の時間計測に使用できるBenchmarkモジュールが用意されています。

from Benchmark import Benchmark

fn exp_evaluator[dt: DType](
    fns: Pointer[fn(ArraySlice[dt]) -> None],
    num: Int
):
    var best_idx = -1
    var best_time = -1
    for i in range(num):
        candidate = fns[i]
        let buf = Buffer[dt]()

        # Benchmark this candidate.
        fn setup():
            buf.fill_random()
        fn wrapper():
            candidate(buf)
        let cur_time = Benchmark(2).run[wrapper, setup]()

        # Track the index of the fastest candidate.
        if best_idx < 0:
            best_idx = i
            best_time = cur_time
        elif best_time > cur_time:
            best_idx = f_idx
            best_time = cur_time

    # Return the fastest implementation.
    return best_idx

オートチューニングには指数関数的なランタイムがあります。これは、Mojoコンパイラスタックの内部実装の詳細(特にMLIR、統合キャッシュ、コンパイルの分散)から恩恵を得ています。これはパワーユーザー向けの機能であり、時間をかけて継続的に開発し、反復する必要があります。

"価値のライフサイクル": バリューの誕生、生、死

関数と型システムを構築するためのさまざまな材料について理解できたので、Mojoで表現したい重要な型をモデル化するために、それらをどのように組み合わせればよいかを見てみましょう。

既存の多くの言語は、さまざまなトレードオフで設計点を表現しています。 例えば、C++は非常に強力ですが、しばしば「デフォルトを間違える」と非難され、バグや誤った機能を引き起こすことがあります。Swiftは扱いやすいですが、値のコピーが多く、パフォーマンスについては「ARCオプティマイザ」に依存するという、あまり予測できないモデルになっています。Rustは、borrowチェッカーを満足させるために強力な値所有の目標から始まりましたが、値が移動可能であることに依存しているため、カスタム移動コンストラクタを表現するのが難しく、memcpyパフォーマンスに多くのストレスを与える可能性があります。Pythonでは、すべてがクラスへの参照であるため、型に関する問題に直面することはありません。

Mojoでは、このような既存のシステムから学ぶことで、非常に強力でありながら、学習や理解が容易なモデルを提供することを目標としています。また、「十分に賢い」コンパイラに組み込まれた「最善の努力」と予測困難な最適化パスを要求することは避けたいと考えています。

これらの問題を解決するために、さまざまな値の分類と、それを表現するための関連するMojoの機能を調べ、ボトムアップで構築します。C++は広く知られているため、例ではC++を主な比較対象としていますが、他の言語がより良い比較対象である場合は、他の言語を参照することもあります。

インスタンス化できない型

Mojoで最も簡素な型は、インスタンスを作成できない型です。この型にはイニシャライザがまったくなく、デストラクタがあっても、それが呼び出されることはありません(デストラクトするインスタンスが存在しないため):

struct NoInstances:
    var state: Int  # Pretty useless

    alias my_int = Int

    @staticmethod
    fn print_hello():
        print("hello world")

Mojo型はデフォルトでデフォルトコンストラクタ、ムーブコンストラクタ、メンバーワイズイニシャライザなどを取得しないので、このNoInstances型のインスタンスを作成することは不可能です。これらを取得するためには、__init__メソッドを定義するか、イニシャライザを合成するデコレータを使用する必要があります。このように、NoInstances.my_intやNoInstances.print_hello()のように、インスタンスを作成できなくても静的メンバを参照できるため、「名前空間」として有用な型です。

移動不可能な型とコピー不可能な型

インスタンス化は可能だが、いったんメモリ上のアドレスに固定されると、暗黙のうちに移動したりコピーしたりすることができない型にたどり着きます。これは、アトミック演算(C++のstd::atomicなど)や、値のメモリアドレスがそのアイデンティティであり、目的にとって重要であるような型を実装する場合に便利です:

struct Atomic:
    var state: Int

    fn __init__(inout self, state: Int = 0):
        self.state = state

    fn __iadd__(inout self, rhs: Int):
        #...atomic magic...

    fn get_value(self) -> Int:
        return atomic_load_int(self.state)

このクラスはイニシャライザを定義していますが、コピーや移動のコンストラクタはないので、一度初期化されると、決して移動したりコピーしたりすることはできません。これは、Mojoの所有権システムが完全に「アドレスの正しい」ものであるため、安全で有用です。これがスタック上または他の型のフィールドで初期化されると、決して移動する必要はありません。

Mojoのアプローチは、a = bコピーや^転送演算子などの組み込みの移動操作のみを制御することに注意してください。上記のAtomicのような独自の型に使える便利なパターンの1つは、明示的なcopy()メソッド(を追加することです。これは、プログラマにとって安全であることが分かっている場合に、インスタンスの明示的なコピーを作成するのに便利です。

ユニークな「移動専用」型

C++では、std::unique_ptrのような型や、POSIXファイルディスクリプタを所有するFileDescriptor型など、「ユニーク」な型に遭遇することがあります。これらの型は、Rustのような、コピーは推奨されないが「移動」は自由な言語において広く普及している。Mojoでは、__moveinit__メソッドを定義して固有の型の所有権を取得することで、この種の移動を実装することができます。例えば、以下のような感じです:

# This is a simple wrapper around POSIX-style fcntl.h functions.
struct FileDescriptor:
    var fd: Int

    # This is how we move our unique type.
    fn __moveinit__(inout self, owned existing: Self):
        self.fd = existing.fd

    # This takes ownership of a POSIX file descriptor.
    fn __init__(inout self, fd: Int):
        self.fd = fd

    fn __init__(inout self, path: String):
        # Error handling omitted, call the open(2) syscall.
        self = FileDescriptor(open(path, ...))

    fn __del__(owned self):
        close(self.fd)   # pseudo code, call close(2)

    fn dup(self) -> Self:
        # Invoke the dup(2) system call.
        return Self(dup(self.fd))
    fn read(...): ...
    fn write(...): ...

消費型移動コンストラクタ (moveinit) は、既存の FileDescriptor の所有権を取得し、その内部実装の詳細を新しいインスタンスに移動させます。これは、FileDescriptor のインスタンスが異なる場所に存在する可能性があり、ある値のボディを盗んで別の値に移動させるという論理的な移動が可能だからです。

以下は、__moveinit__を複数回呼び出すようなひどい例です:

fn egregious_moves(owned fd1: FileDescriptor):
    # fd1 and fd2 have different addresses in memory, but the
    # transfer operator moves unique ownership from fd1 to fd2.
    let fd2 = fd1^

    # Do it again, a use of fd2 after this point will produce an error.
    let fd3 = fd2^

    # We can do this all day...
    let fd4 = fd3^
    fd4.read(...)
    # fd4.__del__() runs here

値を所有する様々な値の間で、前のバインディングを破壊して新しい定数に所有権を移す、postfix-^「transfer」演算子を用いて、値の所有権がどのように移動しているかに注目してください。C++に慣れていると、転送演算子について考える簡単な方法はstd::moveですが、破壊できる状態にリセットすることなく移動できることがわかります。C++では、移動演算子が古い値のfdインスタンスを変更できなかった場合、2回クローズされてしまうのでした。

Mojoは値の有効性を追跡し、カスタムムーブコンストラクタを定義することができます。必要となるケースは稀ですが、必要な場合は非常に強力です。例えば、llvm::SmallVector型のように「インラインストレージ」最適化技術を使用する型もあり、インスタンスへの「内部ポインタ」で実装したい場合があります。これは、mallocメモリアロケータへの負担を減らすためのよく知られたトリックですが、「移動」操作が発生したときにポインタを更新するためのカスタムロジックが必要であることを意味します。

Mojoでは、カスタムの__moveinit__メソッドを実装するだけで、このようなことが簡単にできます。これは、C++でも簡単に実装できますが(ただし、カスタムロジックが不要な場合は定型文が必要です)、他の一般的なメモリセーフ言語では実装が難しいものです。

さらに、Mojoコンパイラは、優れた予測可能性と制御を提供する一方で、非常に洗練されたものでもあるということです。それは、一時的なものとそれに対応するコピー/移動操作を排除する権利を保留します。これが型に不適切な場合、ダンダーメソッドの代わりにcopy()のような明示的なメソッドを使用する必要があります。

"盗む移動 "をサポートする型

メモリセーフ言語の課題の1つは、コンパイラが追跡できる範囲内で予測可能なプログラミングモデルを提供する必要があり、コンパイラでの静的解析は本質的に限界があることです。例えば、以下の最初の例の2つの配列アクセスが異なる配列要素へのアクセスであることをコンパイラが理解することは可能ですが、2番目の例について推論することは(一般的に)不可能です:

std::pair<T, T> getValues1(MutableArray<T> &array) {
    return { std::move(array[0]), std::move(array[1]) };
}
std::pair<T, T> getValues2(MutableArray<T> &array, size_t i, size_t j) {
    return { std::move(array[i]), std::move(array[j]) };
}

この問題は、iとjの動的な値が同じでないことを知る、または証明する方法が(上記の関数本体だけを見ても)ないことです。配列の個々の要素が生きているかどうかを追跡するために動的な状態を維持することは可能ですが、これはしばしば(移動/転送が使用されていない場合でも)重大な実行時費用を引き起こし、Mojoや他のシステムプログラミング言語がやりたがらないことでもあります。この問題に対処する方法はさまざまで、かなり複雑な解決策もありますが、習得するのは必ずしも容易ではありません。

Mojoは、Mojoプログラマがその型システムを回避することなく仕事をこなせるように、実用的なアプローチをとっています。型がコピー可能、移動可能、あるいは構築可能であることを強制しませんが、型がその完全な契約を表現することを望み、プログラマーがC++などの言語から期待する流暢なデザインパターンを可能にしたいと思います。これは、「NULL状態」(オプションの型やNULL可能なポインタなど)を持っていたり、NULL値を持つことで作成が効率的で破壊が不要なためです(例えば、std::vectorはそのデータに対してNULLポインタを持つことができます)。

これらの使用例をサポートするために、^転送演算子は任意のLValueをサポートし、1つに適用されると、"stealing move constructor "を呼び出します。このコンストラクタは、新しい値がライブ状態になるように設定しなければならず、古い値を変異させることができるが、古い値をそのデストラクタがまだ動作する状態にしなければならない。例えば、FileDescriptorをベクターに入れ、そこから移動したい場合、-1がセンチネルであること、つまり「NULL」であることを知るように拡張することを選択するかもしれない。このような実装が可能です:

# This is a simple wrapper around POSIX-style fcntl.h functions.
struct FileDescriptor:
    var fd: Int

    # This is the new key capability.
    fn __moveinit__(inout self, inout existing: Self):
        self.fd = existing.fd
        existing.fd = -1  # neutralize 'existing'.

    fn __moveinit__(inout self, owned existing: Self): # as above
    fn __init__(inout self, fd: Int): # as above
    fn __init__(inout self, path: String): # as above

    fn __del__(owning self):
        if self.fd != -1:
            close(self.fd)   # pseudo code, call close(2)

「盗む移動」コンストラクタが、既存の値からファイルディスクリプタを取り出し、その値を変異させて、デストラクタが何もしないようにしていることに注目してください。このテクニックにはトレードオフがあり、すべての型に最適というわけではありません。センチネルケースをチェックする必要があるため、デストラクタに1つ(安価な)分岐が追加されていることがわかります。また、一般に、このような型をNULL化可能にするのは悪い形式と考えられています。

さらに、Mojo自体にOptional[T]を実装する予定であり、Optionalはこの機能を必要としています。また、Mojoではライブラリ作者は言語設計者よりも自分のドメインの問題を理解していると考えており、一般的に、ライブラリ作者にそのドメインに対する完全な権限を与えることを好みます。そのため、自分の型をオプトイン方式でこの動作に参加させることを選択することができます。

コピー可能なタイプ

可動型の次の段階として、コピー可能な型があります。プログラマは一般に文字列や配列がコピー可能であることを期待していますし、Pythonのすべてのオブジェクト参照はポインタをコピーして参照カウントを調整することでコピー可能です。

コピー可能な型を実装する方法はたくさんあります。PythonやJavaのように、共有ポインタを伝播させる参照セマンティック型を実装することもできますし、一度作成したら決して変異しないので簡単に共有できる不変データ構造を使うこともできますし、Swiftのように遅延コピーオンライトで深い値セマンティクスを実装することもできます。これらのアプローチにはそれぞれ異なるトレードオフがあり、Mojoは、コレクションタイプのいくつかの共通セットを望む一方で、特定のユースケースに焦点を当てた幅広い特殊なものをサポートすることもできるという意見を持っています。

Mojoでは、__copyinit__メソッドを実装することによってこれを行うことができます。ここでは、単純な String を使用したその例を疑似コードで示します:

struct MyString:
    var data: Pointer[Int8]

    # StringRef is a pointer + length and works with StringLiteral.
    def __init__(inout self, input: StringRef):
        self.data = ...

    # Copy the string by deep copying the underlying malloc'd data.
    def __copyinit__(inout self, existing: Self):
        self.data = strdup(existing.data)

    # This isn't required, but optimizes unneeded copies.
    def __moveinit__(inout self, owned existing: Self):
        self.data = existing.data

    def __del__(owned self):
        free(self.data.address)

    def __add__(self, rhs: MyString) -> MyString: ...

この単純な型は、mallocで割り当てられた「NULL終端」の文字列データへのポインターで、わかりやすくするために昔ながらのCのAPIを使用しています。これは copyinit を実装しており、MyString の各インスタンスがその基礎となるポインタを所有し、破壊時にそれを解放するという不変性を維持する。この実装では、上で見たトリックをベースに、__moveinit__コンストラクタを実装し、いくつかの一般的なケースで一時的なコピーを完全に排除することができます。この動作は、次のコードで見ることができます:

fn test_my_string():
    var s1 = MyString("hello ")

    var s2 = s1    # s2.__copyinit__(s1) runs here

    print(s1)

    var s3 = s1^   # s3.__moveinit__(s1) runs here

    print(s2)
    # s2.__del__() runs here
    print(s3)
    # s3.__del__() runs here

この場合、コピーコンストラクタが必要な理由が両方わかります。コピーコンストラクタがないと、s1の値をs2に複製することがエラーになります。同じコピー不可能な型のインスタンスを2つ持つことはできないためです。移動コンストラクタはオプションですが、s3への代入に役立ちます。これがないと、コンパイラはs1からコピーコンストラクタを呼び出し、古いs1インスタンスを破棄してしまいます。これは論理的には正しいのですが、実行時に余分なオーバーヘッドが発生します。

Mojoは値を破棄するので、コピーと破棄のペアを単一の移動操作に変換することができ、std::moveの広範なマイクロマネジメントを必要とせずにC++よりはるかに優れた性能を実現できます。


些細な型

最も柔軟な型は、単なる「ビットの袋」である。このような型は、カスタムコードを呼び出すことなく、コピー、移動、破棄が可能なため、些細です。このような型は、私たちを取り巻く最も一般的な基本型であると言っても過言ではありません。整数や浮動小数点値などは、すべて些細なものです。言語の観点からは、Mojoはこれらを特別にサポートする必要はありません。型作成者がこれらをノーオペレーションとして実装し、インライナーがそれらを消去するようにしても全く問題ありません。

1つは、些細な型にたくさんのメソッドを定義しなければならないという定型文が必要ないこと、もう1つは、たくさんの関数呼び出しを生成して周りに押しやり、インライン化されて何もなくなるというコンパイル時のオーバーヘッドが必要ないことです。さらに、関連する懸念事項として、これらの型の多くは別の意味で些細なものであるということがあります。

そのため、Mojoはこれらの問題をすべて解決する構造体デコレータを提供しています。このデコレータは、Mojoにその型がコピー可能で移動可能であることを伝えますが、そのためのユーザー定義のロジックはありません。また、MojoにCPUレジスタで値を渡すことを好むように指示し、効率的な利点につながることがあります。

値デコレータ(@value decorator)

Mojoのアプローチ(上述)は、Atomicのような低レベルのものを正しく表現する能力を与えるシンプルで予測可能なフックを提供します。これは制御やシンプルなプログラミングモデルには最適ですが、ほとんどの構造体は他の型の単純な集合体であり、そのために多くのボイラープレートを書かなければならないのは避けたいところです。これを解決するために、Mojoは構造体の@valueデコレータを提供し、ボイラープレートを合成してくれるようにしました。valueは、Mojoの新しいメソッドである__moveinit__と__copyinit__を扱うPythonの@dataclassの拡張版と考えることができます。

valueデコレーターは、型のフィールドを調べ、欠けているメンバーを生成します。たとえば、次のような単純な構造体を考えてみましょう:

@value
struct MyPet:
    var name: String
    var age: Int

Mojoは、あなたがメンバー単位のイニシャライザ、ムーブコンストラクタ、コピーコンストラクタを持っていないことに気づき、あなたが書いたかのようにこれらを合成してくれるでしょう:

fn __init__(inout self, owned name: String, age: Int):
    self.name = name^
    self.age = age

fn __copyinit__(inout self, existing: Self):
    self.name = existing.name
    self.age = existing.age

fn __moveinit__(inout self, owned existing: Self):
    self.name = existing.name^
    self.age = existing.age

型に移動専用フィールドが含まれている場合、もちろんコピーコンストラクタを生成することはできません(したがって、生成することもできません)。Mojoは、これらのフィールドが存在しない場合にのみ合成するので、独自のバージョンを定義することによって、Mojoの動作をオーバーライドすることができます。たとえば、カスタムのコピーコンストラクタを定義したいが、デフォルトのメンバーワイズコンストラクタとムーブコンストラクタを使用したいというのはよくあることです。

現時点では、特定のメソッドの生成を抑制したり、生成をカスタマイズしたりする方法はありませんが、需要があれば、@valueジェネレータに引数を追加してこれを行うことができます。

なお、@valueデコレータは、メンバがコピー可能または移動可能な型に対してのみ機能する。構造体の中にAtomicのようなものがある場合、それはおそらく値型ではなく、これらのメンバはいずれにしても必要ありません。

デストラクタの動作

Mojoの構造体にはデストラクタを設定することができます。たとえば、単純な文字列は次のようになります (擬似コード):

struct MyString:
    var data: Pointer[Int8]

    def __init__(inout self, input: StringRef): ...
    def __add__(self, rhs: MyString) -> MyString: ...
    def __del__(owned self):
        free(self.data.address)

Mojoコンパイラは、値が死んだときに自動的にデストラクタを呼び出し、デストラクタが実行されるタイミングについて強力な保証を提供します。Mojoは静的コンパイラ解析を使用して、コードを推論し、デストラクタの呼び出しを挿入するタイミングを決定します。例えば、次のようなものです:

fn use_strings():
    var a = MyString("hello a")
    var b = MyString("hello b")
    print(a)
    # a.__del__() runs here


    print(b)
    # b.__del__() runs here

    a = MyString("temporary a")
    # a.__del__() runs here

    other_stuff()

    a = MyString("final a")
    print(a)
    # a.__del__() runs here

上のコードでは、a と b の値が早い段階で作成され、値の初期化ごとにデストラクタへの呼び出しが行われていることがわかります。また、呼び出しが行われている場所が、b変数の中であることにも注目してください。例えば、Mojoはa変数の(無関係な)コピーからb変数のコピーまで、値を維持し、その呼び出しの直後にそれを破棄します。a 値は、最初のコピーの直後と、新しい (未使用の) 一時的な値を再割り当てした直後、および最後のコピーの直後に破棄されます。

Mojoは、「As Soon As Possible」(ASAP)ポリシーを使用して値を破棄し、呼び出しごとに実行されるアクティブなガベージコレクタのように動作します。内部式(a+b+c+dなど)を使用するコードでは、中間式が不要になると積極的に破棄されます(破棄はC++のように文の終わりまで遅延されることはありません)。Mojoは、ループ、if、try/exceptを含む制御フローを完全に理解します。

さて、これはC++プログラマにとっては驚くべきことかもしれません。C++プログラマが広く使用しているRAIIパターンの使用を無効にしているのです。では、なぜMojoはC++スタイルのスコープ付き破壊を使わず、こんなに積極的に破棄するのでしょうか?

Mojoの設計には、C++のモデルよりも強力な利点がいくつもあります:

Pythonは関数全体を超えるスコープを持たないので、MojoはPythonスタイルのdefが存在しても正しく動作する実行可能なモデルを提供する必要があります。
Pythonはオブジェクトの破壊について強力な保証を提供しないため、RAIIパターンを奨励しません。RAIIパターンを解決するために、Mojo(とPython)はリソースへのスコープ付きアクセスを提供するwith文を提供し、これはRAIIよりも意図的で構文的に明確になっています。
Mojoのアプローチでは、C++のoperator=(const T&)やoperator=(T&&)のように再割当演算子を実装する型が不要になり、型の定義が容易になり、概念もなくなります。
Mojoは、ミュータブル参照が他のミュータブル参照やイミュータブルボローと重なることを許しません。予測可能なプログラミングモデルを提供する主要な方法の1つは、オブジェクトへの参照ができるだけ早く死ぬようにすることです。コピーと削除のペアを「移動」操作に変換する「移動」最適化は、NRVOなどのC++移動最適化手段です。
C++では、スコープ終了時に値を破棄することは、末尾再帰のような一般的なパターンでは問題があります。なぜなら、デストラクタの呼び出しが末尾呼び出しの後に起こるからです。これは、特定の関数型プログラミングのパターンでは、パフォーマンスとメモリの重大な問題になることがあります。
Mojoのアプローチは、RustやSwiftの仕組みに近いもので、どちらも強力な値所有権追跡を持ち、メモリ安全性を提供しているからです。1つの違いは、彼らの実装は、動的な「ドロップフラグ」の使用を必要とすることです。彼らは、安全性を提供するために、値の状態を追跡するために隠されたシャドウ変数を維持します。これらはしばしば最適化されて取り除かれますが、Mojoのアプローチではこのオーバーヘッドが完全に排除されるため、生成されるコードが高速になり、あいまいさが回避されます。

フィールドセンシティブなライフタイム管理

Mojoの寿命分析は、完全に制御フローを意識していることに加え、完全にフィールドを意識しています(構造体の各フィールドは独立して追跡されます)。オブジェクト全体」がイニシャライザで初期化されたか、オブジェクト全体のデストラクタで破壊されたかを個別に追跡します。例えば、次のようなコードを考えてみましょう:

struct TwoStrings:
    var str1: MyString
    var str2: MyString
    fn __init__(inout self): ...
    fn __del__(owned self): ...

fn use_two_strings():
    var ts = TwoStrings()
    # ts.str1.__del__() runs here

    other_stuff()

    ts.str1 = MyString("hello a")     # Overwrite ts.str1
    print(ts.str1)
    # ts.__del__() runs here

ts.str1フィールドは設定された後、すぐに破棄されることに注意してください。これは、例えば転送演算子を使用するときにも見ることができます:

fn consume_and_use_two_strings():
    var ts = TwoStrings()
    consume(ts.str1^)

    # ts is partially initialized here!
    other_stuff()

    ts.str1 = MyString()  # All together now
    use(ts)               # This is ok
    # ts.__del__() runs here

other_stuff()の間、str1フィールドは完全に初期化されていません。これは、所有権がconsume()に移ったためです。もしそうでなければ、Mojoはuninitialized field errorでこのコードを拒否してしまうでしょう。

フィールドは一時的に転送できますが、「オブジェクト全体」は集約型のイニシャライザで構築し、集約型のデストラクタで破棄する必要があります。つまり、フィールドを初期化してオブジェクトを作成することも、フィールドを破棄してオブジェクトを取り壊すこともできないのです:

fn consume_and_use_two_strings():
    var ts = TwoStrings()
    consume(ts.str1^)
    consume(ts.str2^)
    # Error: cannot run the 'ts' destructor without initialized fields.

    var ts2 : TwoStrings
    ts2.str1 = MyString()  # All together now
    ts2.str2 = MyString()  # All together now
    use(ts2) # Error: 'ts2' isn't fully initialized

このようなパターンを許容することもできますが、「価値はその部分の総和以上のものである」ことから、これを否定します。POSIXファイルディスクリプタを整数値として含むFileDescriptorを考えてみましょう。例えば、整数を破棄するのと、FileDescriptorを破棄する(close()システムコールを呼び出すかもしれない)のとでは、大きな違いがあるのです。このため、フルバリューの初期化はすべてイニシャライザを経由して、フルバリューのデストラクタで破棄することを要求しています。

init__におけるフィールドのライフタイム

オブジェクトのフィールドが未初期化であることは知っていますが、オブジェクト全体は初期化されていることを前提とします。つまり、すべてのフィールドが初期化されれば、すぐに「self」をオブジェクト全体として使うことができるのです:

struct TwoStrings:
    var str1: MyString
    var str2: MyString

    fn __init__(inout self, cond: Bool, other: MyString):
        self.str1 = MyString()
        if cond:
            self.str2 = other
            use(self)  # Safe to use immediately!
            # self.str2.__del__(): destroyed because overwritten below.

        self.str2 = self.str1
        use(self)  # Safe to use immediately!

同様に、Mojoのイニシャライザが、他のイニシャライザに委譲するなどして、selfを完全に上書きすることは全く問題ないです:

struct TwoStrings:
    var str1: MyString
    var str2: MyString

    fn __init__(inout self): ...
    fn __init__(inout self, cond: Bool, other: MyString):
        self = TwoStrings()  # basic
        self.str1 = MyString("fancy")

del__と__moveinit__における所有引数のフィールド寿命

デストラクタとムーブイニシャライザの「所有」引数には、最後にちょっとした工夫があります。これらのメソッドは次のように定義されています:

struct TwoStrings:
    var str1: MyString
    var str2: MyString
    fn __init__(...)

    fn __moveinit__(inout self, owned existing: Self): ...
    fn __del__(owned self): ...

これらのメソッドは、興味深い、しかし不明瞭な問題に直面しています。これらのメソッドはいずれも、所有する既存/self値を解体することを担当し、それらと関係のあるサブ要素を破壊するか、それらを使って自分自身の型の削除ロジックを実装することになります。移動コンストラクタは、既存のインスタンスからメンバを盗んで、新しいselfインスタンスを作りたいのです。そのため、どちらも所有する値の要素を所有したり変換したりしたいし、所有する値のデストラクタを実行させたくありません。最悪、__del__メソッドで、無限ループに陥ります。

この問題を解決するために、Mojoはこの2つのメソッドを特別に扱い、メソッドから何らかのリターンが返ってきた時点でその値全体が破棄されると仮定しています。つまり、フィールドの値が転送される前に、オブジェクト全体が使用される可能性があります。例えば、以下のようにすることで期待通りに動作します:

struct TwoStrings:
    var str1: MyString
    var str2: MyString
    fn __init__(...)
    fn __moveinit__(inout self, owned existing: Self): ...

    fn __del__(owned self):
        log(self)       # Self is still whole
        # self.str2.__del__(): Mojo destroys str2 since it isn't used

        consume(str1^)
        # Everything has now been transferred, no destructor is run on self.

しかし、メンバーへの内部ポインタを使ったロジックがある場合、デストラクタやムーブイニシャライザ自体のロジックのために、それらを生かしておく必要がある場合があります。このような場合は、discardパターンに代入してください:

fn __del__(owned self):
    log(self) # Self is still whole

    consume(str1^)
    _ = self.str2
    # self.str2.__del__(): Mojo destroys str2 after its last use.

この場合、consume()が暗黙のうちにstr2の何らかの値を参照していれば、_パターンでアクセスしたときに、最後の使用までstr2が破壊されないようにすることができます。

型の特徴
これは、RustのtraitやSwiftのプロトコル、Haskellのタイプクラスに非常によく似た機能です。

アドバンス/オブスキュアMojo機能
このセクションでは、標準ライブラリの最下層を構築する開発者向けに重要なパワーユーザー機能について説明します。

register_passable構造体デコレータ
値を扱う際のデフォルトのモデルは、値がメモリ上に存在するため、IDを持ち、関数との間で間接的に渡されることを意味します(同等に、マシンレベルでは「参照渡し」されます)。これは、移動できない型には最適で、大きなオブジェクトや高価なコピー操作を必要とするものには安全なデフォルトとなります。一方で、単一の整数や浮動小数点数のような小さなものには非効率的です。

これを解決するために、Mojoでは@register_passableデコレータを使用して、構造体がメモリを通過する代わりにレジスタで渡されることを選択できるようにしています。このデコレータは、標準ライブラリのIntのような型に使用されています:

@register_passable("trivial")
struct Int:
    var value: __mlir_type.`!pop.scalar<index>`

    fn __init__(value: __mlir_type.`!pop.scalar<index>`) -> Self:
        return Self {value: value}
    ...

コピー可能であるためには copyinit メソッドが必要であり、 init メソッドと del メソッドを持つことができます。このデコレーターをつけておくと、コンパイラは可能な限りregister_passable型を通常マシンレジスタにアサインしようとします。

典型的なMojoプログラマにとって、このデコレータの観察可能な効果は微小で、かつてのC言語の register 宣言変数と類似しています。

register_passable型は、それ自身が@register_passableでない型のインスタンスを保持することができません。register_passable型のインスタンスは、予測可能なアイデンティティを持たないため、自己ポインタが安定しない/予測できません。register_passableの引数と結果は、ポインタで渡されるのではなく、CとC++に直接公開されます。この型の initcopyinit メソッドは暗黙のうちに静的で(Pythonの new のように)、inout selfの代わりに値でその結果を返します。このデコレータは、標準ライブラリのコアな型に広く使われることが予想されますが、一般的なアプリケーションレベルのコードでは無視しても大丈夫です。

上のIntの例では、実際にはこのデコレータの「些細な」バリエーションを使用しています。これは、上記のように受け渡し規則を変更しますが、コピーと移動のコンストラクタとデストラクタを禁止します。

マジック演算子
C++のコードには、値のライフサイクルと交差するマジック演算子が多数あり、「placement new」「placement delete」「operator=」のように、既存の値の上に再割り当てするものがあります。Mojoは、言語機能をすべて使用し、安全な構成要素の上で合成すれば安全な言語ですが、どのスタックもCスタイルのポインタの世界であり、安全ではありません。Mojoは実用的な言語であり、C/C++との相互運用と、Mojo自体にStringなどの安全な構成要素を直接実装することに関心があるため、安全でないものを表現する方法が必要です。

Mojo標準ライブラリのPointer[element_type]型は、MLIRMojoでこれらのC++と同等の安全ではない構成要素を実装する方法を望んでいます。最終的には、これらはすべて Pointer 型のメソッドに移行する予定ですが、それまでは、組み込み演算子として公開する必要があるものもあります。


この記事が参加している募集

AIとやってみた

この記事が気に入ったらサポートをしてみませんか?