use crate::fkc::error::FkcError;
use crate::fkc::schema::LayoutSpec;
use crate::kernel::KernelCaps;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Tri {
Required,
Accepted,
Rejected,
NotApplicable,
}
impl Tri {
fn parse(
value: Option<&str>,
default: Tri,
section: &str,
operand: &str,
flag: &str,
) -> Result<Tri, FkcError> {
match value {
None => Ok(default),
Some(v) => match v.trim() {
"required" => Ok(Tri::Required),
"accepted" => Ok(Tri::Accepted),
"rejected" => Ok(Tri::Rejected),
"n/a" | "na" => Ok(Tri::NotApplicable),
other => Err(FkcError::BadLayoutFlag {
section: section.to_string(),
operand: operand.to_string(),
flag: flag.to_string(),
value: other.to_string(),
}),
},
}
}
pub fn is_accepted(self) -> bool {
matches!(self, Tri::Accepted)
}
pub fn is_required(self) -> bool {
matches!(self, Tri::Required)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ResolvedLayout {
pub contiguous: Tri,
pub strided: Tri,
pub broadcast_stride0: Tri,
pub start_offset: Tri,
pub reverse_strides: Tri,
}
impl ResolvedLayout {
pub fn project(&self) -> KernelCaps {
let strided_input = self.strided.is_accepted() && self.broadcast_stride0.is_accepted();
let requires_broadcast = self.broadcast_stride0.is_required();
KernelCaps {
strided_input,
requires_broadcast,
}
}
}
pub fn resolve_layout(
spec: Option<&LayoutSpec>,
section: &str,
operand: &str,
) -> Result<ResolvedLayout, FkcError> {
let s = spec;
Ok(ResolvedLayout {
contiguous: Tri::parse(
s.and_then(|l| l.contiguous.as_deref()),
Tri::NotApplicable,
section,
operand,
"contiguous",
)?,
strided: Tri::parse(
s.and_then(|l| l.strided.as_deref()),
Tri::Rejected,
section,
operand,
"strided",
)?,
broadcast_stride0: Tri::parse(
s.and_then(|l| l.broadcast_stride0.as_deref()),
Tri::Rejected,
section,
operand,
"broadcast_stride0",
)?,
start_offset: Tri::parse(
s.and_then(|l| l.start_offset.as_deref()),
Tri::Rejected,
section,
operand,
"start_offset",
)?,
reverse_strides: Tri::parse(
s.and_then(|l| l.reverse_strides.as_deref()),
Tri::Rejected,
section,
operand,
"reverse_strides",
)?,
})
}
pub fn project_kernel_caps(operand_layouts: &[ResolvedLayout]) -> KernelCaps {
if operand_layouts.is_empty() {
return KernelCaps::empty();
}
let strided_input = operand_layouts.iter().all(|l| l.project().strided_input);
let requires_broadcast = operand_layouts
.iter()
.any(|l| l.broadcast_stride0.is_required());
KernelCaps {
strided_input,
requires_broadcast,
}
}
pub fn is_generic_contract(layouts: &[ResolvedLayout]) -> bool {
!layouts.is_empty()
&& layouts.iter().all(|l| {
l.strided.is_accepted()
&& l.broadcast_stride0.is_accepted()
&& l.contiguous != Tri::Required
})
}
#[cfg(test)]
mod tests {
use super::*;
fn spec(
contiguous: Option<&str>,
strided: Option<&str>,
broadcast: Option<&str>,
start_offset: Option<&str>,
reverse: Option<&str>,
) -> LayoutSpec {
LayoutSpec {
contiguous: contiguous.map(String::from),
strided: strided.map(String::from),
broadcast_stride0: broadcast.map(String::from),
broadcast_axes: None,
start_offset: start_offset.map(String::from),
reverse_strides: reverse.map(String::from),
awkward_layout_strategy: None,
}
}
fn resolve(s: &LayoutSpec) -> ResolvedLayout {
resolve_layout(Some(s), "test", "op").expect("layout resolves")
}
#[test]
fn required_broadcast_projects_requires_broadcast_not_strided() {
let caps = resolve(&spec(
Some("accepted"),
Some("accepted"),
Some("required"),
Some("rejected"),
Some("rejected"),
))
.project();
assert!(
caps.requires_broadcast,
"required broadcast → requires_broadcast"
);
assert!(
!caps.strided_input,
"Required is not Accepted → not strided-generic"
);
let generic = resolve(&spec(
Some("accepted"),
Some("accepted"),
Some("accepted"),
Some("rejected"),
Some("rejected"),
))
.project();
assert!(
!generic.requires_broadcast,
"accepted broadcast is generic, not baked"
);
assert!(generic.strided_input);
}
#[test]
fn project_kernel_caps_requires_broadcast_if_any_operand_does() {
let baked = resolve(&spec(
Some("accepted"),
Some("accepted"),
Some("required"),
Some("rejected"),
Some("rejected"),
));
let generic = resolve(&spec(
Some("accepted"),
Some("accepted"),
Some("accepted"),
Some("rejected"),
Some("rejected"),
));
assert!(
project_kernel_caps(&[generic, baked]).requires_broadcast,
"any baked operand ⇒ the kernel is baked-broadcast",
);
assert!(
!project_kernel_caps(&[generic, generic]).requires_broadcast,
"all-generic ⇒ dense-safe",
);
}
#[test]
fn projection_truth_table() {
let r = resolve(&spec(
Some("accepted"),
Some("accepted"),
Some("accepted"),
Some("rejected"),
Some("rejected"),
));
assert!(r.project().strided_input);
let r = resolve(&spec(
Some("required"),
Some("rejected"),
Some("accepted"),
Some("rejected"),
Some("rejected"),
));
assert!(!r.project().strided_input);
let r = resolve(&spec(None, Some("accepted"), Some("rejected"), None, None));
assert!(!r.project().strided_input);
let r = resolve(&spec(
Some("required"),
Some("rejected"),
Some("rejected"),
Some("rejected"),
Some("rejected"),
));
assert!(!r.project().strided_input);
}
#[test]
fn start_offset_does_not_flip_strided_input() {
let r = resolve(&spec(
Some("required"),
Some("rejected"),
Some("rejected"),
Some("accepted"),
Some("rejected"),
));
assert!(!r.project().strided_input);
assert_eq!(r.start_offset, Tri::Accepted);
}
#[test]
fn reverse_strides_retained_not_projected() {
let r = resolve(&spec(
None,
Some("accepted"),
Some("accepted"),
None,
Some("accepted"),
));
assert!(r.project().strided_input);
assert_eq!(r.reverse_strides, Tri::Accepted);
}
#[test]
fn elementwise_binary_contract_projection_is_false() {
let r = resolve(&spec(
Some("required"),
Some("rejected"),
Some("rejected"),
Some("rejected"),
Some("rejected"),
));
assert!(!r.project().strided_input);
assert_eq!(r.contiguous, Tri::Required);
}
#[test]
fn absent_layout_is_conservative_false() {
let r = resolve_layout(None, "test", "op").unwrap();
assert!(!r.project().strided_input);
assert_eq!(r.strided, Tri::Rejected);
}
#[test]
fn unknown_tristate_value_is_typed_error() {
let s = spec(Some("maybe"), None, None, None, None);
let err = resolve_layout(Some(&s), "test", "op").expect_err("bad value errors");
assert!(matches!(err, FkcError::BadLayoutFlag { .. }), "got {err:?}");
}
#[test]
fn kernel_caps_is_and_of_operands() {
let strided = resolve(&spec(None, Some("accepted"), Some("accepted"), None, None));
let contig = resolve(&spec(
Some("required"),
Some("rejected"),
Some("rejected"),
None,
None,
));
assert!(project_kernel_caps(&[strided, strided]).strided_input);
assert!(!project_kernel_caps(&[strided, contig]).strided_input);
}
#[test]
fn is_generic_contract_classifies_permissive_vs_specialized() {
let generic = resolve(&spec(
Some("n/a"),
Some("accepted"),
Some("accepted"),
None,
None,
));
assert!(is_generic_contract(&[generic, generic]));
let contig = resolve(&spec(
Some("required"),
Some("rejected"),
Some("rejected"),
None,
None,
));
assert!(!is_generic_contract(&[generic, contig]));
assert!(!is_generic_contract(&[contig, contig]));
assert!(!is_generic_contract(&[]));
}
#[test]
fn kernel_caps_field_set_is_pinned_to_the_forward_extension_hook() {
let KernelCaps {
strided_input: _,
requires_broadcast: _,
} = KernelCaps::default();
}
}