Credit: https://mlsdev.com/blog/elixir-programming-facts-to-know-for-better-app-development
Credit: https://mlsdev.com/blog/elixir-programming-facts-to-know-for-better-app-development

Supercharge Your Elixir with Native Implemented Functions (NIFs)

NIFs are a great way to implement fast native code in Elixir, and open the world of external APIs to Elixir developers

I have recently discovered the joy of writing Elixir code. For those who don't know, Elixir is a functional programming language built to run on the same virtual machine that powers Erlang, known as the BEAM. Elixir is lauded for its productivity and development in the domains of web development, embedded software, machine learning, data pipelines, and multimedia processing, just to name a few.  

As I was looking for ways to contribute to the open-source community, a friend of mine suggested that I look into writing NIFs for Elixir to leverage existing libraries written in other languages, such as C, C++, Rust, or Zig. As I was looking for resources about writing NIFs, I found that much of the material only gave a very shallow understanding of what NIFs could do and did not go into depth regarding how to actually use them to make anything interesting in Elixir. My goal with this article is to not only thoroughly explain what NIFs are, but give some tips for working with NIFs as well as show some real-world examples of using NIFs in a library I decided to write which brings the power of the XGBoost library into Elixir

Introduction to NIFs

Native Implemented Functions (NIFs) are functions implemented in code that compiles natively to the target machine and are executed outside of the confines of the BEAM virtual machine.  The :erlang.load_nif/2 functions allow the use of these functions from Elixir (or Erlang for that matter). Some possible use cases of NIFs are to speed up extremely time-sensitive operations, have better control of hardware, or interact with existing APIs written in other languages as we will explore more in-depth later on.

The basic library that you will interface with when writing NIFs (in C/C++ at least – more on this later) is the erl_nif.h C library. The library has fairly comprehensive documentation which you will find yourself referencing quite frequently when writing NIFs, but I would like to go over some of the essential components that you will be interacting with the most.  

In Erlang/Elixir, the simplest form of expression is a term, that is an integer, float, atom, string, list, map, or tuple. The return value is the term itself.  In erl_nif.h, many of the API calls take the form of enif_get_* or enif_make_*, which denotes that the function is either ingesting a term from the Elixir side and converting it to the respective type in C, or taking a C variable of said type and converting it to an Elixir term of the appropriate type.  When converting to an Elixir term, the term and any allocated memory are also registered to the BEAM garbage handler, and at that point is no longer the responsibility of the native code with regards to memory management.

Additionally, erl_nif.h also interacts directly with Elixir binaries, which are represented as an opaque struct where the only two fields you interact with are data, which points to the contiguous block of data where the binary resides, and size, which stores the lengths of the binary.  Binaries are a great way to pass data between the BEAM and the NIF without altering the data.

The last key idea that I will discuss here is the concept of a resource object. As the erl_nif.h documentation states, "The use of resource objects is a safe way to return pointers to native data structures from a NIF."  The main usage of a resource object is when you have a data structure in the native code that cannot be converted to an Elixir term, then you can still pass a reference to that structure to Elixir. That resource will then only be used to pass between different Elixir functions which are mapped to other NIFs. So, from the Elixir side, the resource is represented as a reference and is opaque, and must be passed back to another NIF that acts upon it to do anything useful with it.

NIFs in Practice

Now, let's look at some practical examples of what can be done with NIFs. For these examples, I will be using code from EXGBoost, which is an Elixir library that I wrote to leverage the use of the XGBoost C API in Elixir. The first thing to know is that to use a NIF, your native code implementations must be compiled into a shared library for the appropriate target architecture ( .so for Linux, .dll for Windows, .dylib for MacOS). For EXGBoost, that is named called libexgboost.  Then, in the Elixir module where you want your NIFs to be called from, you will have the following:

defmodule EXGBoost.NIF do
  @on_load :on_load
  def on_load do
    path = :filename.join([:code.priv_dir(:exgboost), "libexgboost"])
    :erlang.load_nif(path, 0)
  end
