ferrum_native_ops/
build_unit.rs

1//! Typed coverage contract between Cargo feature units and native artifacts.
2
3use std::collections::BTreeSet;
4
5use ferrum_types::NativeOperatorBackend;
6use thiserror::Error;
7
8use crate::ResolvedNativeOperatorArtifactSet;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
11pub enum CudaNativeBuildUnit {
12    Marlin,
13    VllmMarlin,
14    VllmMoeMarlin,
15    VllmPagedAttentionV2,
16}
17
18impl CudaNativeBuildUnit {
19    pub const ALL: [Self; 4] = [
20        Self::Marlin,
21        Self::VllmMarlin,
22        Self::VllmMoeMarlin,
23        Self::VllmPagedAttentionV2,
24    ];
25
26    pub const fn as_str(self) -> &'static str {
27        match self {
28            Self::Marlin => "marlin",
29            Self::VllmMarlin => "vllm_marlin",
30            Self::VllmMoeMarlin => "vllm_moe_marlin",
31            Self::VllmPagedAttentionV2 => "vllm_paged_attention_v2",
32        }
33    }
34
35    pub const fn artifact_operator(self) -> &'static str {
36        match self {
37            Self::Marlin => "ferrum.cuda.marlin",
38            Self::VllmMarlin => "ferrum.cuda.vllm_marlin",
39            Self::VllmMoeMarlin => "ferrum.cuda.vllm_moe_marlin",
40            Self::VllmPagedAttentionV2 => "ferrum.cuda.vllm_paged_attention_v2",
41        }
42    }
43
44    pub fn from_artifact_operator(operator: &str) -> Option<Self> {
45        Self::ALL
46            .into_iter()
47            .find(|unit| unit.artifact_operator() == operator)
48    }
49
50    pub const fn required_exports(self) -> &'static [&'static str] {
51        match self {
52            Self::Marlin => &["marlin_cuda", "marlin_cuda_moe"],
53            Self::VllmMarlin => &[
54                "ferrum_marlin_mm",
55                "ferrum_marlin_mm_f16_u4b8",
56                "ferrum_vllm_gptq_marlin_repack",
57            ],
58            Self::VllmMoeMarlin => &[
59                "ferrum_vllm_marlin_moe_clear_profile_config",
60                "ferrum_vllm_marlin_moe_f16",
61                "ferrum_vllm_marlin_moe_set_profile_config",
62            ],
63            Self::VllmPagedAttentionV2 => &[
64                "ferrum_vllm_paged_attention_v1_f16_h128_b16",
65                "ferrum_vllm_paged_attention_v1_f16_h256_b16",
66                "ferrum_vllm_paged_attention_v2_f16_h128_b16",
67                "ferrum_vllm_paged_attention_v2_f16_h256_b16",
68                "ferrum_vnext_vllm_paged_attention_v1_f16_h128_b16_addressed",
69                "ferrum_vnext_vllm_paged_attention_v1_f16_h256_b16_addressed",
70                "ferrum_vnext_vllm_paged_attention_v2_f16_h128_b16_addressed",
71                "ferrum_vnext_vllm_paged_attention_v2_f16_h256_b16_addressed",
72            ],
73        }
74    }
75}
76
77#[derive(Debug, Clone, PartialEq, Eq)]
78pub struct ResolvedCudaNativeBuildCoverage {
79    units: BTreeSet<CudaNativeBuildUnit>,
80}
81
82impl ResolvedCudaNativeBuildCoverage {
83    pub fn resolve(
84        artifacts: &ResolvedNativeOperatorArtifactSet,
85        required: impl IntoIterator<Item = CudaNativeBuildUnit>,
86    ) -> Result<Self, CudaNativeBuildCoverageError> {
87        let views = artifacts
88            .artifacts
89            .iter()
90            .map(|artifact| ArtifactView {
91                operator: artifact.resolved.manifest.operator.as_str(),
92                backend: artifact.resolved.manifest.backend,
93                exports: artifact.resolved.manifest.exports.as_slice(),
94            })
95            .collect::<Vec<_>>();
96        resolve_views(&views, required)
97    }
98
99    pub fn contains(&self, unit: CudaNativeBuildUnit) -> bool {
100        self.units.contains(&unit)
101    }
102
103    pub fn iter(&self) -> impl Iterator<Item = CudaNativeBuildUnit> + '_ {
104        self.units.iter().copied()
105    }
106}
107
108#[derive(Debug, Error, PartialEq, Eq)]
109pub enum CudaNativeBuildCoverageError {
110    #[error(
111        "native artifact set does not provide CUDA build unit {unit} (required operator={operator})"
112    )]
113    MissingArtifact {
114        unit: &'static str,
115        operator: &'static str,
116    },
117    #[error("native artifact for CUDA build unit {unit} has backend {actual:?}, expected cuda")]
118    WrongBackend {
119        unit: &'static str,
120        actual: NativeOperatorBackend,
121    },
122    #[error("native artifact for CUDA build unit {unit} is missing required export {export}")]
123    MissingExport {
124        unit: &'static str,
125        export: &'static str,
126    },
127}
128
129struct ArtifactView<'a> {
130    operator: &'a str,
131    backend: NativeOperatorBackend,
132    exports: &'a [String],
133}
134
135fn resolve_views(
136    artifacts: &[ArtifactView<'_>],
137    required: impl IntoIterator<Item = CudaNativeBuildUnit>,
138) -> Result<ResolvedCudaNativeBuildCoverage, CudaNativeBuildCoverageError> {
139    let required = required.into_iter().collect::<BTreeSet<_>>();
140    for unit in &required {
141        let artifact = artifacts
142            .iter()
143            .find(|artifact| artifact.operator == unit.artifact_operator())
144            .ok_or(CudaNativeBuildCoverageError::MissingArtifact {
145                unit: unit.as_str(),
146                operator: unit.artifact_operator(),
147            })?;
148        if artifact.backend != NativeOperatorBackend::Cuda {
149            return Err(CudaNativeBuildCoverageError::WrongBackend {
150                unit: unit.as_str(),
151                actual: artifact.backend,
152            });
153        }
154        for required_export in unit.required_exports() {
155            if !artifact
156                .exports
157                .iter()
158                .any(|export| export == required_export)
159            {
160                return Err(CudaNativeBuildCoverageError::MissingExport {
161                    unit: unit.as_str(),
162                    export: required_export,
163                });
164            }
165        }
166    }
167    Ok(ResolvedCudaNativeBuildCoverage { units: required })
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    fn exports(unit: CudaNativeBuildUnit) -> Vec<String> {
175        unit.required_exports()
176            .iter()
177            .map(|value| (*value).to_string())
178            .collect()
179    }
180
181    #[test]
182    fn resolves_exact_required_build_units() {
183        let operator_exports = CudaNativeBuildUnit::ALL.map(exports);
184        let artifacts = CudaNativeBuildUnit::ALL
185            .iter()
186            .enumerate()
187            .map(|(index, unit)| ArtifactView {
188                operator: unit.artifact_operator(),
189                backend: NativeOperatorBackend::Cuda,
190                exports: &operator_exports[index],
191            })
192            .collect::<Vec<_>>();
193
194        let coverage =
195            resolve_views(&artifacts, CudaNativeBuildUnit::ALL).expect("complete artifact set");
196
197        assert!(CudaNativeBuildUnit::ALL
198            .into_iter()
199            .all(|unit| coverage.contains(unit)));
200    }
201
202    #[test]
203    fn rejects_missing_artifact_before_source_build_can_fallback() {
204        let error = resolve_views(&[], [CudaNativeBuildUnit::VllmMoeMarlin]).unwrap_err();
205
206        assert_eq!(
207            error,
208            CudaNativeBuildCoverageError::MissingArtifact {
209                unit: "vllm_moe_marlin",
210                operator: "ferrum.cuda.vllm_moe_marlin",
211            }
212        );
213    }
214
215    #[test]
216    fn rejects_artifact_that_does_not_export_the_linked_rust_abi() {
217        let exports = vec!["ferrum_vllm_marlin_moe_f16".to_string()];
218        let artifacts = [ArtifactView {
219            operator: CudaNativeBuildUnit::VllmMoeMarlin.artifact_operator(),
220            backend: NativeOperatorBackend::Cuda,
221            exports: &exports,
222        }];
223
224        let error = resolve_views(&artifacts, [CudaNativeBuildUnit::VllmMoeMarlin]).unwrap_err();
225
226        assert_eq!(
227            error,
228            CudaNativeBuildCoverageError::MissingExport {
229                unit: "vllm_moe_marlin",
230                export: "ferrum_vllm_marlin_moe_clear_profile_config",
231            }
232        );
233    }
234}
ferrum_native_ops/build_unit.rs

ferrum_native_ops/
build_unit.rs
