khal_builder/
lib.rs

1//! Build-time utilities for compiling shader crates to SPIR-V and PTX.
2
3use std::path::{Path, PathBuf};
4use std::process::Command;
5
6/// Configures and runs the SPIR-V and PTX shader compilation pipeline.
7///
8/// Used in `build.rs` scripts to compile a shader crate before the host crate.
9pub struct KhalBuilder {
10    shader_crate: PathBuf,
11    // Useful for unusual crates layout where the src directory isn’t in `shader_crate/src`.
12    shader_src: Option<PathBuf>,
13    // Features to enable when building the library.
14    features: Vec<String>,
15    // The `RUST_MIN_STACK` given to the shader builders.
16    rust_min_stack: u32,
17    /// If the `cuda` feature is enabled and this is `true`, then cuda PTX kernels will be built with cargo-cuda.
18    /// Default: `true`
19    #[allow(dead_code)]
20    build_cuda: bool,
21    /// If this is `true`, then SpirV kernels will be built with cargo-gpu.
22    /// Default: `true`
23    build_spirv: bool,
24}
25
26impl KhalBuilder {
27    /// Creates a new builder for the given shader crate directory.
28    /// If `enable_builtin_features` is true, platform-specific features are auto-detected.
29    pub fn new(shader_crate: impl AsRef<Path>, enable_builtin_features: bool) -> Self {
30        let mut builder = Self {
31            shader_crate: shader_crate.as_ref().to_owned(),
32            shader_src: None,
33            features: Vec::new(),
34            build_cuda: true,
35            build_spirv: true,
36            rust_min_stack: 1024 * 1024 * 32,
37        };
38        if enable_builtin_features {
39            builder = builder.append_builtin_features();
40        }
41        builder
42    }
43
44    /// Sets the `RUST_MIN_STACK` environment variable for the shader compilation processes.
45    pub fn rust_min_stack(mut self, stack: u32) -> Self {
46        self.rust_min_stack = stack;
47        self
48    }
49
50    /// Overrides the shader source directory (defaults to `<shader_crate>/src`).
51    pub fn shader_src(mut self, src: impl AsRef<Path>) -> Self {
52        self.shader_src = Some(src.as_ref().to_owned());
53        self
54    }
55
56    /// Adds a cargo feature to enable when building the shader crate.
57    pub fn feature(mut self, feature: impl ToString) -> Self {
58        let feature = feature.to_string();
59        if !self.features.contains(&feature) {
60            self.features.push(feature);
61        }
62        self
63    }
64
65    /// Compiles the shader crate and writes output files to `output_dir`.
66    pub fn build(self, output_dir: impl AsRef<Path>) {
67        let output_dir = output_dir.as_ref();
68
69        self.setup_change_detection();
70
71        if self.build_spirv {
72            self.build_spirv(output_dir);
73        }
74
75        #[cfg(feature = "cuda")]
76        if self.build_cuda {
77            self.build_ptx(output_dir);
78        }
79    }
80
81    fn append_builtin_features(mut self) -> Self {
82        if cfg!(feature = "unsafe_remove_boundchecks") {
83            self = self.feature("unsafe-remove-boundchecks");
84        }
85
86        self
87    }
88
89    fn setup_change_detection(&self) {
90        println!(
91            "cargo:rerun-if-changed={}",
92            self.shader_crate.to_string_lossy()
93        );
94        let shader_src = self
95            .shader_src
96            .clone()
97            .unwrap_or_else(|| self.shader_crate.join("src"));
98        for entry in walkdir::WalkDir::new(shader_src)
99            .into_iter()
100            .filter_map(|e| e.ok())
101        {
102            println!("cargo:rerun-if-changed={}", entry.path().display());
103        }
104
105        println!("cargo:rerun-if-env-changed=CARGO_FEATURE_PUSH_CONSTANTS"); // TODO: currently unused
106        println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA");
107    }
108
109    fn build_spirv(&self, output_dir: impl AsRef<Path>) {
110        let output_dir = output_dir.as_ref();
111        let mut args = vec![
112            "gpu",
113            "build",
114            "--shader-crate",
115            self.shader_crate
116                .to_str()
117                .expect("Invalid shader crate path"),
118            "--output-dir",
119            output_dir.to_str().expect("Invalid output directory path"),
120            "--multimodule",
121        ];
122
123        let features_str = self.features.join(",");
124        if !features_str.is_empty() {
125            args.push("--features");
126            args.push(&features_str);
127        }
128
129        let status = Command::new("cargo")
130            .args(args)
131            .env("RUST_MIN_STACK", self.rust_min_stack.to_string())
132            .status()
133            .expect("failed to run cargo gpu");
134
135        if !status.success() {
136            panic!("cargo gpu build failed");
137        }
138    }
139
140    /// Compiles the shader crate to PTX for the CUDA backend.
141    #[cfg(feature = "cuda")]
142    fn build_ptx(&self, output_dir: impl AsRef<Path>) {
143        let output_dir = output_dir.as_ref();
144        let features_str = self.features.join(",");
145
146        let mut args = vec![
147            "cuda",
148            "build",
149            "--shader-crate",
150            self.shader_crate
151                .to_str()
152                .expect("Invalid shader crate path"),
153            "--output-dir",
154            output_dir.to_str().expect("Invalid output directory path"),
155        ];
156
157        if !features_str.is_empty() {
158            args.push("--features");
159            args.push(&features_str);
160        }
161
162        let status = Command::new("cargo")
163            .args(args)
164            .env("RUST_MIN_STACK", self.rust_min_stack.to_string())
165            .status()
166            .expect("failed to run cargo cuda");
167
168        if !status.success() {
169            panic!("cargo cuda build failed");
170        }
171    }
172}
khal_builder/lib.rs

khal_builder/
lib.rs
