atlas_kernels/query.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Whole-registry lookups over the compiled targets, split out of `lib.rs` at
4//! the 500-LoC cap. Exact piecewise move — no logic changed.
5//!
6//! Distinct from [`super::resolve`], which picks ONE target for a given
7//! `(model_type, hidden_size)` and needs a tie-break; these two just enumerate
8//! or substring-match, and neither can fail.
9
10use super::{TargetPtxSet, all_ptx_sets};
11
12/// All compiled kernel targets and their PTX module sets.
13///
14/// Returns one entry per target compiled at build time.
15/// Single-target builds return one entry; wildcard builds return all.
16pub fn available_targets() -> Vec<TargetPtxSet> {
17 all_ptx_sets()
18}
19
20/// Find the PTX module set for a target whose model name contains `needle`.
21///
22/// Returns `None` if no compiled target matches.
23pub fn ptx_for_model(needle: &str) -> Option<TargetPtxSet> {
24 all_ptx_sets()
25 .into_iter()
26 .find(|t| t.target.model.contains(needle))
27}