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 include_paths: Vec<PathBuf> = env::var("CMAKE_INCLUDE_PATH")
29        .as_ref()
30        .map(|v| env::split_paths(v).collect())
31        .unwrap_or_default();
32    let mut library_paths: Vec<PathBuf> = env::var("CMAKE_LIBRARY_PATH")
33        .as_ref()
34        .map(|v| env::split_paths(v).collect())
35        .unwrap_or_default();
36    let mut cmake = cmake::Config::new(path);
37    cmake
38        .define("BUILD_CLI", "OFF")
39        .define("BUILD_SHARED_LIBS", "OFF")
40        .define("WITH_MKL", "OFF")
41        .define("OPENMP_RUNTIME", "NONE")
42        .define("CMAKE_POLICY_VERSION_MINIMUM", "3.5");
43    if shared {
44        cmake.define("BUILD_SHARED_LIBS", "ON");
45    }
46    if os == Os::Win {
47        let rustflags = env::var("CARGO_ENCODED_RUSTFLAGS").unwrap_or_default();
48        if !rustflags.contains("target-feature=+crt-static") {
49            println!(
50                "cargo:warning=For Windows compilation, setting the environment variable `RUSTFLAGS=-C target-feature=+crt-static` might be required."
51            );
52        } else {
53            cmake.static_crt(true);
54        }
55
56        println!("cargo::rustc-link-arg=/FORCE:MULTIPLE");
57        cmake.profile("Release").cxxflag("/EHsc");
58    } else if os == Os::Linux {
59        cmake.define("CMAKE_POSITION_INDEPENDENT_CODE", "ON");
60    }
61
62    if cuda {
63        let cuda = cuda_root().expect("CUDA_TOOLKIT_ROOT_DIR is not specified");
64        cmake.define("WITH_CUDA", "ON");
65        cmake.define("CUDA_TOOLKIT_ROOT_DIR", &cuda);
66        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");
67        cmake.define(
68            "CUDA_NVCC_FLAGS",
69            format!(
70                "{}{}",
71                if cuda_small_binary {
72                    "-Xfatbin=-compress-all "
73                } else {
74                    ""
75                },
76                "-Xcompiler=-fPIC"
77            ),
78        );
79
80        if cudnn {
81            cmake.define("WITH_CUDNN", "ON");
82        }
83        if cuda_dynamic_loading {
84            cmake.define("CUDA_DYNAMIC_LOADING", "ON");
85        }
86    }
87    if os == Os::Mac && aarch64 {
88        cmake.define("CMAKE_OSX_ARCHITECTURES", "arm64");
89    }
90
91    if mkl {
92        cmake.define("WITH_MKL", "ON");
93        if let Ok(mklroot) = env::var("DEP_MKL_ROOT") {
94            cmake.env("MKLROOT", mklroot);
95        }
96        if let Ok(include_path) = env::var("DEP_MKL_INCLUDE_PATH") {
97            include_paths.push(PathBuf::from(include_path));
98        }
99        if let Ok(library_path) = env::var("DEP_MKL_LIBRARY_PATH") {
100            library_paths.push(PathBuf::from(library_path));
101        }
102    }
103    if openblas {
104        cmake.define("WITH_OPENBLAS", "ON");
105        if os != Os::Win {
106            include_paths.push(PathBuf::from(env::var("DEP_OPENBLAS_INCLUDE").unwrap()));
107            library_paths.push(PathBuf::from(env::var("DEP_OPENBLAS_LIBRARY").unwrap()));
108        }
109    }
110    if ruy {
111        cmake.define("WITH_RUY", "ON");
112    }
113    if accelarate {
114        cmake.define("WITH_ACCELERATE", "ON");
115    }
116    if tensor_parallel {
117        cmake.define("WITH_TENSOR_PARALLEL", "ON");
118    }
119    if msse4_1 {
120        cmake.define("CMAKE_CXX_FLAGS", "-msse4.1");
121    }
122    if dnnl {
123        cmake.define("WITH_DNNL", "ON");
124        include_paths.push(PathBuf::from(env::var("DEP_DNNL_INCLUDE_PATH").unwrap()));
125        library_paths.push(PathBuf::from(env::var("DEP_DNNL_LIBRARY_PATH").unwrap()));
126    }
127    if openmp_comp {
128        cmake.define("OPENMP_RUNTIME", "COMP");
129    } else if openmp_intel {
130        cmake.define("OPENMP_RUNTIME", "INTEL");
131    }
132    if flash_attention {
133        cmake.define("WITH_FLASH_ATTN", "ON");
134    }
135
136    if !include_paths.is_empty() {
137        cmake.env(
138            "CMAKE_INCLUDE_PATH",
139            env::join_paths(include_paths).unwrap(),
140        );
141    }
142
143    if !library_paths.is_empty() {
144        cmake.env(
145            "CMAKE_LIBRARY_PATH",
146            env::join_paths(library_paths).unwrap(),
147        );
148    }
149
150    let ctranslate2 = cmake.build();
151    ctranslate2.join("build")
152}
153
154// The function below was derived and modified from the `cudarc` crate.
155// Original source: https://github.com/coreylowman/cudarc/blob/main/build.rs
156//
157// Copyright (c) 2024 Corey Lowman
158//
159// Permission is hereby granted, free of charge, to any person obtaining a copy
160// of this software and associated documentation files (the "Software"), to deal
161// in the Software without restriction, including without limitation the rights
162// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
163// copies of the Software, and to permit persons to whom the Software is
164// furnished to do so, subject to the following conditions:
165//
166// The above copyright notice and this permission notice shall be included in all
167// copies or substantial portions of the Software.
168pub fn cuda_root() -> Option<PathBuf> {
169    let env_vars = [
170        "CUDA_PATH",
171        "CUDA_ROOT",
172        "CUDA_TOOLKIT_ROOT_DIR",
173        "CUDNN_LIB",
174    ];
175    let env_vars = env_vars
176        .into_iter()
177        .map(std::env::var)
178        .filter_map(Result::ok);
179
180    let roots = [
181        "/usr",
182        "/usr/local/cuda",
183        "/opt/cuda",
184        "/usr/lib/cuda",
185        "C:/Program Files/NVIDIA GPU Computing Toolkit",
186        "C:/CUDA",
187    ];
188    let roots = roots.into_iter().map(Into::into);
189    env_vars
190        .chain(roots)
191        .map(Into::<PathBuf>::into)
192        .find(|path| path.join("include").join("cuda.h").is_file())
193}
ctranslate2_src_build_support/native.rs

ctranslate2_src_build_support/
native.rs
