small

jinluchang · Jan 12, 2025 · f4b3acb · f4b3acb
1 parent 6fd4859
commit f4b3acb
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/nixpkgs/build-many-qlat-pkgs.sh b/nixpkgs/build-many-qlat-pkgs.sh
@@ -29,6 +29,20 @@ for name in \
     "-pypi" \
     ; do
     time nix-build "$src"/qlat-pkgs.nix -A qlat-jhub-tests"$name" -o result-24-11"$name" --arg nixpkgs 'import (fetchTarball "https://channels.nixos.org/nixos-24.11/nixexprs.tar.xz")' "$@"
+done
+for name in \
+    "" \
+    "-ucxless" \
+    "-std-clang" \
+    "-std-cubaquadless" \
+    "-cuda" \
+    "-cuda-ucxless" \
+    "-std" \
+    "-std-ucxless" \
+    "-std-cuda" \
+    "-std-cuda-ucxless" \
+    "-pypi" \
+    ; do
     time nix-build "$src"/qlat-pkgs.nix -A qlat-jhub-tests"$name" -o result-24-05"$name" --arg nixpkgs 'import (fetchTarball "https://channels.nixos.org/nixos-24.05/nixexprs.tar.xz")' "$@"
 done
 )
diff --git a/nixpkgs/qlat-pkgs.nix b/nixpkgs/qlat-pkgs.nix
@@ -2,6 +2,8 @@
   nixpkgs ? import ./nixpkgs.nix,
   ngpu ? "2", # adjust with actual number of GPUs
   nvcc-arch ? "sm_86", # adjust with actual arch of Nvidia GPU
+  cudaCapability ? "8.6",
+  cudaForwardCompat ? false,
 }:
 
 let
@@ -26,9 +28,13 @@ let
     use-clang = false;
     use-ucx = true;
     use-pypi = false;
+    cudaCapabilities = [ "8.6" ];
+    cudaForwardCompat = false;
   } // {
     ngpu = ngpu;
     nvcc-arch = nvcc-arch;
+    cudaCapabilities = [ cudaCapability ];
+    cudaForwardCompat = cudaForwardCompat;
   };
 
   mk-qlat-name = options:
@@ -409,6 +415,8 @@ let
       config = {
         allowUnfree = opts.use-cuda;
         cudaSupport = opts.use-cudasupport;
+        ${if opts.use-cudasupport then "cudaCapabilities" else null} = opts.cudaCapabilities;
+        ${if opts.use-cudasupport then "cudaForwardCompat" else null} = opts.cudaForwardCompat;
       };
       overlays = [
         (mk-overlay options)