Skip to content

Commit

Permalink
fix: get_tuple NIF helper + compilation configs and speedup (#1479)
Browse files Browse the repository at this point in the history
  • Loading branch information
polvalente authored May 8, 2024
1 parent 0cb116c commit fba9efe
Show file tree
Hide file tree
Showing 4 changed files with 11 additions and 5 deletions.
2 changes: 1 addition & 1 deletion exla/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ ifeq ($(NVCC_TEST),nvcc)
NVCC := nvcc
NVCCFLAGS += -DCUDA_ENABLED
else
NVCC := g++
NVCC := $(CXX)
NVCCFLAGS = $(CFLAGS)
endif

Expand Down
2 changes: 2 additions & 0 deletions exla/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ mix deps.get
mix test
```

By default, EXLA passes `["-jN"]` as a Make argument, where `N` is `System.schedulers_online() - 2`, capped at `1`. `config :exla, :make_args, ...` can be used to override this default setting.

In order to run tests on a specific device, use the `EXLA_TARGET` environment variable, which is a dev-only variable for this project (it has no effect when using EXLA as a dependency). For example, `EXLA_TARGET=cuda` or `EXLA_TARGET=rocm`. Make sure to also specify `XLA_TARGET` to fetch or compile a proper version of the XLA binary.

### Building with Docker
Expand Down
6 changes: 3 additions & 3 deletions exla/c_src/exla/exla_nif_util.cc
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
#include "exla_nif_util.h"

#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "stablehlo/dialect/StablehloOps.h"
#include "xla/primitive_util.h"
#include "xla/shape_util.h"
#include "mlir/IR/Builders.h"
#include "stablehlo/dialect/StablehloOps.h"

namespace exla {
namespace nif {
Expand Down Expand Up @@ -190,7 +190,7 @@ int get_tuple(ErlNifEnv* env, ERL_NIF_TERM tuple, std::vector<int64>& var) {
var.reserve(length);

for (int i = 0; i < length; i++) {
int data;
int64 data;
if (!get(env, terms[i], &data)) return 0;
var.push_back(data);
}
Expand Down
6 changes: 5 additions & 1 deletion exla/mix.exs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@ defmodule EXLA.MixProject do
@version "0.7.1"

def project do
make_args =
Application.get_env(:exla, :make_args) || ["-j#{max(System.schedulers_online() - 2, 1)}"]

[
app: :exla,
version: @version,
Expand Down Expand Up @@ -34,7 +37,8 @@ defmodule EXLA.MixProject do
"MIX_BUILD_EMBEDDED" => "#{Mix.Project.config()[:build_embedded]}",
"CWD_RELATIVE_TO_PRIV_PATH" => cwd_relative_to_priv
}
end
end,
make_args: make_args
]
end

Expand Down

0 comments on commit fba9efe

Please sign in to comment.