From 7edb29bd29dc2d563fa2269671e2bd1e3f421583 Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Mon, 16 Sep 2024 23:18:03 +0000 Subject: [PATCH] cudaPackages.tests: create tests for OpenCV/Torch ordering --- .../tests/opencv-and-torch/default.nix | 81 +++++++++++++++++++ pkgs/development/libraries/opencv/4.x.nix | 2 +- pkgs/top-level/all-packages.nix | 3 + pkgs/top-level/cuda-packages.nix | 41 ++++++++++ 4 files changed, 126 insertions(+), 1 deletion(-) create mode 100644 pkgs/development/cuda-modules/tests/opencv-and-torch/default.nix diff --git a/pkgs/development/cuda-modules/tests/opencv-and-torch/default.nix b/pkgs/development/cuda-modules/tests/opencv-and-torch/default.nix new file mode 100644 index 0000000000000..442bfb8a2dd0e --- /dev/null +++ b/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 diff --git a/pkgs/development/libraries/opencv/4.x.nix b/pkgs/development/libraries/opencv/4.x.nix index e979c9aa38633..3d5c0ac87067c 100644 --- a/pkgs/development/libraries/opencv/4.x.nix +++ b/pkgs/development/libraries/opencv/4.x.nix @@ -46,7 +46,7 @@ , cudaPackages ? {} , nvidia-optical-flow-sdk -, enableLto ? false # TODO: Investigate LTO build failures +, enableLto ? true , enableUnfree ? false , enableIpp ? false , enablePython ? false diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix index 837c1326af1f8..8c810af244d0c 100644 --- a/pkgs/top-level/all-packages.nix +++ b/pkgs/top-level/all-packages.nix @@ -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 { diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix index 639fa70446bee..b2a23363707e2 100644 --- a/pkgs/top-level/cuda-packages.nix +++ b/pkgs/top-level/cuda-packages.nix @@ -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 { }; });