README.md

@@ -1,8 +0,0 @@
-
-Example cuda code compilation in nix.
-
-```bash
-$ nix run github:osbm/cuda-kernel-nix
-Free mem: 22112 MB Total mem: 24078 MB
-Free mem: 21 GB Total mem: 23 GB
-```

flake.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "nixpkgs": {
       "locked": {
-        "lastModified": 1756159630,
-        "narHash": "sha256-ohMvsjtSVdT/bruXf5ClBh8ZYXRmD4krmjKrXhEvwMg=",
+        "lastModified": 1734988233,
+        "narHash": "sha256-Ucfnxq1rF/GjNP3kTL+uTfgdoE9a3fxDftSfeLIS8mA=",
         "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "84c256e42600cb0fdf25763b48d28df2f25a0c8b",
+        "rev": "de1864217bfa9b5845f465e771e0ecb48b30e02d",
         "type": "github"
       },
       "original": {

flake.nix

@@ -29,40 +29,21 @@
     packages."${system}" = {
       # just one main.cu file for now
       # compile it and place it to $out/bin/cuda-kernel
-      cuda-kernel = let
-        # Use CUDA 12.0 which should be compatible with driver 580.76.05
-        cudaPackages = pkgs.cudaPackages;
-        cudatoolkit = cudaPackages.cudatoolkit;
-      in
-        cudaPackages.backendStdenv.mkDerivation {
-          name = "cuda";
-          src = ./main.cu;
-          unpackPhase = "true";
-          buildInputs = [
-            cudatoolkit
-            cudaPackages.cuda_cudart
-          ];
-          nativeBuildInputs = [
-            cudatoolkit
-            pkgs.autoAddDriverRunpath
-          ];
-          LD_LIBRARY_PATH = "${cudatoolkit}/lib:${cudaPackages.cuda_cudart}/lib";
-          buildPhase = ''
-            echo "GCC Version:"
-            gcc --version
-            echo "NVCC Version:"
-            nvcc --version
-
-            echo "Building CUDA program..."
-            nvcc -arch=sm_89 -o main $src
-            mkdir -p $out/bin
-            cp main $out/bin/
-          '';
-          meta = {
-            mainProgram = "main";
-            platforms = ["x86_64-linux"];
-          };
-        };
+      cuda-kernel = pkgs.stdenv.mkDerivation {
+        name = "cuda-kernel";
+        src = ./.;
+        buildInputs = with pkgs; [
+          cudatoolkit
+          cudaPackages.cuda_cudart
+          linuxPackages.nvidia_x11
+        ];
+        LD_LIBRARY_PATH = "${pkgs.linuxPackages.nvidia_x11}/lib";
+        buildPhase = ''
+          nvcc -o main main.cu
+          mkdir -p $out/bin
+          mv main $out/bin/cuda-kernel
+        '';
+      };
       default = self.packages."${system}".cuda-kernel;
     };
   };

main.cu

@@ -1,13 +1,13 @@
+#include <unistd.h>
 #include <iostream>
-#include "cuda.h"
 
-int main() {
-  size_t free_t, total_t;
-  cudaError_t err;
-  err = cudaMemGetInfo(&free_t, &total_t);
-  if (err != cudaSuccess) { printf("cudaMalloc d_steps_arr error: %s\n", cudaGetErrorString(err)); return 1; }
-  std::cout << "Free mem: " << free_t / (1024 * 1024) << " MB Total mem: " << total_t / (1024 * 1024) << " MB" << std::endl;
-  std::cout << "Free mem: " << free_t / (1024 * 1024 * 1024) << " GB Total mem: " << total_t / (1024 * 1024 * 1024) << " GB" << std::endl;
-  return 0;
+
+int main(){
+
+    size_t mf, ma;
+    cudaMemGetInfo(&mf, &ma);
+    std::cout << "Free memory (mb): " << mf/1024/1024 << std::endl;
+    std::cout << "Total memory (mb): " << ma/1024/1024 << std::endl;
+
+    return 0;
 }
-