Juliaの@code_warntypeで「遅いコード」の匂いを見つける

はじめに

Project.tomlやManifest.tomlを使うと、解析コードや自作パッケージの依存関係をかなりきれいに管理できます。
ただ、環境を整えたあとに気になるのが、

Juliaで書いたコードは、本当に速く動いているのか？

という点です。

Juliaは高速な言語として紹介されることが多いですが、何を書いても自動的に速くなるわけではありません。
特に重要なのが、コンパイラが型を推論しやすいコードになっているかです。

この記事では、Juliaの@code_warntypeを使って、型推論がうまくいっていない箇所、つまり「遅いコードの匂い」を見つける方法を紹介します。

Julia公式ドキュメントでも、@code_warntypeは型の不確かさが生じる式を見つけるのに役立つ道具として説明されています。

対象読者

この記事は、以下のような人を想定しています。

• Juliaで簡単な関数を書いたことがある
• 研究やデータ解析でJuliaを使っている
• コードは動くが、なぜか遅いことがある
• @code_warntypeを見たことはあるが、何を読めばいいかわからない
• Vector{Any}やBody::Anyが何となく怖い

なお、この記事では厳密なコンパイラ内部の話までは踏み込みません。
まずは「実用上、どこを見ればよいか」に絞って説明します。

@code_warntypeとは

@code_warntypeは、Juliaが関数をコンパイルするときに、変数や戻り値の型をどのように推論しているかを見るためのマクロです。

使うには、まずInteractiveUtilsを読み込みます。

using InteractiveUtils

例えば、次のように使います。

@code_warntype f(x)

出力には、Juliaが推論した変数の型や、戻り値の型が表示されます。

特に最初のうちは、細かい中間表現を全部読む必要はありません。
まず見るべきなのは、出力の下の方に出てくる

Body::...

の部分です。

ここが例えば

Body::Int64
Body::Float64

のように具体的な型になっていれば、比較的安心です。

一方で、

Body::Any
Body::Union{Int64, String}

のようになっている場合は、戻り値の型がはっきり決まっていない可能性があります。

このようなコードは、場合によっては実行時の動的ディスパッチや余計なメモリアロケーションにつながり、性能が落ちる原因になります。

ただし、ここで大事なのは、

@code_warntypeは「遅さを直接測る道具」ではない

ということです。

実際に遅いかどうかは、@timeやBenchmarkTools.jlなどで測る必要があります。
@code_warntypeは、あくまで「型推論が怪しい場所を見つける道具」です。

例1：戻り値の型が変わる関数

まずは一番わかりやすい例です。

using InteractiveUtils

function check_value(x)
    if x > 0
        return x
    else
        return "negative"
    end
end

@code_warntype check_value(1)

この関数は、x > 0のときは数値を返します。

return x

一方、x <= 0のときは文字列を返します。

return "negative"

つまり、この関数の戻り値は、入力値によって

Int64

になることもあれば、

String

になることもあります。

そのため、@code_warntypeを見ると、環境にもよりますが、戻り値が次のように推論されます。

Body::Union{Int64, String}

これは、

この関数の戻り値はInt64かもしれないし、Stringかもしれない

という意味です。

このように、戻り値の型が入力値によって変わる関数は、型不安定になりやすいです。

修正例：正常系では同じ型を返す

例えば、負の値を許さない関数なら、次のように書くことができます。

function check_value_fixed(x)
    if x <= 0
        error("x must be positive")
    end

    return x
end

@code_warntype check_value_fixed(1)

この場合、正常に値を返すときはxをそのまま返すだけなので、戻り値の型は具体的になります。

Body::Int64

もちろん、エラーを投げるのが常に正解というわけではありません。
重要なのは、

1つの関数の中で、返す値の型が意味もなくバラバラになっていないか

を意識することです。

例えば、失敗を表したい場合は、nothingを返す設計にすることもあります。

function find_positive(x)
    if x > 0
        return x
    else
        return nothing
    end
end

この場合、戻り値は

Union{Nothing, Int64}

のようになります。

小さなUnionは必ずしも大きな問題になるとは限りません。
そのため、Unionが出たら即修正、というよりは、

そのUnionは設計として自然か？

を確認するのがよいと思います。

例2：[]で配列を作るとVector{Any}になる

次は、Julia初心者がかなりやりがちな例です。

function bad_sum(n)
    xs = []

    for i in 1:n
        push!(xs, i)
    end

    return sum(xs)
end

@code_warntype bad_sum(10)

一見すると、1からnまでを配列に入れて合計しているだけです。

しかし、ここで問題なのはこの行です。

xs = []

Juliaで[]と書くと、これは

Vector{Any}

になります。

つまり、xsは「何でも入る配列」です。

整数も入るし、浮動小数点数も入るし、文字列も入ります。

xs = []
push!(xs, 1)
push!(xs, 2.0)
push!(xs, "hello")