end
exgboost/lib/exgboost/nif.ex

This is the minimum needed to load a NIF into Elixir, assuming that a shared library named libexgboost exists in the modules priv directory.  

Basic Example

Now loading a NIF library is not very useful if there are no functions to run, so let's add some basic functions, starting with simply getting the build information of XGBoost that our NIFs are linked against. First, we must ensure that our NIF library (libexgboost) is linked against the XGBoost shared library ( libxgboost ). You can refer to the Makefile to see how that is done – I won't go into details about it here. Let's start by making a utility file that will import relevant libraries into our project, and then making a config.{h,c} to hold the function we are making EXGBoostVersion :

#ifndef EXGBOOST_UTILS_H
#define EXGBOOST_UTILS_H

#include <erl_nif.h>
#include <stdio.h>
#include <string.h>
#include <xgboost/c_api.h>

// Status helpers

ERL_NIF_TERM exg_error(ErlNifEnv *env, const char *msg);

ERL_NIF_TERM ok_atom(ErlNifEnv *env);

ERL_NIF_TERM exg_ok(ErlNifEnv *env, ERL_NIF_TERM term);

#endif

Here, we are importing erl_nif.h which we will need to use the Erlang NIF API functions, and xgboost/c_api.h, which will be used to interact with libxgboost. The interaction between these two libraries is the crux of our library.

#include "utils.h"

// Atoms
ERL_NIF_TERM exg_error(ErlNifEnv *env, const char *msg) {
  ERL_NIF_TERM atom = enif_make_atom(env, "error");
  ERL_NIF_TERM msg_term = enif_make_string(env, msg, ERL_NIF_LATIN1);
  return enif_make_tuple2(env, atom, msg_term);
}

ERL_NIF_TERM ok_atom(ErlNifEnv *env) { return enif_make_atom(env, "ok"); }

ERL_NIF_TERM exg_ok(ErlNifEnv *env, ERL_NIF_TERM term) {
  return enif_make_tuple2(env, ok_atom(env), term);
}

Here, we are setting up some helper functions. Every NIF will return in the form of {:ok, term} | {:error, String.t()} | :ok, which means that one of these three functions will be returned for every NIF.

#ifndef EXGBOOST_CONFIG_H
#define EXGBOOST_CONFIG_H

#include "utils.h"

ERL_NIF_TERM EXGBuildInfo(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);

#endif
#include "config.h"

ERL_NIF_TERM EXGBuildInfo(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
  char const *out = NULL;
  int result = -1;
  ERL_NIF_TERM ret = 0;
  if (argc != 0) {
    ret = exg_error(env, "Wrong number of arguments");
    goto END;
  }
  result = XGBuildInfo(&out);
  if (result == 0) {
    ret = exg_ok(env, enif_make_string(env, out, ERL_NIF_LATIN1));
  } else {
    ret = exg_error(env, XGBGetLastError());
  }
END:
  return ret;
}

Here is the implementation of the actual NIF. Each NIF will have the function signature of ERL_NIF_TERM (*fptr)(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]), where env represents an environment that can host Erlang terms – all terms in an environment are valid as long as the environment is valid, argc is the number of arguments passed to the NIF, and argv is an array of terms passed to the NIF. So, all this function is doing is checking to make sure the correct number of arguments were passed (in this case zero), then invoking the XGBoost API to get the build info (which is returned from the XGBoost API as a JSON-encoded string), and finally either returning it as an Elixir string on success or returning an error on failure.

Now that the NIF implementation is done, it is time to initialize it.  We do this by using the ERL_NIF_INIT function and passing it an array of ErlNifFunc. Each ErlNifFunc defines what the corresponding Elixir function is called, the arity of the function, the native function that is bound to the Elixir function, and a flag variable that can be used to change how the scheduler handles the NIF. In addition the functions array, you also pass the name of the Elixir module that will house your NIFs,  as well as three optional functions load, upgrade, and unload that we will talk about later.  

#ifndef EXGBOOST_H
#define EXGBOOST_H

#include "config.h"

#endif
#include "exgboost.h"

static ErlNifFunc nif_funcs[] = {
	{"xgboost_build_info", 0, EXGBuildInfo}
};
ERL_NIF_INIT(Elixir.EXGBoost.NIF, nif_funcs, load, NULL, upgrade, NULL)

In this example, this will create a function in the Elixir NIF module called xgboost_build_info/0 which, when called, will pass all arguments (in this case there are none) to the native function, run the native function, and return from the native function back to Elixir.  

defmodule EXGBoost.NIF do
  @on_load :on_load
  def on_load do
    path = :filename.join([:code.priv_dir(:exgboost), "libexgboost"])
    :erlang.load_nif(path, 0)
  end
  
  def xgboost_build_info, do: :erlang.nif_error(:not_implemented)
end

Now, we can run the NIF by doing EXGBoost.NIF.xgboost_build_info(). Congratulations! You have now used Elixir NIFs to run an external API. Next, let's look at a more advanced example.

Using Resource Objects

For the advanced example, let's fast forward to performing a prediction using XGBoost. The two main data structures used in XGBoost are DMatrix which represents the data matrix holding the input data to the model, and Booster which represents the model itself. These two structures will be initialized and created from the XGBoost API, but we need to allow the Elixir NIF module to interact with them in some way. This is where resource objects come in. We use resources as a handle to the DMatrix and Booster structs that we can pass back and forth between NIFs. There are a few things that must be done for each resource type. We must declare each unique resource type, we must register with ERL_NIF_INIT what resource types we will be using, and since the XGBoost API documentation tells us to use the custom XGDMatrixFree and XGBoosterFree functions to free their respective structs, we must also tell the BEAM garbage handler how to properly dispose of the resources.

#ifndef EXGBOOST_UTILS_H
#define EXGBOOST_UTILS_H

#include <erl_nif.h>
#include <xgboost/c_api.h>

ErlNifResourceType *DMatrix_RESOURCE_TYPE;
ErlNifResourceType *Booster_RESOURCE_TYPE;

void DMatrix_RESOURCE_TYPE_cleanup(ErlNifEnv *env, void *arg);
void Booster_RESOURCE_TYPE_cleanup(ErlNifEnv *env, void *arg);

#endif

Here we declare both resource types as well as their cleanup functions.

#include "utils.h"

// Resource type helpers
void DMatrix_RESOURCE_TYPE_cleanup(ErlNifEnv *env, void *arg) {
  DMatrixHandle handle = *((DMatrixHandle *)arg);
  XGDMatrixFree(handle);
}

void Booster_RESOURCE_TYPE_cleanup(ErlNifEnv *env, void *arg) {
  BoosterHandle handle = *((BoosterHandle *)arg);
  XGBoosterFree(handle);
}

Here we define the cleanup functions to use the custom freeing functions that the XGBoost documentation tells us to use.

#include "exgboost.h"

static int load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info) {
  DMatrix_RESOURCE_TYPE = enif_open_resource_type(
      env, NULL, "DMatrix_RESOURCE_TYPE", DMatrix_RESOURCE_TYPE_cleanup,
      (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER), NULL);
  Booster_RESOURCE_TYPE = enif_open_resource_type(
      env, NULL, "Booster_RESOURCE_TYPE", Booster_RESOURCE_TYPE_cleanup,
      (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER), NULL);
  if (DMatrix_RESOURCE_TYPE == NULL || Booster_RESOURCE_TYPE == NULL) {
    return 1;
  }
  return 0;
}

