CUDA.jlにおけるdot演算子の挙動に関して

Question

実現したいこと

GPUを用いた行列演算において、dot演算子を用いることで少しでもメモリ消費を抑えて高速化したい。
後述する通りdot演算子を使わなければ問題ないが、dot演算子を使わないとメモリ消費量が増えてメモリ不足となるため。

前提

Julia(v1.8.5), CUDA.jl(v4.3.2)を用いています。
また実際のスクリプトでは数万＊数万の行列を複数使っており、nn_predict内の処理ももう少し複雑ですが、エラーが発生する最低限のスクリプトとして記載しました。

発生している問題・エラーメッセージ

後述するコードにおいて、p_buf.z2が処理ごとに変化します。
本来は、0.0 → 0.0 → 何らかの値(＊) → (＊)と同じ値　になってほしいですが
例えば　0.0 → 1.2 → 22 → -8.7　となります。

該当のソースコード

using CSV, DataFrames
using CUDA
using Random, Distributions

Random.seed!(48)

mutable struct TrainParams
    w::CuArray{Float32, 2}
    b::CuArray{Float32, 2}
end
mutable struct BufParams
    z1::CuArray{Float32, 2}
    z2::CuArray{Float32, 2}
    z3::CuArray{Float32, 2}
end

function nn_predict(p_buf::BufParams, p1::TrainParams, p2::TrainParams, p3::TrainParams, x_train::CuArray{Float32})
    println(sum(p_buf.z2))
    p_buf.z1 .= p1.w * x_train .+ p1.b;
    println(sum(p_buf.z2))
    p_buf.z2 .= p2.w* p_buf.z1 .+ p2.b;
    println(sum(p_buf.z2))
    p_buf.z3 .= p3.w * p_buf.z2 .+ p3.b;
    println(sum(p_buf.z2))
    
    return nothing
end


dist = Normal(0, sqrt(2.0f0/1000))
x_train = cu(rand(dist, 1000, 10))

p1 = TrainParams(
                cu(rand(dist, 1000, 1000)),
                cu(rand(dist, 1000, 1))
        )
p2 = TrainParams(
                cu(rand(dist, 1000, 1000)),
                cu(rand(dist, 1000, 1))
        )
p3 = TrainParams(
                cu(rand(dist, 1000, 1000)),
                cu(rand(dist, 1000, 1))
        )

buf = CUDA.zeros(1000, 10)
p_buf = BufParams(
                        buf,
                        buf,
                        buf
                )

nn_predict(p_buf, p1, p2, p3, x_train)

試したこと

dot演算子ではなく代入にすれば、きちんと
0.0 → 0.0 → 何らかの値(＊) → (＊)と同じ値
となります。
また、CUDAを使わず全てArray{}で行うとdot演算子でもきちんと
0.0 → 0.0 → 何らかの値(＊) → (＊)と同じ値
となります。

質問

前述したようにdot演算子を使わないとメモリ不足になるため出来るだけdot演算子を使いたいのですが

1. なぜdot演算子を使うと値が処理ごとに変化するのでしょうか？
2. どのように修正すればdot演算子を使っても想定通りの挙動になるのでしょうか？
3. もしくはdot演算子を使わなくてもメモリ消費量を抑える方法があれば教えてください

以上です。

よろしくお願いします。

Answer 1

.=で値がそれぞれ変化するのは，例示されているプログラムからだけで言うと，

buf = CUDA.zeros(1000, 10)
p_buf = BufParams(
                        buf,
                        buf,
                        buf
                )

の部分で，z1, z2, z3が同じバッファに割り当てられているからだと思います。
この部分をそれぞれ領域を確保する

p_buf = BufParams(
                        CUDA.zeros(1000, 10),
                        CUDA.zeros(1000, 10),
                        CUDA.zeros(1000, 10)
                )

にすると改善するのではないでしょうか?

=なら，その度に新規でメモリが割り合てられるので最初z1,z2,z3が同じバッファでも
代入時にそれぞれ別のメモリとなるので問題が露見しないことになります。

.=だと最初に確保したメモリをそのまま使い続けるので，z1,z2,z3が同じバッファを利用している場合，
それぞれの値の更新時に，z1,z2,z3の値が変化することになります。

蛇足

この後に書く事は余計なお世話なのですが，外部コンストラクタを使って

BufParams(r::Int, c::Int) = BufParams(
    CUDA.zeros(r, c), CUDA.zeros(r, c), CUDA.zeros(r, c)
)

のように定義しておくと，

p_buf = BufParams(1000, 10)

のような形で利用できて便利かもしれません。

Comments

[NULNUL]

ありがとうございます。
`BufParams()`の初期化の時に代入としてメモリが確保されるのだと勝手に思っていました。
ご教示いただいた通りに修正したところ、想定通りの挙動となりました。

また、外部コンストラクタの書き方を知らなかったのでご教示くださりありがとうございました。
こちらに変更しようと思います。

その上で恐縮なのですが、差し支えなければ以下についてご教示いただけますと幸いです。
+ 原因が「z1,z2,z3が同じバッファであること」でしたが、なぜCPUの時にはdot演算子でも問題なく動いたのでしょうか？どのようなキーワードで調べればいいかも分からないため、参考になるサイトを教えていただけるだけでもありがたいです。

[ujimushi_sradjp]

基本的には返答しないのですが，所用で休暇をとっているので一言だけ。
Base.zeros(1000, 10)は型指定なしだとFloat64型の行列になる(CUDA.zerosはデフォルトがFoat32型)ので，
コンストラクタの引数に渡した時にFloat64→Float32の型変換が発生し，Float32型の行列が新規に作成されているので，
見かけ上うまくいっていたち見えたのではないかと。

CPUでもbuf =zeros(Float32, 1000, 10)とかだと同じ問題が発生すると想像できます。

[NULNUL]

理解できました。ありがとうございました。