From 7b0cdd17e4659c8260ad9909c360a9604df57048 Mon Sep 17 00:00:00 2001
From: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Date: Wed, 5 Mar 2025 16:19:13 -0300
Subject: [PATCH 1/2] feat: allow explicitly disabling CUDA step

---
 exla/Makefile  | 16 +++++++++-------
 exla/README.md |  2 ++
 exla/mix.exs   |  8 ++++++++
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/exla/Makefile b/exla/Makefile
index 77b863dd7c..1842c1b728 100644
--- a/exla/Makefile
+++ b/exla/Makefile
@@ -37,26 +37,28 @@ else
 endif
 
 LDFLAGS = -L$(XLA_EXTENSION_LIB) -lxla_extension -shared
+NVCC := $(CXX)
+NVCCFLAGS = $(CFLAGS)
 
 ifeq ($(CROSSCOMPILE),)
 	# Interrogate the system for local compilation
 	UNAME_S = $(shell uname -s)
 
+ifeq ($(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)
diff --git a/exla/README.md b/exla/README.md
index 2ffe144ff7..7b3bd71761 100644
--- a/exla/README.md
+++ b/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=true` environment variable to disable custom CUDA functionality in EXLA.
 
 ### Usage with Nerves
 
diff --git a/exla/mix.exs b/exla/mix.exs
index 2d1980bfd5..90c1e290c8 100644
--- a/exla/mix.exs
+++ b/exla/mix.exs
@@ -33,7 +33,15 @@ defmodule EXLA.MixProject do
         priv_path = Path.join(Mix.Project.app_path(), "priv")
         cwd_relative_to_priv = relative_to(File.cwd!(), priv_path)
 
+        exla_cpu_only =
+          if System.get_env("EXLA_CPU_ONLY") in ["1", "true"] do
+            "true"
+          else
+            ""
+          end
+
         %{
+          "EXLA_CPU_ONLY" => exla_cpu_only,
           "FINE_INCLUDE_DIR" => Fine.include_dir(),
           "MIX_BUILD_EMBEDDED" => "#{Mix.Project.config()[:build_embedded]}",
           "CWD_RELATIVE_TO_PRIV_PATH" => cwd_relative_to_priv,

From 2d42cc38b908c704507e5ca44fd4fff2b53f53e2 Mon Sep 17 00:00:00 2001
From: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Date: Thu, 6 Mar 2025 15:18:21 -0300
Subject: [PATCH 2/2] chore: changes due to code rewview

---
 exla/Makefile  | 2 +-
 exla/README.md | 2 +-
 exla/mix.exs   | 8 --------
 3 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/exla/Makefile b/exla/Makefile
index 1842c1b728..9678fa78a2 100644
--- a/exla/Makefile
+++ b/exla/Makefile
@@ -44,7 +44,7 @@ ifeq ($(CROSSCOMPILE),)
 	# Interrogate the system for local compilation
 	UNAME_S = $(shell uname -s)
 
-ifeq ($(EXLA_CPU_ONLY),)
+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))
diff --git a/exla/README.md b/exla/README.md
index 7b3bd71761..d837b5ec73 100644
--- a/exla/README.md
+++ b/exla/README.md
@@ -44,7 +44,7 @@ EXLA relies on the [XLA](https://github.com/elixir-nx/xla) package to provide th
     * 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=true` environment variable to disable custom CUDA functionality in EXLA.
+    * 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
 
diff --git a/exla/mix.exs b/exla/mix.exs
index 90c1e290c8..2d1980bfd5 100644
--- a/exla/mix.exs
+++ b/exla/mix.exs
@@ -33,15 +33,7 @@ defmodule EXLA.MixProject do
         priv_path = Path.join(Mix.Project.app_path(), "priv")
         cwd_relative_to_priv = relative_to(File.cwd!(), priv_path)
 
-        exla_cpu_only =
-          if System.get_env("EXLA_CPU_ONLY") in ["1", "true"] do
-            "true"
-          else
-            ""
-          end
-
         %{
-          "EXLA_CPU_ONLY" => exla_cpu_only,
           "FINE_INCLUDE_DIR" => Fine.include_dir(),
           "MIX_BUILD_EMBEDDED" => "#{Mix.Project.config()[:build_embedded]}",
           "CWD_RELATIVE_TO_PRIV_PATH" => cwd_relative_to_priv,