Rust から Intel MKL を使う

諸事情で、バッチのサイズ分だけ sgemm (単精度一般行列積) を呼ぶ必要があった。 最初は OpenBLAS の sgemm を FFI で呼び出すというのをバッチサイズ回やってみたりしてて、効率よくなさそうだなーと思ってたところ、こんなのを見つけた:

software.intel.com

つまり、この cblas_sgemm_batch を使うためには OpenBLAS から MKL に移行する必要がある。 しかし Rust には blasバインディング (-sys) は充実しているが MKL のものはない、が、辛うじてバインディングを作るための -src クレートは見つかった:

github.com

これを使って MKL の C-API を Rust から呼び出す。手順:

extern crate libc;
extern crate intel_mkl_src;

// libc クレートとかを使って頑張って cblas_sgemm_batch を extern "C" する

#[repr(C)]
enum CblasTranspose {
    CblasNoTrans = 111,
    CblasTrans = 112,
    CblasConjTrans = 113,
}

#[repr(C)]
enum CblasLayout {
    CblasRowMajor=101, 
    CblasColMajor=102
}

type MklInt = i64;

extern "C" {
    fn cblas_sgemm_batch(
        layout: CblasLayout,
        transa_array: *const CblasTranspose,
        transb_array: *const CblasTranspose,
        m_array: *const MklInt,
        n_array: *const MklInt,
        k_array: *const MklInt, 
        alpha_array: *const libc::c_float,
        a_array: *const *const libc::c_float,
        lda_array: *const MklInt,
        b_array: *const *const libc::c_float,
        ldb_array: *const MklInt,
        beta_array: *const libc::c_float,
        c_array: *mut *mut libc::c_float,
        ldc_array: *const MklInt,
        group_count: MklInt,
        group_size: *const MklInt,
    );
}

これで、cblas_sgemm_batch(.. が呼び出せる。 OpenBLAS は大抵リンク時に fortran コンパイラを要求してきて面倒なので、MKL で済むならその方がいいと思った。