このような配列は柔軟ですが、数値計算には向いていません。
各要素の型が具体的に決まらないため、Juliaのコンパイラが最適化しづらくなります。

Julia公式ドキュメントでも、抽象型パラメータを持つコンテナは避けるのがよいと説明されています。例えばVector{Real}のような配列は任意のReal型を入れられるため、要素を効率よく連続したメモリに並べられない場合があります。一方で、Vector{Float64}のように具体型にすると効率的に扱いやすくなります。

修正例：配列の型を最初に決める

この例では整数だけを入れたいので、次のように書きます。

function good_sum(n)
    xs = Int[]

    for i in 1:n
        push!(xs, i)
    end

    return sum(xs)
end

@code_warntype good_sum(10)

Int[]は、

Vector{Int}

を作ります。

これにより、xsにはIntの値が入ることが最初からわかります。
そのため、Juliaはより効率的なコードを生成しやすくなります。

また、配列の長さが最初からわかっているなら、次のように事前確保することもできます。

function better_sum(n)
    xs = Vector{Int}(undef, n)

    for i in 1:n
        xs[i] = i
    end

    return sum(xs)
end

ただし、この例だけなら、そもそも配列を作らずに

sum(1:n)

と書くのが一番簡単です。

ここで言いたいのは、

[]で雑に配列を作るとVector{Any}になりやすい

という点です。

研究コードでは、あとから値を詰めていく配列をよく作ります。
そのときに何も考えずに

result = []

と書いていると、Vector{Any}が増えていきます。

数値だけを入れるなら、

result = Float64[]

整数だけなら、

result = Int[]

のように、最初に型を決めておくのがよいです。

例3：構造体のフィールドが抽象型になっている

次は、少し実践寄りの例です。

例えば、解析パラメータを入れる構造体を作るとします。

struct BadParam
    value::AbstractFloat
end

function double_value(p)
    return 2p.value
end

p = BadParam(1.0)

@code_warntype double_value(p)

ここで、フィールドの型を

value::AbstractFloat

としています。

一見すると、

浮動小数点数を入れたいからAbstractFloatでよさそう

と思うかもしれません。

しかし、AbstractFloatは抽象型です。
Float64、Float32、BigFloatなどをまとめる上位の型です。

つまり、BadParamという型だけを見ても、valueの中身が具体的に何型なのかはわかりません。

これも、Juliaが最適化しづらくなる原因になります。

Julia公式ドキュメントでも、構造体のフィールドを抽象型にすると、コンパイラが高性能なコードを生成しにくくなるため、型パラメータを使う方法が紹介されています。

修正例：型パラメータを使う

この場合は、次のように書くのがJuliaらしいです。

struct GoodParam{T<:AbstractFloat}
    value::T
end

function double_value(p)
    return 2p.value
end

p = GoodParam(1.0)

@code_warntype double_value(p)

このとき、pの型は

GoodParam{Float64}

になります。

つまり、構造体の型そのものに、

このvalueはFloat64です

という情報が含まれます。

Float32を渡せば、

GoodParam{Float32}

になります。

このように、型パラメータを使うと、抽象的な設計を保ちながら、具体的な型情報もコンパイラに伝えられます。

型注釈を増やせば速くなる、ではない

ここで注意したいのは、

Juliaでは、型注釈をたくさん書けば速くなるわけではない

ということです。

例えば、次のように書いても、本質的な改善にならないことがあります。

function f(x::Int)
    y::Any = x
    return y
end

引数に::Intを付けても、途中でAnyにしてしまえば、型情報は失われます。

Juliaでは、型注釈で無理やり速くするというより、

• 関数を小さく分ける
• 戻り値の型が自然に決まるようにする
• 配列の要素型を具体的にする
• 構造体のフィールドを具体型または型パラメータにする
• グローバル変数に依存しすぎない

といった書き方が重要です。

例4：グローバル変数に依存する

研究用のスクリプトでは、つい次のようなコードを書きがちです。

scale = 2.0

function apply_scale(x)
    return scale * x
end

@code_warntype apply_scale(3.0)

トップレベルにあるscaleはグローバル変数です。

グローバル変数はあとから別の型の値に変更される可能性があります。

scale = "large"

このような変更があり得るため、Juliaはグローバル変数の型を安定して推論しづらくなります。

修正方法の1つは、定数として定義することです。

const SCALE = 2.0

function apply_scale_const(x)
    return SCALE * x
end

@code_warntype apply_scale_const(3.0)

または、関数の引数として渡します。

function apply_scale_arg(x, scale)
    return scale * x
end

@code_warntype apply_scale_arg(3.0, 2.0)

研究コードでは、解析パラメータをグローバル変数として置きたくなることがあります。
しかし、性能や見通しを考えると、パラメータは構造体や関数の引数として渡す方が扱いやすいです。

例えば、

struct AnalysisConfig{T<:Real}
    scale::T
end

function apply_scale_config(x, config)
    return config.scale * x
end

config = AnalysisConfig(2.0)

