Benchmarks for EXLA.Backend doesn't work well

[zacky1972]

$ elixir -v
Erlang/OTP 25 [erts-13.0.3] [source] [64-bit] [smp:64:64] [ds:64:64:10] [async-threads:1] [jit:ns]

Elixir 1.13.4 (compiled with Erlang/OTP 25)

mix.exs:

defmodule NxScaffoldingTest.MixProject do
  use Mix.Project

  def project do
    [
      app: :nx_scaffolding_test,
      version: "0.1.0",
      elixir: "~> 1.14-rc",
      start_permanent: Mix.env() == :prod,
      deps: deps()
    ]
  end

  # Run "mix help compile.app" to learn about applications.
  def application do
    [
      extra_applications: [:logger]
    ]
  end

  # Run "mix help deps" to learn about dependencies.
  defp deps do
    [
      # {:dep_from_hexpm, "~> 0.3.0"},
      # {:dep_from_git, git: "https://github.com/elixir-lang/my_dep.git", tag: "0.1.0"}
      {:nx_scaffolding, "~> 0.1", github: "cocoa-xu/nx_scaffolding", runtime: false, only: [:dev, :test]},
      {:nx, "~> 0.3"},
      {:exla, "~> 0.3"}
    ]
  end

lib/mix/tasks/exla/backend.ex

defmodule Mix.Tasks.Exla.Benchmark do
  use Mix.Tasks.NxScaffolding.Benchmark
end

Run the benchmark and I got the following error:

$ mix exla.benchmark --backend EXLA.Backend
** (exit) exited in: GenServer.call(EXLA.Client, {:client, :host, [platform: :host]}, :infinity)
    ** (EXIT) no process: the process is not alive or there's no process currently associated with the given name, possibly because its application isn't started
    (elixir 1.13.4) lib/gen_server.ex:1019: GenServer.call/3
    (exla 0.3.0) lib/exla/backend.ex:137: EXLA.Backend.client_and_device_id/1
    (exla 0.3.0) lib/exla/backend.ex:34: EXLA.Backend.from_binary/3
    (nx_scaffolding_test 0.1.0) lib/mix/tasks/exla/backend.ex:2: Mix.Tasks.Exla.Benchmark.run/1
    (mix 1.13.4) lib/mix/task.ex:397: anonymous fn/3 in Mix.Task.run_task/3
    (mix 1.13.4) lib/mix/cli.ex:84: Mix.CLI.run_task/2

[cocoa-xu]

Hi @zacky1972! Thanks for reporting the issue, and sorry I was busy with other things in the past 4 days. I'll look into this later today!

[cocoa-xu]

In my test I found that EXLA seems to require some setup when running in mix test, but it works as expected when running with:

$ iex -S mix
iex> Mix.Tasks.Exla.Benchmark.run(["--backend", "EXLA.Backend"])

07:25:34.349 [info] XLA service 0x7f50544cc500 initialized for platform Host (this does not guarantee that XLA will be used). Devices:

07:25:34.354 [info]   StreamExecutor device (0): Host, Default Version
+---------------------+--------------+--------------+------------------+-----------------+----------------+-----------------+---------------+
| :callback           | :data_type   | :succeeded   | :mean            | :std            | :min           | :max            | :note         |
+---------------------+--------------+--------------+------------------+-----------------+----------------+-----------------+---------------+
| "Nx.dot/2"          | :f32         | 5            | "805.279 µs"     | "1.545 ms"      | "27.4 µs"      | "3.895 ms"      | "100x100"     |
| "Nx.dot/2"          | :f64         | 5            | "22.644 µs"      | "4.586 µs"      | "18.98 µs"     | "31.2 µs"       | "100x100"     |
| "Nx.constant/3"     | :f32         | 5            | "39.6 µs"        | "23.672 µs"     | "25.14 µs"     | "86.441 µs"     | nil           |
| "Nx.constant/3"     | :f64         | 5            | "17.502 µs"      | "10.831 µs"     | "7.48 µs"      | "34.991 µs"     | nil           |
+---------------------+--------------+--------------+------------------+-----------------+----------------+-----------------+---------------+

:ok

[cocoa-xu]

I'm unsure how to initialise EXLA (the XLA service) when running a mix task as it seems that erlang is complaining that the GenServer process is not alive or not started.

I saw that there is an EXLAHelper in the EXLA tests, but I don’t know which code are the necessary steps to set up the XLA service…maybe @josevalim knows how to do that properly?

[josevalim]

Make sure your mix tasks has:

@requirements [“app.start”]

So it starts the current app and all dependencies. :) does that work?

[cocoa-xu]

Oh, it works! Thank you @josevalim!

[zacky1972]

It works! Thanks!

$ mix exla.benchmark --backend EXLA.Backend
==> nx_scaffolding
Compiling 1 file (.ex)
Generated nx_scaffolding app
==> nx_scaffolding_test
Compiling 2 files (.ex)
Generated nx_scaffolding_test app

10:28:55.772 [info]  XLA service 0x7f3abc7267e0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:

10:28:55.777 [info]    StreamExecutor device (0): Host, Default Version
+---------------------+--------------+--------------+-----------------+-----------------+----------------+-----------------+---------------+
| :callback           | :data_type   | :succeeded   | :mean           | :std            | :min           | :max            | :note         |
+---------------------+--------------+--------------+-----------------+-----------------+----------------+-----------------+---------------+
| "Nx.dot/2"          | :f32         | 5            | "1.576 ms"      | "2.979 ms"      | "57.3 µs"      | "7.535 ms"      | "100x100"     |
| "Nx.dot/2"          | :f64         | 5            | "78.314 µs"     | "5.013 µs"      | "71.53 µs"     | "87.041 µs"     | "100x100"     |
| "Nx.constant/3"     | :f32         | 5            | "37.722 µs"     | "20.53 µs"      | "26.61 µs"     | "78.751 µs"     | nil           |
| "Nx.constant/3"     | :f64         | 5            | "41.408 µs"     | "14.234 µs"     | "31.32 µs"     | "68.92 µs"      | nil           |
+---------------------+--------------+--------------+-----------------+-----------------+----------------+-----------------+---------------+