Skip to content

Commit 5f4d27b

Browse files
danieldkdanieldk
committed
Factor out CMake helper file generation
1 parent 0c9b25d commit 5f4d27b

5 files changed

Lines changed: 41 additions & 105 deletions

File tree

‎build2cmake/src/torch/common.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,12 @@ use crate::metadata::Metadata;
99
use crate::FileSet;
1010

1111
static REGISTRATION_H: &str = include_str!("../templates/registration.h");
12+
static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
13+
static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
14+
static WINDOWS_UTILS: &str = include_str!("../templates/windows.cmake");
15+
static HIPIFY: &str = include_str!("../templates/cuda/hipify.py");
16+
static COMPILE_METAL_CMAKE: &str = include_str!("../templates/metal/compile-metal.cmake");
17+
static METALLIB_TO_HEADER_PY: &str = include_str!("../templates/metal/metallib_to_header.py");
1218

1319
pub fn write_pyproject_toml(
1420
env: &Environment,
@@ -104,3 +110,30 @@ pub fn write_ops_py(
104110

105111
Ok(())
106112
}
113+
114+
/// Helper function to write a file to the cmake subdirectory
115+
pub fn write_cmake_file(file_set: &mut FileSet, filename: &str, content: &[u8]) {
116+
let mut path = PathBuf::new();
117+
path.push("cmake");
118+
path.push(filename);
119+
file_set.entry(path).extend_from_slice(content);
120+
}
121+
122+
/// Writes all CMake helper files that any backend might need.
123+
/// Each backend will use only the files it references in its CMakeLists.txt.
124+
pub fn write_cmake_helpers(file_set: &mut FileSet) {
125+
write_cmake_file(file_set, "utils.cmake", CMAKE_UTILS.as_bytes());
126+
write_cmake_file(file_set, "kernel.cmake", CMAKE_KERNEL.as_bytes());
127+
write_cmake_file(file_set, "windows.cmake", WINDOWS_UTILS.as_bytes());
128+
write_cmake_file(file_set, "hipify.py", HIPIFY.as_bytes());
129+
write_cmake_file(
130+
file_set,
131+
"compile-metal.cmake",
132+
COMPILE_METAL_CMAKE.as_bytes(),
133+
);
134+
write_cmake_file(
135+
file_set,
136+
"metallib_to_header.py",
137+
METALLIB_TO_HEADER_PY.as_bytes(),
138+
);
139+
}

‎build2cmake/src/torch/cpu.rs‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use minijinja::{context, Environment};
66

77
use crate::config::{Backend, Build, Torch};
88
use crate::fileset::FileSet;
9+
use crate::torch::common::write_cmake_helpers;
910
use crate::torch::common::write_metadata;
1011
use crate::torch::common::write_ops_py;
1112
use crate::torch::common::write_pyproject_toml;
@@ -14,9 +15,6 @@ use crate::torch::kernel::render_kernel_components;
1415
use crate::torch::kernel_ops_identifier;
1516
use crate::version::Version;
1617

17-
static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
18-
static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
19-
2018
pub fn write_torch_ext_cpu(
2119
env: &Environment,
2220
build: &Build,
@@ -68,19 +66,7 @@ fn write_cmake(
6866
ops_name: &str,
6967
file_set: &mut FileSet,
7068
) -> Result<()> {
71-
let mut utils_path = PathBuf::new();
72-
utils_path.push("cmake");
73-
utils_path.push("utils.cmake");
74-
file_set
75-
.entry(utils_path.clone())
76-
.extend_from_slice(CMAKE_UTILS.as_bytes());
77-
78-
let mut kernel_path = PathBuf::new();
79-
kernel_path.push("cmake");
80-
kernel_path.push("kernel.cmake");
81-
file_set
82-
.entry(kernel_path.clone())
83-
.extend_from_slice(CMAKE_KERNEL.as_bytes());
69+
write_cmake_helpers(file_set);
8470

8571
let cmake_writer = file_set.entry("CMakeLists.txt");
8672

‎build2cmake/src/torch/cuda.rs‎

Lines changed: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use minijinja::{context, Environment};
88

99
use crate::config::{Backend, Build, Dependency, Torch};
1010
use crate::torch::common::prefix_and_join_includes;
11+
use crate::torch::common::write_cmake_helpers;
1112
use crate::torch::common::write_metadata;
1213
use crate::torch::common::write_ops_py;
1314
use crate::torch::common::write_pyproject_toml;
@@ -17,11 +18,6 @@ use crate::torch::kernel_ops_identifier;
1718
use crate::version::Version;
1819
use crate::FileSet;
1920

20-
static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
21-
static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
22-
static WINDOWS_UTILS: &str = include_str!("../templates/windows.cmake");
23-
static HIPIFY: &str = include_str!("../templates/cuda/hipify.py");
24-
2521
pub fn write_torch_ext_cuda(
2622
env: &Environment,
2723
backend: Backend,
@@ -103,33 +99,7 @@ fn write_cmake(
10399
ops_name: &str,
104100
file_set: &mut FileSet,
105101
) -> Result<()> {
106-
let mut utils_path = PathBuf::new();
107-
utils_path.push("cmake");
108-
utils_path.push("utils.cmake");
109-
file_set
110-
.entry(utils_path.clone())
111-
.extend_from_slice(CMAKE_UTILS.as_bytes());
112-
113-
let mut kernel_path = PathBuf::new();
114-
kernel_path.push("cmake");
115-
kernel_path.push("kernel.cmake");
116-
file_set
117-
.entry(kernel_path.clone())
118-
.extend_from_slice(CMAKE_KERNEL.as_bytes());
119-
120-
let mut windows_utils_path = PathBuf::new();
121-
windows_utils_path.push("cmake");
122-
windows_utils_path.push("windows.cmake");
123-
file_set
124-
.entry(windows_utils_path.clone())
125-
.extend_from_slice(WINDOWS_UTILS.as_bytes());
126-
127-
let mut hipify_path = PathBuf::new();
128-
hipify_path.push("cmake");
129-
hipify_path.push("hipify.py");
130-
file_set
131-
.entry(hipify_path.clone())
132-
.extend_from_slice(HIPIFY.as_bytes());
102+
write_cmake_helpers(file_set);
133103

134104
let cmake_writer = file_set.entry("CMakeLists.txt");
135105

‎build2cmake/src/torch/metal.rs‎

Lines changed: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use minijinja::{context, Environment};
66

77
use crate::config::{Backend, Build, Torch};
88
use crate::fileset::FileSet;
9+
use crate::torch::common::write_cmake_helpers;
910
use crate::torch::common::write_metadata;
1011
use crate::torch::common::write_ops_py;
1112
use crate::torch::common::write_pyproject_toml;
@@ -14,11 +15,6 @@ use crate::torch::kernel::render_kernel_components;
1415
use crate::torch::kernel_ops_identifier;
1516
use crate::version::Version;
1617

17-
static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
18-
static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
19-
static COMPILE_METAL_CMAKE: &str = include_str!("../templates/metal/compile-metal.cmake");
20-
static METALLIB_TO_HEADER_PY: &str = include_str!("../templates/metal/metallib_to_header.py");
21-
2218
pub fn write_torch_ext_metal(
2319
env: &Environment,
2420
build: &Build,
@@ -70,33 +66,7 @@ fn write_cmake(
7066
ops_name: &str,
7167
file_set: &mut FileSet,
7268
) -> Result<()> {
73-
let mut utils_path = PathBuf::new();
74-
utils_path.push("cmake");
75-
utils_path.push("utils.cmake");
76-
file_set
77-
.entry(utils_path.clone())
78-
.extend_from_slice(CMAKE_UTILS.as_bytes());
79-
80-
let mut kernel_path = PathBuf::new();
81-
kernel_path.push("cmake");
82-
kernel_path.push("kernel.cmake");
83-
file_set
84-
.entry(kernel_path.clone())
85-
.extend_from_slice(CMAKE_KERNEL.as_bytes());
86-
87-
let mut compile_metal_path = PathBuf::new();
88-
compile_metal_path.push("cmake");
89-
compile_metal_path.push("compile-metal.cmake");
90-
file_set
91-
.entry(compile_metal_path)
92-
.extend_from_slice(COMPILE_METAL_CMAKE.as_bytes());
93-
94-
let mut metallib_to_header_path = PathBuf::new();
95-
metallib_to_header_path.push("cmake");
96-
metallib_to_header_path.push("metallib_to_header.py");
97-
file_set
98-
.entry(metallib_to_header_path)
99-
.extend_from_slice(METALLIB_TO_HEADER_PY.as_bytes());
69+
write_cmake_helpers(file_set);
10070

10171
let cmake_writer = file_set.entry("CMakeLists.txt");
10272

‎build2cmake/src/torch/xpu.rs‎

Lines changed: 2 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use itertools::Itertools;
77
use minijinja::{context, Environment};
88

99
use crate::config::{Backend, Build, Dependency, Torch};
10+
use crate::torch::common::write_cmake_helpers;
1011
use crate::torch::common::write_metadata;
1112
use crate::torch::common::write_ops_py;
1213
use crate::torch::common::write_pyproject_toml;
@@ -16,10 +17,6 @@ use crate::torch::kernel_ops_identifier;
1617
use crate::version::Version;
1718
use crate::FileSet;
1819

19-
static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
20-
static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
21-
static WINDOWS_UTILS: &str = include_str!("../templates/windows.cmake");
22-
2320
pub fn write_torch_ext_xpu(
2421
env: &Environment,
2522
build: &Build,
@@ -98,27 +95,7 @@ fn write_cmake(
9895
ops_name: &str,
9996
file_set: &mut FileSet,
10097
) -> Result<()> {
101-
let mut utils_path = PathBuf::new();
102-
utils_path.push("cmake");
103-
utils_path.push("utils.cmake");
104-
file_set
105-
.entry(utils_path.clone())
106-
.extend_from_slice(CMAKE_UTILS.as_bytes());
107-
108-
// Add windows.cmake for Windows-specific install targets
109-
let mut windows_utils_path = PathBuf::new();
110-
windows_utils_path.push("cmake");
111-
windows_utils_path.push("windows.cmake");
112-
file_set
113-
.entry(windows_utils_path.clone())
114-
.extend_from_slice(WINDOWS_UTILS.as_bytes());
115-
116-
let mut kernel_path = PathBuf::new();
117-
kernel_path.push("cmake");
118-
kernel_path.push("kernel.cmake");
119-
file_set
120-
.entry(kernel_path.clone())
121-
.extend_from_slice(CMAKE_KERNEL.as_bytes());
98+
write_cmake_helpers(file_set);
12299

123100
let cmake_writer = file_set.entry("CMakeLists.txt");
124101

0 commit comments

Comments
 (0)