RustでOSを書いた

はじめに

RISC-V CPU を FPGA 上に実装して、マイクロカーネル OS を Rust で書いて動かしてみました。

CPU について

RISC-VとChiselで学ぶ はじめてのCPU自作 に沿って RISC-V の CPU を作り、機能をエンハンスしました。

• 乗除算命令、RVC命令、ビット拡張命令の一部を追加
• 7段パイプライン化
• DRAM コントローラ
• 4KB命令キャッシュ、8KBデータキャッシュ
• 2ビット分岐予測
• 周辺コントローラ実装(SDC、UART、タイマー、割込コントローラ)

Arty A7-35T という FPGA ボード上で動作させています。

スーパーバイザーモードは実装していないので、仮想メモリは使えません。みんなで仲良くメモリを共有します。

CPU の実装はこちらに置いてあります。書籍のサポートリポジトリの fpga 実装版を fork して機能追加しています。
https://github.com/kinodjnz/riscv-chisel-book/tree/fpga-work/chisel-template/src/main/scala/fpga

OS について

ベアメタルでプログラムを動かすこともできるんですが、やっぱりいろいろ不便なので簡単な OS を動かすことにしました。

Resea という C で書かれたコンパクトでわかりやすいマイクロカーネル OS がありますので、これを Rust で移植してみることにしました。というかまだごく一部しか実装していないので移植と言うのはアレですが、とにかく一部実装しました。

Rust で静的メモリ領域を使う

カーネルで動的メモリを確保しようとすると複雑になるので、Resea では静的にメモリを確保しています。

例えば、タスク（一般的な OS ではプロセスと呼ばれるもの）の状態を管理するメモリは最初から BSS（起動時に0初期化される領域）に固定のタスク上限数分だけ確保しておいて、必要になったら初期化して使うということをします。

これを Rust で実装する場合には static mut な領域を使うことになります。しかし、static mut というのは unsafe の中でも特に危険なもののようです： static mutの危険性

static mut の変数からは、同じオブジェクトに対して p: &T と q: &mut T を作ったりできてしまうため、通常の借用ルールを破ってしまうことになります。

今回は Cell<T> という、領域をラップしてくれるやつを使うことにしました。Cell<T> はメモリレイアウトは T と同じですが、pub fn set(&self, val: T) を使うと中身を書き換えることができます。

self が mut でないのになぜ書き換えられるのかというと、Cell の中身の immutable な参照が取れないようになっているのと、中身が書き換えられることを前提にコンパイラが最適化しすぎないように特別扱いされているためであるようです。

カーネルのコードでの使用例です。

pub struct TaskPool {
    pub tasks: [UnsafeCell<Task>; config::NUM_TASKS as usize],
}
staic mut TASK_POOL = TaskPool {
    tasks: zeroed_array!(UnsafeCell<Task>, config::NUM_TASKS as usize),
};
pub struct Task {
    stack_ptr: Cell<usize>,
    // (略)
};

fn arch_task_init(tid: u32, task: &Task, pc: usize) {
    task.stack_ptr.set(init_stack(tid, pc));
}

arch_task_init の task 引数は immutable のはずですが、set を呼ぶことで中身が書き換えられます。

タスクスイッチ

いわゆるマルチタスク的なことを実現するために、タイマー割り込みで一定時間ごとにレジスタとスタックを別のタスクに切り替えます。そうすると同時に複数のタスクが走っているように見えます。

タスクスイッチ前のスタックにはここに至るまでに保存してきた情報が入っています。

レジスタ退避前のスタック状態

スイッチする際には callee saved レジスタをスタックに保存します。RISC-V では関数が破壊して良いレジスタと破壊してはいけないレジスタが定義されていて、後者を callee saved レジスタと呼びます。

レジスタ退避後のスタック状態

callee saved レジスタ以外にも ra と mepc を保存します。

ra はプロシージャからの戻りアドレスが入っているレジスタです。RISC CPU では戻りアドレスをスタックではなくレジスタに入れるアーキテクチャが多いですが、RISC-V も ra レジスタに戻りアドレスを入れることになっています。

mepc とはトラップからの戻りアドレスが入っているレジスタです。トラップは割り込み、システムコール、フォールトなどに起因して発生します。トラップからは専用の命令 mret でリターンしますが、そのリターン先アドレスが mepc になります。mret でリターンするとカーネルからユーザータスクに切り替わります。^[1]

レジスタを退避したら、スタックポインタを入れ替えてレジスタを復元すると別のタスクに切り替わった状態になります。

ここのレジスタやスタックを細かくいじる処理は Rust では書けないのでアセンブリで書くしかありません。

