Skip to content

Instantly share code, notes, and snippets.

@Munksgaard
Last active October 20, 2023 07:44
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Munksgaard/cf4978cacd293cd55b29113facf33453 to your computer and use it in GitHub Desktop.
Save Munksgaard/cf4978cacd293cd55b29113facf33453 to your computer and use it in GitHub Desktop.

Futlixir

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

Nx.global_default_backend(EXLA.Backend)

Code.compile_file("foo.ex", "/home/munksgaard/src/futlixir/present")

defmodule RandomMatrix do
  def new(n, m) do
    for _ <- 1..n do
      for _ <- 1..m do
        :rand.uniform()
      end
    end
  end
end

Creating a Futhark context

# Create a Futhark context
{:ok, cfg} = Foo.futhark_context_config_new()
{:ok, ctx} = Foo.futhark_context_new(cfg)

Some simple computations

# A few random matrices
a = [[1, 2], [3, 4]]
b = [[5, 6], [7, 8]]

# Move to Futhark
{:ok, a_fut} = Foo.futhark_f64_2d_from_list(ctx, a)
{:ok, b_fut} = Foo.futhark_f64_2d_from_list(ctx, b)

# Compute matrix multiplication
{:ok, c_fut} = Foo.futhark_entry_matmul(ctx, a_fut, b_fut)

# Return a list
{:ok, result} = Foo.futhark_f64_2d_to_list(ctx, c_fut)

# Remember to free arrays
:ok = Foo.futhark_free_f64_2d(ctx, a_fut)
:ok = Foo.futhark_free_f64_2d(ctx, b_fut)
:ok = Foo.futhark_free_f64_2d(ctx, c_fut)

result

Benchmarks

Not representative!

a = RandomMatrix.new(2048, 2048)
b = RandomMatrix.new(2048, 2048)

{:ok, a_fut} = Foo.futhark_f64_2d_from_list(ctx, a)
{:ok, b_fut} = Foo.futhark_f64_2d_from_list(ctx, b)

# Just to be safe
:ok = Foo.futhark_context_sync(ctx)

{futlixir_time, :ok} =
  :timer.tc(fn ->
    {:ok, c_fut} = Foo.futhark_entry_matmul(ctx, a_fut, b_fut)
    {:ok, _} = Foo.futhark_f64_2d_to_list(ctx, c_fut)
    :ok
  end)

a_nx = Nx.tensor(a)
b_nx = Nx.tensor(b)

{nx_time, :ok} =
  :timer.tc(fn ->
    _ = Nx.dot(a_nx, b_nx) |> Nx.to_list()
    :ok
  end)

%{futlixir_time: futlixir_time, nx_time: nx_time}

Futlixir

**Work-in-progress**

A tool to run Futhark code from Elixir.

What is Futhark?

  • Functional array language, targetting GPUs.
  • Strongly, statically typed, syntax like ML/Haskell (or Elixir).
  • Completely pure (e.g. no IO).
  • Specialized for parallel computations.
  • Programs compile to C or Python libraries (among other things).
  • Fast: Competitive with highly optimized hand-written CUDA/OpenCL.

Link: https://futhark-lang.org/

**Disclaimer:** Stockholm syndrome

What does Futhark look like?

def dotprod [n] (xs: [n]f64) (ys: [n]f64): [n]f64 =
  map2 (\x y -> x * y) xs ys
  |> reduce (+) 0.0

def matmul [n][m][p] (a: [n][m]f64) (b: [m][p]f64): [n][p]f64 =
  map (\a_row ->
         map (\b_col -> dotprod a_row b_col)
             (transpose b))
      n
  • Second-order array combinators (map, scan, reduce).
  • Arrays must be regular.
  • Size types.

Futlixir is a NIF-generator

Link: https://github.com/Munksgaard/futlixir

$ futhark opencl --library foo.fut
$ ls
foo.c  foo.fut  foo.h  foo.json

Futlixir takes a compiled Futhark library and turns it into a NIF:

$ futlixir foo.json Foo
$ ls
...  foo.ex  foo_nif.c

Now we can compile the NIF to a shared object file:

$ gcc -Wall -shared -o foo_nif.so -fPIC foo_nif.c -lOpenCL -lm
$ ls
...  foo_nif.so

(Could be handled automatically by Futlixir)

And now we can use it, simply by including foo.ex in our Elixir programs!

In IEx you can use c("foo.ex"), in Livebook you can use Code.compile_file("foo.ex", "/absolute/path/to/parent/dir") (with some caveats).

I guess you can just dump the .ex and .so files in your application?

Livebook demo

What’s next for Futlixir?

  • Some exotic features haven’t been tested (opaque values)
  • Cut a release 0.1.0
  • After that, work on ergonomics:
    • Automatically compiling NIF shared objects
    • Automatically compile Futhark files?
    • Some sort of macro for writing Futhark code directly in Elixir!?
      • That’s a pipe dream for sure.
    • Using GenServers to handle contexts and arrays, to make sure they are used and freed correctly?
  • Also, I should probably do some actual work instead.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment