Rustでallocを使わずにページ境界に合わせたメモリアロケーションをするには

κeenです。こちらの記事が面白かったので関連したネタを。
https://qiita.com/moriai/items/67761b3c0d83da3b6bb5

コメントで書くか迷ったのですがコメントで書かれても「だから何？」となりかねないので記事に通知がいくけど直接はコメントしない奥ゆかしい方法で投稿します。

さて、タイトルの通りページ境界に合わせたメモリアロケーションをする方法です。つまり特定のアドレスにアラインメント（境界合わせ）されたメモリを確保します。
Rustには repr(align(N))というアトリビュートがあるのでそれを使うとコンパイラがアラインメントを勝手にやってくれます。
なのでこういう構造体を確保して、確保した領域を使えばアラインメントされたことになります。4096要素なのは配列で並べたときに隙間がないようにするためです

#[repr(align(4096))]
struct AlignedData {
    data: [u8; 4096],
}

一見すると任意長のメモリは Dynamically Sized Typesで表現できるので、4096要素の配列である必要はなさそうです。すなわち #[repr(align(4096))] struct AlignedData {data: [u8]} でも十分じゃないかという疑問が湧きます。しかしDSTsと各種アロケーションのAPIの噛み合わせが悪く、データを確保する手段として使えそうにないので大人しくSized Typesにします。

さて、これを用いて aligned_alloc は以下のように書けます

use std::mem;

const ALIGN: usize = 4096;

#[repr(align(4096))]
struct AlignedData {
    #[allow(unused)]
    data: [u8; ALIGN],
}

// 4096要素の配列にはDefaultがないので手で実装する
impl Default for AlignedData {
    fn default() -> Self {
        AlignedData { data: [0; ALIGN] }
    }
}

fn aligned_alloc(size: usize) -> Box<[u8]> {
    // unsafeは一連のコードを囲んで、ブロック内のコードを見れば
    // 正しいかどうか分かるようにする
    unsafe {
        // 余りを切り上げながら4096で割る
        let size = (size + ALIGN - 1) / ALIGN;
        // メモリ確保
        let mut vec = Vec::<AlignedData>::with_capacity(size);
        // データの書き込み
        vec.resize_with(size, Default::default);
        // アライン付き確保が終わったのでデータをu8にする
        let mut data = mem::transmute::<_, Vec<u8>>(vec);
        // そのままだと4096分の1の要素しかないのでメタデータを変更する
        data.set_len(size * ALIGN);
        data.into_boxed_slice()
    }
}

Vecを使って領域を確保します。また、Rustから扱いやすいように返り型は Box<[u8]> にしておきました。

これを使うコードを書いてみましょう。

fn print_info(mem: &[u8]) {
    // 要素にアクセスしなくてもアドレスはとれる
    let addr = mem as *const _ as *const u8 as u64;
    let mut bound: u64 = 1;
    while addr & bound == 0 {
        bound <<= 1;
    }
    println!(
        "size: {:>10}  addr: 0x{:>012x}  bound: {:>7}",
        mem.len(),
        addr,
        bound
    );
}


fn main() {
    const ONEGB: usize = 1024 * 1024 * 1024;
    let mut size: usize = 4;
    while size <= ONEGB {
        let mem0 = aligned_alloc(size);
        print_info(&mem0);
        let mut mem1 = aligned_alloc(size);
        mem1[0] = 0;
        print_info(&mem1);
        size *= 4;
    }
}

実行します。環境はUbuntu 20.10、glibc 2.32、Rust 1.48.0です。

size:       4096  addr: 0x563f83673000  bound:    4096
size:       4096  addr: 0x563f83675000  bound:    4096
size:       4096  addr: 0x563f83673000  bound:    4096
size:       4096  addr: 0x563f83677000  bound:    4096
size:       4096  addr: 0x563f83673000  bound:    4096
size:       4096  addr: 0x563f83679000  bound:    4096
size:       4096  addr: 0x563f83673000  bound:    4096
size:       4096  addr: 0x563f8367b000  bound:    4096
size:       4096  addr: 0x563f83673000  bound:    4096
size:       4096  addr: 0x563f8367d000  bound:    4096
size:       4096  addr: 0x563f83673000  bound:    4096
size:       4096  addr: 0x563f8367f000  bound:    4096
size:      16384  addr: 0x563f83681000  bound:    4096
size:      16384  addr: 0x563f83686000  bound:    8192
size:      65536  addr: 0x563f8368b000  bound:    4096
size:      65536  addr: 0x563f8369c000  bound:   16384
size:     262144  addr: 0x7fa56fced000  bound:    4096
size:     262144  addr: 0x7fa56fcab000  bound:    4096
size:    1048576  addr: 0x7fa56fc2d000  bound:    4096
size:    1048576  addr: 0x7fa56fb2b000  bound:    4096
size:    4194304  addr: 0x7fa56f92d000  bound:    4096
size:    4194304  addr: 0x7fa56f52b000  bound:    4096
size:   16777216  addr: 0x7fa56ed2d000  bound:    4096
size:   16777216  addr: 0x7fa56dd2b000  bound:    4096
size:   67108864  addr: 0x7fa56bd2d000  bound:    4096
size:   67108864  addr: 0x7fa567d2b000  bound:    4096
size:  268435456  addr: 0x7fa55fd2d000  bound:    4096
size:  268435456  addr: 0x7fa54fd2b000  bound:    4096
size: 1073741824  addr: 0x7fa52fd2d000  bound:    4096
size: 1073741824  addr: 0x7fa4efd2b000  bound:    4096

[u8; 4096] でデータを確保している影響で最初の方が4096の倍数に切り上げられてしまっていますが、実際に少量のデータをアラインメントして確保したいことはそんなに多くなさそうなのでそこまで問題にならないんじゃないでしょうか。

因みに元の記事からの明らかなデグレとして、アラインメントを動的に設定できないというのがあります。
多くの場合ページサイズは固定ですが、 sysconf(_SC_PAGESIZE) を使って動的に取得する方が行儀のいいコードなのであまり褒められたものではないですね。
そういった意味である種のネタ、 repr(align(N)) という機能があるんですよという紹介でした。