static int upgrade(ErlNifEnv *env, void **priv_data, void **old_priv_data,
                   ERL_NIF_TERM load_info) {
  DMatrix_RESOURCE_TYPE = enif_open_resource_type(
      env, NULL, "DMatrix_RESOURCE_TYPE", DMatrix_RESOURCE_TYPE_cleanup,
      ERL_NIF_RT_TAKEOVER, NULL);
  Booster_RESOURCE_TYPE = enif_open_resource_type(
      env, NULL, "Booster_RESOURCE_TYPE", Booster_RESOURCE_TYPE_cleanup,
      ERL_NIF_RT_TAKEOVER, NULL);
  if (DMatrix_RESOURCE_TYPE == NULL || Booster_RESOURCE_TYPE == NULL) {
    return 1;
  }
  return 0;
}
static ErlNifFunc nif_funcs[] = {
	{"xgboost_build_info", 0, EXGBuildInfo}
};
ERL_NIF_INIT(Elixir.EXGBoost.NIF, nif_funcs, load, NULL, upgrade, NULL)

Here is where we register to ERL_NIF_INIT that we are using these two resource types, and which cleanup function to use for each. One of load or upgrade is called to initialize the library. unload (which is null in this case) is called to release the library. By passing the appropriate cleanup function to enif_open_resource, we are registering which freeing function to use when the appropriate resource is discarded by the garbage collector.

Now we are ready to use resource types to make resources from the DMatrix and Booster structs. I like to make helper functions that take in the BoosterHandle and DMatrixHandle types and returns the resource or error, so let's go ahead and make those functions:

#include "booster.h"

static ERL_NIF_TERM make_Booster_resource(ErlNifEnv *env,
                                          BoosterHandle handle) {
  ERL_NIF_TERM ret = -1;
  BoosterHandle **resource =
      enif_alloc_resource(Booster_RESOURCE_TYPE, sizeof(BoosterHandle *));
  if (resource != NULL) {
    *resource = handle;
    ret = exg_ok(env, enif_make_resource(env, resource));
    enif_release_resource(resource);
  } else {
    ret = exg_error(env, "Failed to allocate memory for XGBoost DMatrix");
  }
  return ret;
}
#include "dmatrix.h"

static ERL_NIF_TERM make_DMatrix_resource(ErlNifEnv *env,
                                          DMatrixHandle handle) {
  ERL_NIF_TERM ret = -1;
  DMatrixHandle **resource =
      enif_alloc_resource(DMatrix_RESOURCE_TYPE, sizeof(DMatrixHandle *));
  if (resource != NULL) {
    *resource = handle;
    ret = exg_ok(env, enif_make_resource(env, resource));
    enif_release_resource(resource);
  } else {
    ret = exg_error(env, "Failed to allocate memory for XGBoost DMatrix");
  }
  return ret;
}

Great! Now we can just use either of these functions when we need to return the appropriate resource to the Elixir NIF module. So, let's make the NIF that will allow us to create a new Booster from Elixir. The NIF will accept an Elixir list of DMatrix resources to initialize the Booster from, and if the list is empty it creates a Booster that's not associated with any DMatrix.  First, here is a helper function that takes in a list of DMatrix resource objects that are passed from Elixir and populates the array dmats with all of the corresponding DMatrixHandle structs:

int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term,
                         DMatrixHandle **dmats, unsigned *len) {
  ERL_NIF_TERM head, tail;
  int i = 0;
  if (!enif_get_list_length(env, term, len)) {
    return 0;
  }
  *dmats = (DMatrixHandle *)enif_alloc(*len * sizeof(DMatrixHandle));
  if (NULL == dmats) {
    return 0;
  }
  while (enif_get_list_cell(env, term, &head, &tail)) {
    DMatrixHandle **resource = NULL;
    if (!enif_get_resource(env, head, DMatrix_RESOURCE_TYPE,
                           (void *)&(resource))) {
      return 0;
    }
    memcpy(&((*dmats)[i]), resource, sizeof(DMatrixHandle));
    term = tail;
    i++;
  }
  return 1;
}

Next, we add the NIF implementation:

#include "booster.h"

