tjh.dev
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 cmake,
6 llvmPackages, # openmp
7 withMkl ? false,
8 mkl,
9 withCUDA ? false,
10 withCuDNN ? false,
11 cudaPackages,
12 # Enabling both withOneDNN and withOpenblas is broken
13 # https://github.com/OpenNMT/CTranslate2/issues/1294
14 withOneDNN ? false,
15 oneDNN,
16 withOpenblas ? true,
17 openblas,
18 withRuy ? true,
19
20 # passthru tests
21 libretranslate,
22 wyoming-faster-whisper,
23}:
24
25let
26 cmakeBool = b: if b then "ON" else "OFF";
27in
28stdenv.mkDerivation rec {
29 pname = "ctranslate2";
30 version = "4.6.0";
31
32 src = fetchFromGitHub {
33 owner = "OpenNMT";
34 repo = "CTranslate2";
35 rev = "v${version}";
36 hash = "sha256-EM2kunqtxo0BTIzrEomfaRsdav7sx6QEOhjDtjjSoYY=";
37 fetchSubmodules = true;
38 };
39
40 nativeBuildInputs = [
41 cmake
42 ]
43 ++ lib.optionals withCUDA [
44 cudaPackages.cuda_nvcc
45 ];
46
47 cmakeFlags = [
48 # https://opennmt.net/CTranslate2/installation.html#build-options
49 # https://github.com/OpenNMT/CTranslate2/blob/54810350e662ebdb01ecbf8e4a746f02aeff1dd7/python/tools/prepare_build_environment_linux.sh#L53
50 # https://github.com/OpenNMT/CTranslate2/blob/59d223abcc7e636c1c2956e62482bc3299cc7766/python/tools/prepare_build_environment_macos.sh#L12
51 "-DOPENMP_RUNTIME=COMP"
52 "-DWITH_CUDA=${cmakeBool withCUDA}"
53 "-DWITH_CUDNN=${cmakeBool withCuDNN}"
54 "-DWITH_DNNL=${cmakeBool withOneDNN}"
55 "-DWITH_OPENBLAS=${cmakeBool withOpenblas}"
56 "-DWITH_RUY=${cmakeBool withRuy}"
57 "-DWITH_MKL=${cmakeBool withMkl}"
58 ]
59 ++ lib.optional stdenv.hostPlatform.isDarwin "-DWITH_ACCELERATE=ON";
60
61 buildInputs =
62 lib.optionals withMkl [
63 mkl
64 ]
65 ++ lib.optionals withCUDA [
66 cudaPackages.cuda_cccl # <nv/target> required by the fp16 headers in cudart
67 cudaPackages.cuda_cudart
68 cudaPackages.libcublas
69 cudaPackages.libcurand
70 ]
71 ++ lib.optionals (withCUDA && withCuDNN) [
72 cudaPackages.cudnn
73 ]
74 ++ lib.optionals withOneDNN [
75 oneDNN
76 ]
77 ++ lib.optionals withOpenblas [
78 openblas
79 ]
80 ++ lib.optionals stdenv.hostPlatform.isDarwin [
81 llvmPackages.openmp
82 ];
83
84 passthru.tests = {
85 inherit
86 libretranslate
87 wyoming-faster-whisper
88 ;
89 };
90
91 meta = with lib; {
92 description = "Fast inference engine for Transformer models";
93 mainProgram = "ct2-translator";
94 homepage = "https://github.com/OpenNMT/CTranslate2";
95 changelog = "https://github.com/OpenNMT/CTranslate2/blob/${src.rev}/CHANGELOG.md";
96 license = licenses.mit;
97 maintainers = with maintainers; [
98 hexa
99 misuzu
100 ];
101 broken = (cudaPackages.cudaOlder "11.4") || !(withCuDNN -> withCUDA);
102 };
103}