push sheeet

pkgs/development/cuda-modules/tests/flags.nix

@@ -0,0 +1,79 @@
+{
+  _cuda,
+  cudaNamePrefix,
+  lib,
+  runCommand,
+}:
+let
+  inherit (builtins) deepSeq toJSON tryEval;
+  inherit (_cuda.bootstrapData) cudaCapabilityToInfo;
+  inherit (_cuda.lib) formatCapabilities;
+  inherit (lib.asserts) assertMsg;
+in
+# When changing names or formats: pause, validate, and update the assert
+assert assertMsg (
+  cudaCapabilityToInfo ? "7.5" && cudaCapabilityToInfo ? "8.6"
+) "The following test requires both 7.5 and 8.6 be known CUDA capabilities";
+assert
+  let
+    expected = {
+      cudaCapabilities = [
+        "7.5"
+        "8.6"
+      ];
+      cudaForwardCompat = true;
+
+      # Sorted alphabetically
+      archNames = [
+        "Ampere"
+        "Turing"
+      ];
+
+      realArches = [
+        "sm_75"
+        "sm_86"
+      ];
+
+      virtualArches = [
+        "compute_75"
+        "compute_86"
+      ];
+
+      arches = [
+        "sm_75"
+        "sm_86"
+        "compute_86"
+      ];
+
+      gencode = [
+        "-gencode=arch=compute_75,code=sm_75"
+        "-gencode=arch=compute_86,code=sm_86"
+        "-gencode=arch=compute_86,code=compute_86"
+      ];
+
+      gencodeString = "-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86";
+
+      cmakeCudaArchitecturesString = "75;86";
+    };
+    actual = formatCapabilities {
+      inherit cudaCapabilityToInfo;
+      cudaCapabilities = [
+        "7.5"
+        "8.6"
+      ];
+      cudaForwardCompat = true;
+    };
+    actualWrapped = (tryEval (deepSeq actual actual)).value;
+  in
+  assertMsg (expected == actualWrapped) ''
+    Expected: ${toJSON expected}
+    Actual: ${toJSON actualWrapped}
+  '';
+runCommand "${cudaNamePrefix}-tests-flags"
+  {
+    __structuredAttrs = true;
+    strictDeps = true;
+  }
+  ''
+    touch "$out"
+  ''

pkgs/development/cuda-modules/tests/opencv-and-torch/default.nix

@@ -0,0 +1,81 @@
+{
+  cudaPackages,
+  lib,
+  writeGpuTestPython,
+  # Configuration flags
+  openCVFirst,
+  useOpenCVDefaultCuda,
+  useTorchDefaultCuda,
+}:
+let
+  inherit (lib.strings) optionalString;
+
+  openCVBlock = ''
+
+    import cv2
+    print("OpenCV version:", cv2.__version__)
+
+    # Ensure OpenCV can access the GPU.
+    assert cv2.cuda.getCudaEnabledDeviceCount() > 0, "No CUDA devices found for OpenCV"
+    print("OpenCV CUDA device:", cv2.cuda.printCudaDeviceInfo(cv2.cuda.getDevice()))
+
+    # Ensure OpenCV can access the GPU.
+    print(cv2.getBuildInformation())
+
+    a = cv2.cuda.GpuMat(size=(256, 256), type=cv2.CV_32S, s=1)
+    b = cv2.cuda.GpuMat(size=(256, 256), type=cv2.CV_32S, s=1)
+    c = int(cv2.cuda.sum(cv2.cuda.add(a, b))[0]) # OpenCV returns a Scalar float object.
+
+    assert c == 2 * 256 * 256, f"Expected {2 * 256 * 256} OpenCV, got {c}"
+
+  '';
+
+  torchBlock = ''
+
+    import torch
+    print("Torch version:", torch.__version__)
+
+    # Set up the GPU.
+    torch.cuda.init()
+    # Ensure the GPU is available.
+    assert torch.cuda.is_available(), "CUDA is not available to Torch"
+    print("Torch CUDA device:", torch.cuda.get_device_properties(torch.cuda.current_device()))
+
+    a = torch.ones(256, 256, dtype=torch.int32).cuda()
+    b = torch.ones(256, 256, dtype=torch.int32).cuda()
+    c = (a + b).sum().item()
+    assert c == 2 * 256 * 256, f"Expected {2 * 256 * 256} for Torch, got {c}"
+
+  '';
+
+  content = if openCVFirst then openCVBlock + torchBlock else torchBlock + openCVBlock;
+
+  torchName = "torch" + optionalString useTorchDefaultCuda "-with-default-cuda";
+  openCVName = "opencv4" + optionalString useOpenCVDefaultCuda "-with-default-cuda";
+in
+# TODO: Ensure the expected CUDA libraries are loaded.
+# TODO: Ensure GPU access works as expected.
+writeGpuTestPython {
+  name = if openCVFirst then "${openCVName}-then-${torchName}" else "${torchName}-then-${openCVName}";
+  libraries =
+    # NOTE: These are purposefully in this order.
+    pythonPackages:
+    let
+      effectiveOpenCV = pythonPackages.opencv4.override (prevAttrs: {
+        cudaPackages = if useOpenCVDefaultCuda then prevAttrs.cudaPackages else cudaPackages;
+      });
+      effectiveTorch = pythonPackages.torchWithCuda.override (prevAttrs: {
+        cudaPackages = if useTorchDefaultCuda then prevAttrs.cudaPackages else cudaPackages;
+      });
+    in
+    if openCVFirst then
+      [
+        effectiveOpenCV
+        effectiveTorch
+      ]
+    else
+      [
+        effectiveTorch
+        effectiveOpenCV
+      ];
+} content