pyrox.dev
1{ lib
2, config
3, fetchFromGitHub
4, stdenv
5, cmake
6, pkg-config
7, cudaPackages ? { }
8, symlinkJoin
9, tbb
10, hostSystem ? "CPP"
11, deviceSystem ? if config.cudaSupport then "CUDA" else "OMP"
12}:
13
14# Policy for device_vector<T>
15assert builtins.elem deviceSystem [
16 "CPP" # Serial on CPU
17 "OMP" # Parallel with OpenMP
18 "TBB" # Parallel with Intel TBB
19 "CUDA" # Parallel on GPU
20];
21
22# Policy for host_vector<T>
23# Always lives on CPU, but execution can be made parallel
24assert builtins.elem hostSystem [ "CPP" "OMP" "TBB" ];
25
26let
27 pname = "nvidia-thrust";
28 version = "1.16.0";
29
30 inherit (cudaPackages) backendStdenv cudaFlags;
31 cudaCapabilities = map cudaFlags.dropDot cudaFlags.cudaCapabilities;
32
33 tbbSupport = builtins.elem "TBB" [ deviceSystem hostSystem ];
34 cudaSupport = deviceSystem == "CUDA";
35
36 # TODO: Would like to use this:
37 cudaJoined = symlinkJoin {
38 name = "cuda-packages-unsplit";
39 paths = with cudaPackages; [
40 cuda_nvcc
41 cuda_nvrtc # symbols: cudaLaunchDevice, &c; notice postBuild
42 cuda_cudart # cuda_runtime.h
43 libcublas
44 ];
45 postBuild = ''
46 ln -s $out/lib $out/lib64
47 '';
48 };
49in
50stdenv.mkDerivation {
51 inherit pname version;
52
53 src = fetchFromGitHub {
54 owner = "NVIDIA";
55 repo = "thrust";
56 rev = version;
57 fetchSubmodules = true;
58 hash = "sha256-/EyznxWKuHuvHNjq+SQg27IaRbtkjXR2zlo2YgCWmUQ=";
59 };
60
61 # NVIDIA's "compiler hacks" seem like work-arounds for legacy toolchains and
62 # cause us errors such as:
63 # > Thrust's test harness uses CMAKE_CXX_COMPILER for the CUDA host compiler.
64 # > Refusing to overwrite specified CMAKE_CUDA_HOST_COMPILER
65 # So we un-fix cmake after them:
66 postPatch = ''
67 echo > cmake/ThrustCompilerHacks.cmake
68 '';
69
70 buildInputs = lib.optionals tbbSupport [ tbb ];
71
72 nativeBuildInputs = [
73 cmake
74 pkg-config
75 ] ++ lib.optionals cudaSupport [
76 # Goes in native build inputs because thrust looks for headers
77 # in a path relative to nvcc...
78 cudaJoined
79 ];
80
81 cmakeFlags = [
82 "-DTHRUST_INCLUDE_CUB_CMAKE=${if cudaSupport then "ON" else "OFF"}"
83 "-DTHRUST_DEVICE_SYSTEM=${deviceSystem}"
84 "-DTHRUST_HOST_SYSTEM=${hostSystem}"
85 "-DTHRUST_AUTO_DETECT_COMPUTE_ARCHS=OFF"
86 "-DTHRUST_DISABLE_ARCH_BY_DEFAULT=ON"
87 ] ++ lib.optionals cudaFlags.enableForwardCompat [
88 "-DTHRUST_ENABLE_COMPUTE_FUTURE=ON"
89 ] ++ map (sm: "THRUST_ENABLE_COMPUTE_${sm}") cudaCapabilities;
90
91 passthru = {
92 inherit cudaSupport cudaPackages cudaJoined;
93 };
94
95 meta = with lib; {
96 description = "A high-level C++ parallel algorithms library that builds on top of CUDA, TBB, OpenMP, etc";
97 homepage = "https://github.com/NVIDIA/thrust";
98 license = licenses.asl20;
99 platforms = platforms.unix;
100 maintainers = with maintainers; [ SomeoneSerge ];
101 };
102}