ctranslate2_src_build_support/
native.rs

1use std::{
2    env,
3    path::{Path, PathBuf},
4};
5
6use crate::Os;
7
8pub fn build_native(
9    path: &Path,
10    os: Os,
11    cuda: bool,
12    cudnn: bool,
13    cuda_dynamic_loading: bool,
14    aarch64: bool,
15    mkl: bool,
16    openblas: bool,
17    ruy: bool,
18    accelarate: bool,
19    tensor_parallel: bool,
20    msse4_1: bool,
21    dnnl: bool,
22    openmp_comp: bool,
23    openmp_intel: bool,
24    flash_attention: bool,
25    cuda_small_binary: bool,
26    shared: bool,
27) -> PathBuf {
28    let mut cmake = cmake::Config::new(path);
29    cmake
30        .define("BUILD_CLI", "OFF")
31        .define("BUILD_SHARED_LIBS", "OFF")
32        .define("WITH_MKL", "OFF")
33        .define("OPENMP_RUNTIME", "NONE")
34        .define("CMAKE_POLICY_VERSION_MINIMUM", "3.5");
35    if shared {
36        cmake.define("BUILD_SHARED_LIBS", "ON");
37    }
38    if os == Os::Win {
39        let rustflags = env::var("CARGO_ENCODED_RUSTFLAGS").unwrap_or_default();
40        if !rustflags.contains("target-feature=+crt-static") {
41            println!(
42                "cargo:warning=For Windows compilation, setting the environment variable `RUSTFLAGS=-C target-feature=+crt-static` might be required."
43            );
44        } else {
45            cmake.static_crt(true);
46        }
47
48        println!("cargo::rustc-link-arg=/FORCE:MULTIPLE");
49        cmake.profile("Release").cxxflag("/EHsc");
50    } else if os == Os::Linux {
51        cmake.define("CMAKE_POSITION_INDEPENDENT_CODE", "ON");
52    }
53
54    if cuda {
55        let cuda = cuda_root().expect("CUDA_TOOLKIT_ROOT_DIR is not specified");
56        cmake.define("WITH_CUDA", "ON");
57        cmake.define("CUDA_TOOLKIT_ROOT_DIR", &cuda);
58        cmake.define("CUDA_ARCH_LIST", "5.3;6.0;6.2;7.0;7.2;7.5;8.0;8.6;8.9;9.0");
59        cmake.define(
60            "CUDA_NVCC_FLAGS",
61            format!(
62                "{}{}",
63                if cuda_small_binary {
64                    "-Xfatbin=-compress-all "
65                } else {
66                    ""
67                },
68                "-Xcompiler=-fPIC"
69            ),
70        );
71
72        if cudnn {
73            cmake.define("WITH_CUDNN", "ON");
74        }
75        if cuda_dynamic_loading {
76            cmake.define("CUDA_DYNAMIC_LOADING", "ON");
77        }
78    }
79    if os == Os::Mac && aarch64 {
80        cmake.define("CMAKE_OSX_ARCHITECTURES", "arm64");
81    }
82
83    if mkl {
84        cmake.define("WITH_MKL", "ON");
85    }
86    if openblas {
87        cmake.define("WITH_OPENBLAS", "ON");
88    }
89    if ruy {
90        cmake.define("WITH_RUY", "ON");
91    }
92    if accelarate {
93        cmake.define("WITH_ACCELERATE", "ON");
94    }
95    if tensor_parallel {
96        cmake.define("WITH_TENSOR_PARALLEL", "ON");
97    }
98    if msse4_1 {
99        cmake.define("CMAKE_CXX_FLAGS", "-msse4.1");
100    }
101    if dnnl {
102        cmake.define("WITH_DNNL", "ON");
103    }
104    if openmp_comp {
105        cmake.define("OPENMP_RUNTIME", "COMP");
106    } else if openmp_intel {
107        cmake.define("OPENMP_RUNTIME", "INTEL");
108    }
109    if flash_attention {
110        cmake.define("WITH_FLASH_ATTN", "ON");
111    }
112
113    let ctranslate2 = cmake.build();
114    ctranslate2.join("build")
115}
116
117// The function below was derived and modified from the `cudarc` crate.
118// Original source: https://github.com/coreylowman/cudarc/blob/main/build.rs
119//
120// Copyright (c) 2024 Corey Lowman
121//
122// Permission is hereby granted, free of charge, to any person obtaining a copy
123// of this software and associated documentation files (the "Software"), to deal
124// in the Software without restriction, including without limitation the rights
125// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
126// copies of the Software, and to permit persons to whom the Software is
127// furnished to do so, subject to the following conditions:
128//
129// The above copyright notice and this permission notice shall be included in all
130// copies or substantial portions of the Software.
131pub fn cuda_root() -> Option<PathBuf> {
132    let env_vars = [
133        "CUDA_PATH",
134        "CUDA_ROOT",
135        "CUDA_TOOLKIT_ROOT_DIR",
136        "CUDNN_LIB",
137    ];
138    let env_vars = env_vars
139        .into_iter()
140        .map(std::env::var)
141        .filter_map(Result::ok);
142
143    let roots = [
144        "/usr",
145        "/usr/local/cuda",
146        "/opt/cuda",
147        "/usr/lib/cuda",
148        "C:/Program Files/NVIDIA GPU Computing Toolkit",
149        "C:/CUDA",
150    ];
151    let roots = roots.into_iter().map(Into::into);
152    env_vars
153        .chain(roots)
154        .map(Into::<PathBuf>::into)
155        .find(|path| path.join("include").join("cuda.h").is_file())
156}