@code_warntype apply_scale_config(3.0, config)

のようにすると、設定値をまとめつつ、型情報も保ちやすくなります。

@code_warntypeを見るときの基本方針

最初のうちは、@code_warntypeの出力を全部読もうとしなくてよいです。

まずは以下の順番で見るとわかりやすいです。

1. Body::...を見る

最初に見るべきなのは、戻り値の型です。

Body::Float64

のように具体型なら、まずは安心です。

一方で、

Body::Any

になっていたら、かなり怪しいです。

Body::Union{...}

の場合は、そのUnionが自然なものか確認します。

例えば、Union{Nothing, Int64}のような小さなUnionは、設計として自然な場合もあります。
一方で、巨大なUnionやAnyが出ている場合は、見直した方がよいことが多いです。

2. Anyが出ていないか見る

Anyは、Juliaにとって「何でもあり」の型です。

柔軟ですが、数値計算では性能上の問題になりやすいです。

特に、

Vector{Any}

や

Body::Any

が出てきたら注意します。

3. 抽象型のコンテナを疑う

次のような型は注意です。

Vector{Any}
Vector{Real}
Vector{AbstractFloat}

数値計算では、できるだけ

Vector{Float64}
Vector{Int}
Matrix{Float64}

のような具体的な型にした方が扱いやすいです。

4. 構造体のフィールドを見る

構造体で

value::Real
value::AbstractFloat
array::AbstractVector

のように書いている場合は、本当にそれでよいか確認します。

多くの場合、型パラメータを使って

struct MyType{T}
    value::T
end

のようにした方が、Juliaの性能を引き出しやすいです。

赤い表示を全部消す必要はない

@code_warntypeを見ると、赤色や警告っぽい表示が出て不安になることがあります。

しかし、赤い表示を全部消す必要はありません。

@code_warntypeは、性能上問題になる可能性がある型を強調して表示します。
ただし、表示されたものがすべて実際の性能問題になるとは限りません。

公式ドキュメントでも、code_warntypeは潜在的に問題になりうる非具象型を強調する一方で、実際には影響が小さいものも含まれ得ることが説明されています。

大事なのは、

• Body::Anyになっていないか
• 数値計算の内側のループでAnyが出ていないか
• 配列がVector{Any}になっていないか
• 構造体のフィールドが抽象型になっていないか

を確認することです。

@code_warntypeとベンチマークは役割が違う

最後に、@code_warntypeとベンチマークの違いを整理します。

@code_warntypeは、型推論の結果を見る道具です。

@code_warntype f(x)

一方で、実際の実行時間を見るには、例えば@timeを使います。

@time f(x)

より正確に測りたい場合は、BenchmarkTools.jlを使います。

using BenchmarkTools

@btime f($x)

つまり、

• @code_warntype：型推論がうまくいっているかを見る
• @time：ざっくり実行時間とメモリ確保を見る
• @btime：より正確にベンチマークする

という使い分けです。

遅いコードを見つけるときは、まず@timeや@btimeで遅い場所を確認し、その関数に対して@code_warntypeを見るのがよいと思います。

研究コードでの使い方

研究用の解析コードでは、最初からきれいなパッケージを書くのは難しいです。

最初は、どうしても

analysis.jl

のような巨大スクリプトになりがちです。

その場合でも、重い処理を少しずつ関数に切り出して、

@code_warntype analyze_event(event)
@code_warntype fit_spectrum(data)
@code_warntype calculate_response(params)

のように確認していくと、問題のある場所を見つけやすくなります。

特に、以下のようなコードは見直し候補です。

result = []

params = Dict()

struct Config
    value::Any
end

function f(x)
    if condition
        return 1.0
    else
        return "failed"
    end
end

これらは全部、Juliaでは型情報が失われやすい書き方です。

もちろん、柔軟性が必要な場面ではAnyを使っても構いません。
ただ、数値計算の中心部分では、できるだけ具体的な型で書く方が性能を出しやすいです。

まとめ

この記事では、Juliaの@code_warntypeを使って、型不安定なコードを見つける基本を紹介しました。

ポイントは以下の通りです。

• @code_warntypeは「遅さを測る道具」ではなく、「型推論の結果を見る道具」
• まずはBody::...を見る
• Body::AnyやVector{Any}は注意する
• []で作った配列はVector{Any}になる
• 戻り値の型が場合によって変わる関数は型不安定になりやすい
• 構造体のフィールドに抽象型を使うと、性能が落ちる原因になることがある
• 赤い表示を全部消す必要はない
• 実際の速度は@timeやBenchmarkTools.jlで測る

Juliaでは、型をガチガチに書くというより、

コンパイラが自然に型を推論できるように書く

ことが大切です。

研究やデータ解析のコードでも、巨大なスクリプトを一気に最適化する必要はありません。
まずは重い処理を小さな関数に分けて、@code_warntypeで型の流れを確認するところから始めるとよいと思います。