[arrayfire 3.7.2] RTX3080ではブロードキャスト(broadcasting)処理が使えない？

六花です。

グラフィックボードをGTX1080 tiからRTX3080に更新したらarrayfireでエラーが発生しました。
その時に対処した備忘録です。

■af::matmulのブロードキャスト処理がエラーの原因

まず、以下のコードがあったとします。

af::array A = af::matmul(B, C);

これは、場合によっては通る時もあれば、通らないときもあります。
通らない場合とは、3番目と4番目の次元を使った行列の時です。

ArrayFire matmul
https://arrayfire.org/docs/group__blas__func__matmul.htm#ga63306b6ed967bd1055086db862fe885b

行列というのはm×nの2次元までしか使いませんが、arrayfireにおいては並列処理のために3次元、4次元までデータを入力できます。
ニューラルネットワークにおいては上のリンクにおける表の一番下、「{M,K,b2,b3}{K,N,1,1}{M,N,b2,b3}」のパターンを一番使うと思います。
b2、b3の数だけ並列に処理を行ってくれるので、格段に早く処理が終わります。

ただ、これがRTX3080だとエラーになってしまいます。
そのため、以下のように編集する必要があります。

af::array A = (size_1, size_2, size_3);
for(int i = 0; i < size_3; ++i)
{
    A(af::span, af::span, i) = af::matmul(B(af::span, af::span, size_3), C(af::span, af::span, size_3));
}