Merge build-cuda and build-linux actions

.github/actions/build-cuda/action.yml

@@ -1,26 +0,0 @@
-name: 'Build and Test with CUDA'
-description: 'Build and test MLX with CUDA'
-
-inputs:
-  toolkit:
-    description: 'The CUDA toolkit'
-    required: true
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Python package
-      shell: bash
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
-      run: pip install --no-build-isolation -e ".[dev]" -v
-
-    - name: Build CPP only
-      shell: bash
-      run: |
-        cmake . -B build \
-          -DMLX_BUILD_CUDA=ON \
-          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc \
-          -DCMAKE_BUILD_TYPE=DEBUG
-        cmake --build build -j $(nproc)

.github/actions/build-linux/action.yml

@@ -1,13 +1,30 @@
 name: 'Build and Test on Linux'
-description: 'Build and test MLX on Linux'
+
+inputs:
+  toolkit:
+    description: 'The toolkit to build with'
+    required: false
+    default: 'cpu'
 
 runs:
   using: "composite"
   steps:
+    - name: Set cmake args
+      shell: bash
+      env:
+        CMAKE_ARGS_CPU: >
+          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
+          -DMLX_BUILD_CUDA=OFF
+        CMAKE_ARGS_CUDA: >
+          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
+          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
+          -DMLX_BUILD_CUDA=ON
+      run: |
+        echo CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}" > $GITHUB_ENV
+
     - name: Install Python package
       shell: sh
       env:
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
         DEBUG: 1
       run: pip install --no-build-isolation -e ".[dev]" -v
 
@@ -20,6 +37,5 @@ runs:
     - name: Build CPP only
       shell: bash
       run: |
-        mkdir -p build && cd build
+        cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ env.CMAKE_ARGS }}
-        cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
+        cmake --build build -j $(nproc)
-        make -j $(nproc)

.github/workflows/pull_request.yml

@@ -66,7 +66,7 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: ${{ matrix.toolkit }}
-      - uses: ./.github/actions/build-cuda
+      - uses: ./.github/actions/build-linux
         with:
           toolkit: ${{ matrix.toolkit }}
       - uses: ./.github/actions/test-linux