Commit 7edb29bd authored by Connor Baker's avatar Connor Baker
Browse files

cudaPackages.tests: create tests for OpenCV/Torch ordering

parent aa117cc2

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

0 → 100644
+81 −0
Original line number Diff line number Diff line 
{

  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

pkgs/development/libraries/opencv/4.x.nix

+1 −1
Original line number Diff line number Diff line
@@ -46,7 +46,7 @@ 
, cudaPackages ? {}

, nvidia-optical-flow-sdk



, enableLto ? false # TODO: Investigate LTO build failures

, enableLto ? true

, enableUnfree ? false

, enableIpp ? false

, enablePython ? false

pkgs/top-level/all-packages.nix

+3 −0
Original line number Diff line number Diff line
@@ -22594,6 +22594,9 @@ with pkgs; 
    pythonPackages = python3Packages;

    # TODO(@connorbaker): OpenCV 4.9 only supports up to CUDA 12.3.

    cudaPackages = cudaPackages_12_3;

    # TODO: LTO does not work.

    # https://github.com/NixOS/nixpkgs/issues/343123

    enableLto = false;

  };















  opencv4WithoutCuda = opencv4.override {

























pkgs/top-level/cuda-packages.nix



+41
−0






























Original line number
Diff line number
Diff line








@@ -33,7 +33,9 @@ let













    attrsets















    customisation















    fixedPoints













    lists















    strings













    trivial















    versions















    ;















  # Backbone










@@ -81,6 +83,45 @@ let













    nccl = final.callPackage ../development/cuda-modules/nccl { };















    nccl-tests = final.callPackage ../development/cuda-modules/nccl-tests { };





























    tests =













      let













        bools = [













          true













          false













        ];













        configs = {













          openCVFirst = bools;













          useOpenCVDefaultCuda = bools;













          useTorchDefaultCuda = bools;













        };













        builder =













          {













            openCVFirst,













            useOpenCVDefaultCuda,













            useTorchDefaultCuda,













          }@config:













          {













            name = strings.concatStringsSep "-" (













              [













                "test"













                (if openCVFirst then "opencv" else "torch")













              ]













              ++ lists.optionals (if openCVFirst then useOpenCVDefaultCuda else useTorchDefaultCuda) [













                "with-default-cuda"













              ]













              ++ [













                "then"













                (if openCVFirst then "torch" else "opencv")













              ]













              ++ lists.optionals (if openCVFirst then useTorchDefaultCuda else useOpenCVDefaultCuda) [













                "with-default-cuda"













              ]













            );













            value = final.callPackage ../development/cuda-modules/tests/opencv-and-torch config;













          };













      in













      attrsets.listToAttrs (attrsets.mapCartesianProduct builder configs);





























    writeGpuTestPython = final.callPackage ../development/cuda-modules/write-gpu-test-python.nix { };















  });