feat: allow explicitly disabling CUDA step (#1588)

Author: polvalente

exla/Makefile

@@ -36,27 +36,29 @@ else
 	CFLAGS += -O3
 endif
 
+NVCC := $(CXX)
+NVCCFLAGS = $(CFLAGS)
 LDFLAGS = -L$(XLA_EXTENSION_LIB) -lxla_extension -shared -fvisibility=hidden
 
 ifeq ($(CROSSCOMPILE),)
 	# Interrogate the system for local compilation
 	UNAME_S = $(shell uname -s)
 
+ifdef ($(EXLA_CPU_ONLY),)
+$(info EXLA_CPU_ONLY is not set, checking for nvcc availability)
 	NVCC_RESULT := $(shell which nvcc 2> /dev/null)
 	NVCC_TEST := $(notdir $(NVCC_RESULT))
 
-ifeq ($(NVCC_TEST),nvcc)
-	NVCC := nvcc
-	NVCCFLAGS += -DCUDA_ENABLED
+	ifeq ($(NVCC_TEST),nvcc)
+		NVCC := nvcc
+		NVCCFLAGS += -DCUDA_ENABLED
+	endif
 else
-	NVCC := $(CXX)
-	NVCCFLAGS = $(CFLAGS)
+$(info EXLA_CPU_ONLY is set, skipping nvcc step)
 endif
 else
 	# Determine settings for cross-compiled builds like for Nerves
 	UNAME_S = Linux
-	NVCC := $(CXX)
-	NVCCFLAGS = $(CFLAGS)
 endif
 
 ifeq ($(UNAME_S), Darwin)

exla/README.md

@@ -43,6 +43,8 @@ EXLA relies on the [XLA](https://github.com/elixir-nx/xla) package to provide th
   * Incompatible protocol buffer versions
     * Error message: "this file was generated by an older version of protoc which is incompatible with your Protocol Buffer headers".
     * If you have `protoc` installed on your machine, it may conflict with the `protoc` precompiled inside XLA. Uninstall, unlink, or remove `protoc` from your path to continue.
+  * Missing CUDA symbols
+    * In some cases, you might be compiling a CPU-only version of `:xla` in an environment that has CUDA available. For these cases, you can set the `EXLA_CPU_ONLY` environment variable to any value to disable custom CUDA functionality in EXLA.
 
 ### Usage with Nerves