Run CPP tests for CUDA build in CI

.circleci/config.yml

@@ -232,6 +232,7 @@ jobs:
           name: Install Python package
           command: |
             uv venv
+            uv pip install cmake
             DEBUG=1 CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
               uv pip install -e ".[dev]" -v
       - run:
@@ -240,6 +241,18 @@ jobs:
             source .venv/bin/activate
             LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
             LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
+      - run:
+          name: Build CPP only
+          command: |
+            source .venv/bin/activate
+            cmake . -B build \
+              -DMLX_BUILD_CUDA=ON \
+              -DCMAKE_CUDA_COMPILER=`which nvcc` \
+              -DCMAKE_BUILD_TYPE=DEBUG
+            cmake --build build -j `nproc`
+      - run:
+          name: Run CPP tests
+          command: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
       - run:
           name: CCache report
           command: |

CMakeLists.txt

@@ -140,6 +140,12 @@ elseif(MLX_BUILD_METAL)
   target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
 endif()
 
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  # With newer clang/gcc versions following libs are implicitly linked, but when
+  # building on old distributions they need to be explicitly listed.
+  target_link_libraries(mlx PRIVATE dl pthread)
+endif()
+
 if(WIN32)
   if(MSVC)
     # GGUF does not build with MSVC.

mlx/backend/cuda/jit_module.cpp

@@ -67,9 +67,11 @@ const std::string& cccl_dir() {
       return path.string();
     }
     // Finally check the environment variable.
-    path = std::getenv("MLX_CCCL_DIR");
-    if (!path.empty() && std::filesystem::exists(path)) {
-      return path.string();
+    if (const char* env = std::getenv("MLX_CCCL_DIR"); env) {
+      path = env;
+      if (!path.empty() && std::filesystem::exists(path)) {
+        return path.string();
+      }
     }
     return std::string();
   }();

mlx/mlx.h

@@ -4,6 +4,7 @@
 
 #include "mlx/array.h"
 #include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/gpu/available.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/compile.h"
 #include "mlx/device.h"

tests/device_tests.cpp

@@ -10,7 +10,7 @@ using namespace mlx::core;
 
 TEST_CASE("test device placement") {
   auto device = default_device();
-  Device d = metal::is_available() ? Device::gpu : Device::cpu;
+  Device d = gpu::is_available() ? Device::gpu : Device::cpu;
   if (std::getenv("DEVICE") == nullptr) {
     CHECK_EQ(device, d);
   }
@@ -18,7 +18,7 @@ TEST_CASE("test device placement") {
   array x(1.0f);
   array y(1.0f);
   auto z = add(x, y, default_device());
-  if (metal::is_available()) {
+  if (gpu::is_available()) {
     z = add(x, y, Device::gpu);
     z = add(x, y, Device(Device::gpu, 0));
   } else {

tests/scheduler_tests.cpp

@@ -16,7 +16,7 @@ TEST_CASE("test stream management") {
   CHECK_NE(s1, s2);
 
   // Check that default streams have the correct devices
-  if (metal::is_available()) {
+  if (gpu::is_available()) {
     auto s_gpu = default_stream(Device::gpu);
     CHECK_EQ(s_gpu.device, Device::gpu);
   } else {
@@ -28,7 +28,7 @@ TEST_CASE("test stream management") {
   s_cpu = new_stream(Device::cpu);
   CHECK_EQ(s_cpu.device, Device::cpu);
 
-  if (metal::is_available()) {
+  if (gpu::is_available()) {
     auto s_gpu = new_stream(Device::gpu);
     CHECK_EQ(s_gpu.device, Device::gpu);
   } else {