【Rust】型黒魔術+associated constによる計算速度改善
Rustのtraitでは、associated type(関連型)とassociated const(関連定数)という機能が使えます。
trait SomeTrait {
const ASSOCIATED_CONST: usize;
type AssociatedType;
}
struct StructA;
impl SomeTrait for StructA{
const ASSOCIATED_CONST: usize = 10;
type AssociatedType = i32;
}
struct StructB;
impl SomeTrait for StructB{
const ASSOCIATED_CONST: usize = 100;
type AssociatedType = i128;
}
このように、traitを実装する構造体によって振る舞いを変えることができます。 associated typeについては、演算子オーバーロードをしたことがある人は見たことがあると思います。
さて、このassociated typeですが、ジェネリクスと組み合わせることで黒魔術的なことができます。
struct Math;
struct T;
struct F;
trait TAnd<A, B> { type Result; }
impl TAnd<F, F> for Math { type Result = F; }
impl TAnd<F, T> for Math { type Result = F; }
impl TAnd<T, F> for Math { type Result = F; }
impl TAnd<T, T> for Math { type Result = T; }
このような定義をすることで、 <Math as TAnd<F, T>>::Result というような書き方で構造体 F を表すことができます。
これの何が嬉しいかというと、これは繰り返し適用可能です。例えば、 <Math as TAnd<<Math as TAnd<F, T>>::Result, T>>::Result みたいなことができます。
さらに、trait境界を利用することでジェネリクスの利点を大きく活用できます。
例えば以下のような感じです。
trait TAnd3<A, B, C> { type Result; }
impl<A, B, C> TAnd3<A, B, C> for Math
where Math: TAnd<A, C>,
Math: TAnd<<Math as TAnd<A, B>>::Result, C>
{
type Result = <Math as TAnd<<Math as TAnd<A, B>>::Result, C>>::Result;
}
だんだん黒魔術じみたコードになってまいりました。
これで、Andが定義されている構造体に対して3変数のAndをとることができます(多分)。
また、associated constに対しても同じようにジェネリクスでいろいろできます。
さて、これを使えば何ができるかというと、行列のサイズをコンパイル時に決定できます。 これによって、コンパイラによる最適化時に行列のサイズがわかるため少し計算速度が上がると予測しました。 また、行列計算のときは行列のサイズが等しいなどの計算が定義されるための制約がありますが、これがコンパイル時に満たすことを保障できるメリットはあります。
ということで、実際に書いて検証してみました。コードは以下のGitHubにあげてあります。
型計算の黒魔術メイン部分については src/lib.rs を、行列演算の実装と測定プログラムについては ct_math_test/ をご覧ください
ある程度の速度で計算可能な行列の和と積の計算速度について、サイズをコンパイル時に決定したものがそうでないものの何%の時間で処理を完了したかというグラフがこちらです
なんと、和演算においては最大50%、積演算においては最大75%もの処理時間削減に成功しています。 極値でなくとも、10%程度の削減は安定してできているようで、十分な効果があるのではないでしょうか。 おそらくコンパイル時にループ展開がされることが効いているのだと思われますが、これほどまでとは。 ということで、コンパイラの最適化はすごいという話でした。
測定プログラムについて補足
計算のプログラムについては、両者ほぼコピペの同じものを利用しています。 積演算については、あまりにも遅かったのでキャッシュ最適化をしました。 その関係でボトルネックがキャッシュから移動している可能性があります。 和演算は特に最適化をしていないので、キャッシュがボトルネックになって最適化が効きにくくなっている可能性がありますので、二つの演算に対する比較は単純にはできません。 また、利用したプログラムはシングルスレッドですべての計算を行うものです。
