Nx速度競争 def vs defn

この記事で、@zacky1972 さんから、defnには最適化高速化の効果もある。とコメントいただきました。
kino_bencheeも使ってみたかったので、勉強も兼ねてやってみます。

速度比較結果

defn、compiler指定無し（青色）でも速くなってます。
Nx.Defn.jit(dot_fn, compiler: EXLA)でjitコンパイルした場合（オレンジ色）は飛びぬけて速い！

default_defn_optionsにcompiler: EXLAを指定しておけば、デフォルトのcompileがEXLAになります。
defnで記述するだけで、この効果が得られます。
compiler: EXLAをつけましょう！効果は絶大です

import Config
config :nx, :default_backend, EXLA.Backend
config :nx, :default_defn_options, compiler: EXLA

この例は、jitでコンパイルしましたが、defnで記述した場合も同じ効果が得られるとおもいます。
そういえば、Axonのコードを実行する場合、この記述を付ける事でかなり速くなってました。この働きですね。

以下は、測定した時のlivebookです

速度比較のlivemd

Mix.install(
  [
    {:nx, "~> 0.3.0"},
    {:exla, "~> 0.3.0"},
    {:kino, github: "livebook-dev/kino", branch: :main, override: true},
    {:kino_benchee, "~> 0.1", github: "livebook-dev/kino_benchee", branch: "main"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend
      # default_defn_options: [compiler: EXLA] #比較kのためデフォルトでは無効にする
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda111"
  ],
  # XLA_TARGETの変更で、exlaを再コンパイルしたい時はtrueにする
  force: true
)

速度比較

dot_fn = fn ->
  a = Nx.random_uniform({100, 100})
  b = Nx.random_uniform({100, 100})
  Nx.dot(a, b)
end

dot_jit = Nx.Defn.jit(dot_fn)
dot_jit_exla = Nx.Defn.jit(dot_fn, compiler: EXLA)

Benchee.run(
  %{
    "def" => dot_fn,
    "Nx.Defn.jit" => dot_jit,
    "Nx.Defn.jit_exla" => dot_jit_exla
  },
  time: 10,
  memory_time: 2
)

参考
https://qiita.com/zacky1972/items/c8deaf2413efa726ea79
https://qiita.com/RyoWakabayashi/items/afbd4d2187368346d5a9