ERL_NIF_TERM EXGBoosterCreate(ErlNifEnv *env, int argc,
                              const ERL_NIF_TERM argv[]) {
  DMatrixHandle *dmats = NULL;
  ERL_NIF_TERM ret = -1;
  int result = -1;
  unsigned dmats_len = 0;
  BoosterHandle booster = NULL;
  if (1 != argc) {
    ret = exg_error(env, "Wrong number of arguments");
    goto END;
  }
  if (!exg_get_dmatrix_list(env, argv[0], &dmats, &dmats_len)) {
    ret = exg_error(env, "Invalid list of DMatrix");
    goto END;
  }
  if (0 == dmats_len) {
    result = XGBoosterCreate(NULL, 0, &booster);
    if (result == 0) {
      ret = make_Booster_resource(env, booster);
      goto END;
    } else {
      ret = exg_error(env, "Error making booster");
      goto END;
    }
  }

  result = XGBoosterCreate(dmats, dmats_len, &booster);
  if (result == 0) {
    ret = make_Booster_resource(env, booster);
    goto END;
  } else {
    ret = exg_error(env, XGBGetLastError());
  }
END:
  return ret;
}

This implementation accepts a list of DMatrix resources, uses them to create a new BoosterHandle struct, then creates a resource for the Booster that is then passed back to the calling Elixir module. Now, we register the new NIF to its corresponding Elixir function:

static ErlNifFunc nif_funcs[] = {
	{"xgboost_build_info", 0, EXGBuildInfo},
    {"booster_create", 1, EXGBoosterCreate},
};

Finally, we add the NIF to our Elixir module. Now, we can create a new Booster from Elixir using EXGBoost.NIF.booster_create/1.

defmodule EXGBoost.NIF do
  @on_load :on_load
  def on_load do
    path = :filename.join([:code.priv_dir(:exgboost), "libexgboost"])
    :erlang.load_nif(path, 0)
  end
  
  def xgboost_build_info, do: :erlang.nif_error(:not_implemented)
  def booster_create(_handles), do: :erlang.nif_error(:not_implemented)
end

Advanced Example

Now that we can pass the DMatrix and Booster structs between NIFs, let's pretend we have a trained Booster and want to make predictions using it. So let's go ahead and implement the EXGBoosterPredictFromDMatrix NIF. This will accept a Booster resource, a DMatrix resource, and a JSON-encoded config string and return a 2-tuple of the shape of the prediction results and the flat array of the prediction results.

ERL_NIF_TERM EXGBoosterPredictFromDMatrix(ErlNifEnv *env, int argc,
                                          const ERL_NIF_TERM argv[]) {
  BoosterHandle booster;
  BoosterHandle **booster_resource = NULL;
  DMatrixHandle dmatrix;
  DMatrixHandle **dmatrix_resource = NULL;
  char *config = NULL;
  bst_ulong *out_shape = NULL;
  bst_ulong out_dim = 0;
  float *out_result = NULL;

  ERL_NIF_TERM ret = -1;
  int result = -1;
  if (3 != argc) {
    ret = exg_error(env, "Wrong number of arguments");
    goto END;
  }
  if (!enif_get_resource(env, argv[0], Booster_RESOURCE_TYPE,
                         (void *)&(booster_resource))) {
    ret = exg_error(env, "Invalid Booster");
    goto END;
  }
  if (!enif_get_resource(env, argv[1], DMatrix_RESOURCE_TYPE,
                         (void *)&(dmatrix_resource))) {
    ret = exg_error(env, "Invalid DMatrix");
    goto END;
  }
  if (!exg_get_string(env, argv[2], &config)) {
    ret = exg_error(env, "Config must be a JSON-encoded string");
    goto END;
  }
  booster = *booster_resource;
  dmatrix = *dmatrix_resource;
  result = XGBoosterPredictFromDMatrix(booster, dmatrix, config, &out_shape,
                                       &out_dim, &out_result);
  if (result == 0) {
    ret = collect_prediction_results(env, out_shape, out_dim, out_result);
  } else {
    ret = exg_error(env, XGBGetLastError());
  }
END:
  if (config != NULL) {
    enif_free(config);
  }
  return ret;
}