// extern "C" fn arch_task_switch(prev_sp: *mut usize, next_sp: usize);
arch_task_switch:
        // 切り替え元タスクのスタックにレジスタを退避
        addi    sp, sp, -64
        sw      ra, 60(sp)
        sw      gp, 56(sp)
        sw      tp, 52(sp)
        sw      s0, 48(sp)
        sw      s1, 44(sp)
        sw      s2, 40(sp)
        sw      s3, 36(sp)
        sw      s4, 32(sp)
        sw      s5, 28(sp)
        sw      s6, 24(sp)
        sw      s7, 20(sp)
        sw      s8, 16(sp)
        sw      s9, 12(sp)
        sw      s10, 8(sp)
        sw      s11, 4(sp)
        csrr    a2, mepc
        sw      a2, 0(sp)

        // prev_sp のアドレスにスタックポインタを退避
        sw      sp, 0(a0)
        // next_sp をスタックポインタに設定
        mv      sp, a1

        // 切り替え先タスクのスタックからレジスタを復元
        lw      a2, 0(sp)
        csrw    mepc, a2
        lw      s11, 4(sp)
        lw      s10, 8(sp)
        lw      s9, 12(sp)
        lw      s8, 16(sp)
        lw      s7, 20(sp)
        lw      s6, 24(sp)
        lw      s5, 28(sp)
        lw      s4, 32(sp)
        lw      s3, 36(sp)
        lw      s2, 40(sp)
        lw      s1, 44(sp)
        lw      s0, 48(sp)
        lw      tp, 52(sp)
        lw      gp, 56(sp)
        lw      ra, 60(sp)
        addi    sp, sp, 64

        // 切り替え先タスクにリターン
        ret

一度 arch_task_switch で退避されたタスクに再度切り替える場合は前回退避したスタックがあるのでこれで良いのですが、新規のタスクに切り替えたい場合は困ります。そこで、新規にタスクを作るときに空のスタック領域を割り当てるのと同時にダミーのスタックフレームを作ります。

この部分は Rust で書いてますが完全に unsafe なので別に Rust じゃなくて良い感があります。

static mut EXCEPTION_STACKS: ExceptionStack = ExceptionStack {
    stack: [[0; STACK_COUNT]; config::NUM_TASKS as usize],
};

fn init_stack(tid: u32, pc: usize) -> usize {
    unsafe {
        let stack: *mut u32 = EXCEPTION_STACKS.stack.get_unchecked_mut(tid as usize) as *mut u32;
        let sp = stack.add(STACK_COUNT).sub(16);
        extern "C" {
            fn arch_start_task();
        }
        write_volatile(sp, pc as u32); // mepc
        for i in 0..7 {
            // gp, tp, s0-s11
            write_volatile(sp.add(i * 2 + 1), 0);
            write_volatile(sp.add(i * 2 + 2), 0);
        }
        write_volatile(sp.add(15), arch_start_task as u32); // ra

        sp as usize
    }
}

arch_start_task:
        // mret 後の割り込みを許可する
        li      a0, 0x80
        csrs    mstatus, a0

        // カーネルからリターン
        mret

ra に初期化処理(arch_start_task)のアドレスを設定しておいて、mepc にユーザータスクのプログラムアドレスを設定しておくと、うまくタスクが起動します。

動かしてみる

OS といってもタスクスイッチしかできないので、シリアルコンソールにメッセージを吐くだけのテストプログラムを動かしてみました。

pub fn init_task() -> ! {
    printk!(b"init task started\n");
    loop {
        printk!(b"Hello, Resea\n");
        cycle::wait(cycle::clock_hz());
    }
}

pub fn worker_task() -> ! {
    cycle::wait(cycle::clock_hz() / 2);
    printk!(b"worker task started\n");
    loop {
        printk!(b"Hello, RISC-V\n");
        cycle::wait(cycle::clock_hz());
    }
}

コンパイルしてできるカーネルイメージを SD カードに書いてブートすると、無限ループする2つのタスクが並行に動きます。

ソースコード全体はこちらです。
https://github.com/kinodjnz/resea-rust

ただ自作した CPU のペリフェラルにも依存しているので簡単には動かせません…

おわりに

スレッド機能のようなものを CPU には特に実装していないのですが、ちゃんと並行処理が走るのは面白いですね。

OS 作りは先人の成果を借りているだけで特に新しいことはしていないのですが、イチから作れるという自己満足感があります。といってもまだ未実装機能だらけなので残りの実装もしていきたいところです。

脚注

1. 今回実装した CPU ではマシンモードしか実装していないため mret を使用しています。スーパーバイザーモードを使用する場合は命令やレジスタが異なります。