GPU Benchmarks

# NOTE: see environment variables defined in secrets 
# check installed with nvidia-smi --list-gpus 
# see https://elixirforum.com/t/compiling-xla-with-cuda-11-2-support-needed/42528 

# Mix.install([
#   {:nx, "~> 0.6"},
#   {:exla, "~> 0.6"},
#   {:benchee, "~> 1.2"}
# ])

# Application.put_env(:exla, :clients,
#   cuda: [platform: :cuda],
#   host: [platform: :host]
# )

Mix.install(
  [
    {:nx, "~> 0.6"},
    {:exla, "~> 0.6"},
    {:benchee, "~> 1.2"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA]
    ],
    exla: [
      default_client: :cuda,
      clients: [
        host: [platform: :host],
        cuda: [platform: :cuda]
      ]
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda120"
  ]
)

Section

defmodule Softmax do
  import Nx.Defn
  defn(softmax(n), do: Nx.exp(n) / Nx.sum(Nx.exp(n)))
end

key = Nx.Random.key(42)
{tensor, _key} = Nx.Random.uniform(key, shape: {1_000_000})

Benchee.run(
  %{
    "JIT with EXLA" => fn ->
      apply(EXLA.jit(&Softmax.softmax/1), [tensor])
    end,
    "Regular Elixir" => fn ->
      Softmax.softmax(tensor)
    end
  },
  time: 10
)