まとめて計算(GPU)

SDPA の探索方向(dX, dY, dz)を求めるときの計算式を添付した。この辺の計算を GPU で計算する可能性について考えてみる。
まずは Bdz = r の線形方程式は、行列 B が対称半正定値なので Cholesky 分解で処理する。B_ij を要素を計算しながら計算が終わった部分は先に GPU に転送してしまう。でも現時点では GPU 側のメモリが少ないので、あまり効果は期待できない。
Bdz = r を解いてベクトル dz が求まったら、dz を用いて行列 dY を求める。dY の計算は GPU の必要は無い。
最後に行列 dX の計算だが、R Y^{-1} と X dY Y^{-1} を個別に計算してから引き算を行っている。特に後者は２回連続の行列積なので一度に GPU 側で処理した方が良い。
あと全般的に言えることは行列がブロック対角構造をしているので、現在と計算手順を変更してブロック別に計算を行う必要がある。つまり第一ブロックの計算 → 第二ブロックの計算 → というようにブロックに関するループを一番外側にする必要がある。