Here, you can see that we take in the two resources, pass their underlying structs to the XGBoost XGBoosterPredictFromDMatrix API call, then use collect_prediction_results to return the desired output, so let's take a look at collect_prediction_results:

static ERL_NIF_TERM collect_prediction_results(ErlNifEnv *env,
                                               bst_ulong *out_shape,
                                               bst_ulong out_dim,
                                               float *out_result) {
  bst_ulong out_len = 1;
  ERL_NIF_TERM shape_arr[out_dim];
  for (bst_ulong j = 0; j < out_dim; ++j) {
    shape_arr[j] = enif_make_int(env, out_shape[j]);
    out_len *= out_shape[j];
  }
  ERL_NIF_TERM shape = enif_make_tuple_from_array(env, shape_arr, out_dim);
  ERL_NIF_TERM result_arr[out_len];
  for (bst_ulong i = 0; i < out_len; ++i) {
    result_arr[i] = enif_make_double(env, out_result[i]);
  }
  return exg_ok(env, enif_make_tuple2(
                         env, shape,
                         enif_make_list_from_array(env, result_arr, out_len)));
}

This function will convert the shape to a tuple, convert the output predictions to a list, and return a 2-tuple containing both.

After registering this function to both the C and Elixir sides like we did before, we can use the function. First, we can use Elixir structs to wrap the resource neatly so that instead of interacting with a reference() type (which is the Elixir type of the resource), we can use the Booster and DMatrix structs:

defmodule EXGBoost.Booster do
	@enforce_keys [:ref]
    defstruct [:ref, :best_iteration, :best_score]
end
defmodule EXGBoost.DMatrix do
	@enforce_keys [
    :ref,
    :format
  ]
  defstruct [
    :ref,
    :format,
    :label,
    :weight,
    :base_margin,
    :group,
    :label_upper_bound,
    :label_lower_bound,
    :feature_weights,
    :missing,
    :nthread,
    :feature_names,
    :feature_types
  ]
end

We can use @enforce_keys to require that the resource be passed to the struct. Now, we can use the .ref key in each struct to pass their corresponding resources:

def predict(%Booster{} = booster, %DMatrix{} = data, opts \\ []) do
    ...
    # Refer to source for full example
    # https://github.com/acalejos/exgboost/blob/main/lib/exgboost/booster.ex
    {shape, preds} =
      EXGBoost.NIF.booster_predict_from_dmatrix(booster.ref, data.ref, Jason.encode!(config))
      |> Internal.unwrap!()

    Nx.tensor(preds) |> Nx.reshape(shape)
  end

Just like that, you can make predictions on a trained Booster using XGBoost from Elixir!

NIFs in the Wild

As I alluded to before, there are other languages that you can implement Elixir NIFs in. One of the downsides of writing NIFs (which you will become familiar with if you write them, and is heavily caveated on the erl_nif.h documentation) is that you open your program up to very unsafe code, where the BEAM cannot protect you. This means gasp SEGFAULTS can happen in your program. Or more insidious even -- memory leaks. It is very important to write your native code diligently and to be cognizant of the risks that NIFs incur, so as to only use them when and where appropriate.

Fortunately, the open-source community is to the rescue (as it often is), with projects such as Rustler and Zigler aiming to bring the power of NIFs to Elixir using much safer languages (Rust and Zig respectively).  Although I have not used either of these projects myself, I would encourage you to try to use these when possible rather than using the erl_nif.h C API directly.

If you're looking for some projects in the wild that use NIFs, I think you would be surprised to see how many well-known projects use them under the hood, but here are just a few examples:

Conclusion

NIFs are a great way to implement fast native code in Elixir, and open the world of external APIs to Elixir developers, but as with writing anything in C – with great power comes great responsibility. Some open-source alternatives have sprouted up to help address these problems, and as they become more mature, there might be increasingly fewer reasons to use C NIFs, although I doubt they will ever go extinct.  Comment below with any other interesting projects that use NIFs!

Comments