use super::structure_key::FdxOperandDesc;
use fuel_ir::DType;
const MAX_RANK: usize = 8;
const MAX_OPERANDS: usize = 8;
const MAX_TOKEN_BYTES: usize = 4096;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReduceAxes {
All,
TrailingAxis,
Keepdim(u8),
}
impl ReduceAxes {
fn mask(self, rank: usize) -> u8 {
let full_mask: u8 = if rank == 0 {
0
} else if rank >= 8 {
0xFF
} else {
(1u8 << rank) - 1
};
let trailing_bit: u8 = if rank == 0 { 0 } else { 1u8 << (rank - 1) };
match self {
ReduceAxes::All => full_mask,
ReduceAxes::TrailingAxis => trailing_bit,
ReduceAxes::Keepdim(m) => m,
}
}
}
fn reduce_field(mask: u8, rank: usize) -> String {
let full_mask: u8 = if rank >= 8 { 0xFF } else { (1u8 << rank) - 1 };
let trailing_bit: u8 = if rank == 0 { 0 } else { 1u8 << (rank - 1) };
if mask == 0 {
"-".to_string()
} else if rank > 0 && mask == full_mask {
"rall".to_string()
} else if rank > 0 && mask == trailing_bit {
"rlast".to_string()
} else {
format!("x{mask:02x}")
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GemMathPrecision {
BitStable,
ReducedMantissa,
}
impl GemMathPrecision {
fn code(self) -> &'static str {
match self {
GemMathPrecision::BitStable => "st",
GemMathPrecision::ReducedMantissa => "rm",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GemCell {
pub m: i64,
pub n: i64,
pub k: i64,
pub batch: Option<i64>,
pub weight_dtype: DType,
pub acc_dtype: DType,
pub out_dtype: DType,
pub math_precision: GemMathPrecision,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccMp {
pub acc_dtype: DType,
pub math_precision: GemMathPrecision,
}
impl AccMp {
const DEFAULT_MP: GemMathPrecision = GemMathPrecision::BitStable;
fn deviates_from(&self, compute_dtype: DType) -> bool {
self.acc_dtype != compute_dtype || self.math_precision != Self::DEFAULT_MP
}
fn field(&self) -> Option<String> {
Some(format!(
"{}/{}",
sk5_dtype_token(self.acc_dtype)?,
self.math_precision.code(),
))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CellClass {
DenseContraction,
IndexedAccess,
Reduction,
ElementwiseOne,
ElementwiseTwo,
ElementwiseThree,
Other,
}
impl CellClass {
pub fn code(self) -> &'static str {
match self {
CellClass::DenseContraction => "ctr",
CellClass::IndexedAccess => "ndx",
CellClass::Reduction => "rdc",
CellClass::ElementwiseOne => "ew1",
CellClass::ElementwiseTwo => "ew2",
CellClass::ElementwiseThree => "ew3",
CellClass::Other => "oth",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DequantForm {
Linear,
Codebook,
}
impl DequantForm {
fn code(self) -> &'static str {
match self {
DequantForm::Linear => "dl",
DequantForm::Codebook => "dc",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OperandRole {
Input,
InPlace,
Output,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct CellOperand {
pub desc: FdxOperandDesc,
pub role: OperandRole,
pub dequant_form: Option<DequantForm>,
}
impl CellOperand {
fn new(desc: FdxOperandDesc, role: OperandRole) -> Self {
Self {
desc,
role,
dequant_form: None,
}
}
pub fn input(desc: FdxOperandDesc) -> Self {
Self::new(desc, OperandRole::Input)
}
pub fn in_place(desc: FdxOperandDesc) -> Self {
Self::new(desc, OperandRole::InPlace)
}
pub fn output(desc: FdxOperandDesc) -> Self {
Self::new(desc, OperandRole::Output)
}
#[must_use]
pub fn with_dequant_form(mut self, form: DequantForm) -> Self {
self.dequant_form = Some(form);
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct DeclaredFacts {
pub contraction: Option<GemCell>,
pub index_operand: Option<usize>,
pub reduced_axes: Option<ReduceAxes>,
}
impl DeclaredFacts {
pub fn none() -> Self {
Self::default()
}
pub fn contraction(cell: GemCell) -> Self {
Self {
contraction: Some(cell),
..Self::default()
}
}
pub fn indexed(slot: usize) -> Self {
Self {
index_operand: Some(slot),
..Self::default()
}
}
pub fn reduction(axes: ReduceAxes) -> Self {
Self {
reduced_axes: Some(axes),
..Self::default()
}
}
}
fn is_valid_namespaced_target(target: &str) -> bool {
let mut parts = target.split(':');
let (Some(ns), Some(rest)) = (parts.next(), parts.next()) else {
return false;
};
if parts.next().is_some() {
return false; }
if ns.is_empty() || rest.is_empty() {
return false;
}
!target
.chars()
.any(|c| c == '|' || c == ';' || c == '/' || c.is_whitespace() || c.is_control())
}
fn iteration_frame(operands: &[CellOperand]) -> Option<Vec<i64>> {
if operands.is_empty() || operands.len() > MAX_OPERANDS {
return None; }
if operands
.iter()
.any(|o| o.desc.shape.len() > MAX_RANK || o.desc.shape.len() != o.desc.strides.len())
{
return None; }
let rank = operands
.iter()
.map(|o| o.desc.shape.len())
.max()
.unwrap_or(0);
let mut frame = vec![0i64; rank];
for o in operands {
let off = rank - o.desc.shape.len();
for (i, &e) in o.desc.shape.iter().enumerate() {
frame[off + i] = frame[off + i].max(e);
}
}
Some(frame)
}
pub fn derive_cell_class(facts: &DeclaredFacts, operands: &[CellOperand]) -> Option<CellClass> {
let frame = iteration_frame(operands)?;
cell_class_in_frame(facts, operands, &frame)
}
fn is_canonical_order(operands: &[CellOperand]) -> bool {
match operands.iter().position(|o| o.role == OperandRole::Output) {
Some(p) => operands[p..].iter().all(|o| o.role == OperandRole::Output),
None => true,
}
}
fn is_pointwise(operands: &[CellOperand], frame: &[i64]) -> bool {
let rank = frame.len();
operands.iter().all(|o| {
let off = rank - o.desc.shape.len();
let extents_ok = frame.iter().enumerate().all(|(i, &f)| {
let e = if i < off { 1 } else { o.desc.shape[i - off] };
e == 1 || e == f
});
let written = matches!(o.role, OperandRole::Output | OperandRole::InPlace);
let written_ok = !written || o.desc.shape == frame;
extents_ok && written_ok
})
}
fn cell_class_in_frame(
facts: &DeclaredFacts,
operands: &[CellOperand],
frame: &[i64],
) -> Option<CellClass> {
if !is_canonical_order(operands) {
return None;
}
if let Some(slot) = facts.index_operand {
match operands.get(slot) {
Some(o) if o.role != OperandRole::Output => {}
_ => return None,
}
}
if facts.contraction.is_some() {
return Some(CellClass::DenseContraction);
}
if facts.index_operand.is_some() {
return Some(CellClass::IndexedAccess);
}
if let Some(axes) = facts.reduced_axes
&& axes.mask(frame.len()) != 0
{
return Some(CellClass::Reduction);
}
let n_inputs = operands
.iter()
.filter(|o| o.role != OperandRole::Output)
.count();
Some(match (is_pointwise(operands, frame), n_inputs) {
(true, 1) => CellClass::ElementwiseOne,
(true, 2) => CellClass::ElementwiseTwo,
(true, 3) => CellClass::ElementwiseThree,
_ => CellClass::Other,
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum DeriveDecline {
OperandCount,
BadShape,
ScaleDtypeAtOperand0,
ReservedDtype,
UnknownDtype,
CollapsedReduction,
ReducedAxisOutOfRange,
IndexSlotOutOfRange,
NotCanonicalOrder,
InvalidTarget,
AccMpOnContraction,
InvalidContraction,
TokenTooLong,
}
impl DeriveDecline {
pub fn wire(self) -> &'static str {
match self {
DeriveDecline::OperandCount => "OperandCount",
DeriveDecline::BadShape => "BadShape",
DeriveDecline::ScaleDtypeAtOperand0 => "ScaleDtypeAtOperand0",
DeriveDecline::ReservedDtype => "ReservedDtype",
DeriveDecline::UnknownDtype => "UnknownDtype",
DeriveDecline::CollapsedReduction => "CollapsedReduction",
DeriveDecline::ReducedAxisOutOfRange => "ReducedAxisOutOfRange",
DeriveDecline::IndexSlotOutOfRange => "IndexSlotOutOfRange",
DeriveDecline::NotCanonicalOrder => "NotCanonicalOrder",
DeriveDecline::InvalidTarget => "InvalidTarget",
DeriveDecline::AccMpOnContraction => "AccMpOnContraction",
DeriveDecline::InvalidContraction => "InvalidContraction",
DeriveDecline::TokenTooLong => "TokenTooLong",
}
}
}
fn dtype_decline(dt: DType) -> DeriveDecline {
if matches!(dt, DType::F8E6M2) {
DeriveDecline::ReservedDtype
} else {
DeriveDecline::UnknownDtype
}
}
pub fn derive_structure_key_token(
facts: &DeclaredFacts,
operands: &[CellOperand],
target: &str,
acc_mp: Option<AccMp>,
) -> Option<String> {
derive_structure_key_token_typed(facts, operands, target, acc_mp).ok()
}
struct ValidCell {
frame: Vec<i64>,
dtype: &'static str,
class: CellClass,
}
fn validate_shape(operands: &[CellOperand]) -> Result<Vec<i64>, DeriveDecline> {
if operands.is_empty() || operands.len() > MAX_OPERANDS {
return Err(DeriveDecline::OperandCount);
}
iteration_frame(operands).ok_or(DeriveDecline::BadShape)
}
fn validate_primary(first: DType, target: &str) -> Result<&'static str, DeriveDecline> {
if matches!(first, DType::F8E8M0) {
return Err(DeriveDecline::ScaleDtypeAtOperand0);
}
let dtype = sk5_dtype_token(first).ok_or_else(|| dtype_decline(first))?;
if !is_valid_namespaced_target(target) {
return Err(DeriveDecline::InvalidTarget); }
Ok(dtype)
}
fn validate_roles(facts: &DeclaredFacts, operands: &[CellOperand]) -> Result<(), DeriveDecline> {
if !is_canonical_order(operands) {
return Err(DeriveDecline::NotCanonicalOrder);
}
match facts.index_operand.map(|slot| operands.get(slot)) {
Some(Some(o)) if o.role != OperandRole::Output => Ok(()),
Some(_) => Err(DeriveDecline::IndexSlotOutOfRange),
None => Ok(()),
}
}
fn validate_cell(
facts: &DeclaredFacts,
operands: &[CellOperand],
target: &str,
acc_mp: Option<AccMp>,
) -> Result<ValidCell, DeriveDecline> {
if acc_mp.is_some() && facts.contraction.is_some() {
return Err(DeriveDecline::AccMpOnContraction);
}
let frame = validate_shape(operands)?;
let first = operands.first().ok_or(DeriveDecline::OperandCount)?;
let dtype = validate_primary(first.desc.dtype, target)?;
validate_roles(facts, operands)?;
let class = cell_class_in_frame(facts, operands, &frame).ok_or(DeriveDecline::BadShape)?;
Ok(ValidCell {
frame,
dtype,
class,
})
}
fn check_reduction(
class: CellClass,
facts: &DeclaredFacts,
operands: &[CellOperand],
rank: usize,
) -> Result<(), DeriveDecline> {
let (CellClass::Reduction, Some(axes)) = (class, facts.reduced_axes) else {
return Ok(());
};
let mask = axes.mask(rank);
if rank < 8 && (u32::from(mask) >> rank) != 0 {
return Err(DeriveDecline::ReducedAxisOutOfRange);
}
let collapsed = |o: &CellOperand| {
o.desc.shape.len() != rank
|| (0..rank).any(|d| (mask >> d) & 1 == 1 && o.desc.shape[d] != 1)
};
if operands
.iter()
.any(|o| o.role == OperandRole::Output && collapsed(o))
{
return Err(DeriveDecline::CollapsedReduction);
}
Ok(())
}
fn index_width(operands: &[CellOperand]) -> &'static str {
let max_touched: i128 = operands
.iter()
.map(|o| {
o.desc
.strides
.iter()
.zip(o.desc.shape.iter())
.map(|(&s, &e)| i128::from(s.unsigned_abs()) * i128::from(e.max(1) - 1))
.sum::<i128>()
})
.max()
.unwrap_or(0);
if max_touched >= (1i128 << 31) {
"ix64"
} else {
"ix32"
}
}
fn work_class(frame: &[i64]) -> &'static str {
let work_elems: i128 = frame.iter().map(|&e| i128::from(e)).product();
if work_elems <= 32 {
"warp"
} else if work_elems <= 1024 {
"block"
} else {
"grid"
}
}
fn reduced_mask_of(class: CellClass, facts: &DeclaredFacts, rank: usize) -> u8 {
match (class, facts.reduced_axes) {
(CellClass::Reduction, Some(axes)) => axes.mask(rank),
_ => 0,
}
}
pub fn derive_structure_key_token_typed(
facts: &DeclaredFacts,
operands: &[CellOperand],
target: &str,
acc_mp: Option<AccMp>,
) -> Result<String, DeriveDecline> {
let ValidCell {
frame,
dtype,
class,
} = validate_cell(facts, operands, target, acc_mp)?;
let rank = frame.len();
check_reduction(class, facts, operands, rank)?;
let reduced_mask = reduced_mask_of(class, facts, rank);
let innermost_reduced = rank >= 1 && (reduced_mask >> (rank - 1)) & 1 == 1;
let operand_keys: Vec<String> = operands
.iter()
.map(|o| operand_sub_key(o, &frame, innermost_reduced))
.collect();
let mut token = format!(
"sk5|{class}|{dtype}|{target}|{idx}|{work}|r{rank}|{ops}|{reduce}",
class = class.code(),
idx = index_width(operands),
work = work_class(&frame),
ops = operand_keys.join(";"),
reduce = reduce_field(reduced_mask, rank),
);
append_precision_fields(&mut token, facts, class, acc_mp, operands[0].desc.dtype)?;
if token.len() > MAX_TOKEN_BYTES {
return Err(DeriveDecline::TokenTooLong); }
Ok(token)
}
fn append_precision_fields(
token: &mut String,
facts: &DeclaredFacts,
class: CellClass,
acc_mp: Option<AccMp>,
compute: DType,
) -> Result<(), DeriveDecline> {
if let Some(a) = acc_mp.filter(|a| a.deviates_from(compute)) {
token.push('|');
token.push_str(&a.field().ok_or_else(|| dtype_decline(a.acc_dtype))?);
}
if let (CellClass::DenseContraction, Some(cell)) = (class, facts.contraction) {
token.push('|');
token.push_str(&contraction_field(&cell).ok_or(DeriveDecline::InvalidContraction)?);
}
Ok(())
}
fn contraction_field(cell: &GemCell) -> Option<String> {
let m = size_class(cell.m)?;
let n = size_class(cell.n)?;
let k = size_class(cell.k)?;
let kdiv = div_bucket(cell.k);
let batch = match cell.batch {
Some(b) => format!("/b{}", size_class(b)?),
None => String::new(),
};
let wdt = sk5_dtype_token(cell.weight_dtype)?;
let acc = sk5_dtype_token(cell.acc_dtype)?;
let out = sk5_dtype_token(cell.out_dtype)?;
Some(format!(
"c{m}{n}{k}/{kdiv}{batch}/{wdt}/{acc}/{out}/{mp}",
mp = cell.math_precision.code(),
))
}
fn size_class(extent: i64) -> Option<char> {
if extent < 0 {
return None;
}
Some(if extent <= 8 {
't'
} else if extent <= 128 {
's'
} else if extent <= 2048 {
'm'
} else {
'l'
})
}
pub fn sk5_dtype_token(dt: DType) -> Option<&'static str> {
match dt {
DType::F8E6M2 => None,
other => fuel_ir::sk4_token(other),
}
}
fn operand_sub_key(co: &CellOperand, frame: &[i64], innermost_reduced: bool) -> String {
let o = &co.desc;
let rank = frame.len();
let off = rank - o.shape.len();
let mut mask = 0u8;
for (i, &f) in frame.iter().enumerate().take(rank.min(8)) {
let stride_is_zero = if i < off {
true
} else {
o.strides[i - off] == 0 || o.shape[i - off] == 1
};
if f > 1 && stride_is_zero {
mask |= 1 << i;
}
}
let layout = layout_code(&o.shape, &o.strides);
let (inner_extent, inner_stride) = match o.shape.len().checked_sub(1) {
Some(i) => (o.shape[i], Some(o.strides[i])),
None => (1, None),
};
let div = div_bucket(inner_extent);
let dsz = o.dtype.size_in_bytes();
let vec = if layout == "br"
|| innermost_reduced
|| dsz == 0
|| o.align_bytes == 0
|| inner_stride != Some(1)
{
"v1"
} else {
let mut picked = "v1";
for &l in &[8i64, 4, 2] {
let vbytes = (l as u32) * (dsz as u32);
if vbytes <= 16
&& o.align_bytes.is_multiple_of(vbytes)
&& inner_extent >= l
&& inner_extent % l == 0
{
picked = match l {
8 => "v8",
4 => "v4",
_ => "v2",
};
break;
}
}
picked
};
let mut key = format!(
"{layout}/{mask:02x}/{vec}/{div}/{flip}",
flip = if o.flipped { "r" } else { "f" },
);
if let Some(form) = co.dequant_form {
key.push('/');
key.push_str(form.code());
}
key
}
fn layout_code(ext: &[i64], strides: &[i64]) -> &'static str {
if ext.iter().zip(strides).any(|(&e, &s)| e > 1 && s == 0) {
return "br";
}
let mut p: i128 = 1;
let mut contiguous = true;
for i in (0..ext.len()).rev() {
let e = ext[i];
if e <= 1 {
continue; }
if i128::from(strides[i].unsigned_abs()) != p {
contiguous = false;
break;
}
p *= i128::from(e);
}
if contiguous {
return "co";
}
if let Some(i) = (0..ext.len()).rev().find(|&i| ext[i] > 1)
&& strides[i].unsigned_abs() == 1
{
return "ic";
}
"st"
}
fn div_bucket(e: i64) -> &'static str {
if e >= 16 && e % 16 == 0 {
"d16"
} else if e >= 8 && e % 8 == 0 {
"d8"
} else if e >= 4 && e % 4 == 0 {
"d4"
} else if e >= 2 && e % 2 == 0 {
"d2"
} else {
"da"
}
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::{DType, Layout, Shape, StrideVec};
fn co(dims: &[usize], dtype: DType) -> FdxOperandDesc {
FdxOperandDesc::from_layout(&Layout::contiguous(Shape::from_dims(dims)), dtype)
}
fn f32c(dims: &[usize]) -> FdxOperandDesc {
co(dims, DType::F32)
}
fn canon(descs: &[FdxOperandDesc]) -> Vec<CellOperand> {
let n = descs.len();
descs
.iter()
.enumerate()
.map(|(i, d)| {
if n >= 2 && i == n - 1 {
CellOperand::output(d.clone())
} else {
CellOperand::input(d.clone())
}
})
.collect()
}
fn derive(facts: DeclaredFacts, descs: &[FdxOperandDesc], target: &str) -> Option<String> {
derive_structure_key_token(&facts, &canon(descs), target, None)
}
fn derive_acc_mp(
facts: DeclaredFacts,
descs: &[FdxOperandDesc],
target: &str,
acc_mp: Option<AccMp>,
) -> Option<String> {
derive_structure_key_token(&facts, &canon(descs), target, acc_mp)
}
#[test]
fn every_spelled_dtype_is_also_derivable_at_operand_zero() {
let mut spelled = 0usize;
let mut not_derivable: Vec<(DType, &'static str)> = Vec::new();
let mut mismatched: Vec<(DType, &'static str, String)> = Vec::new();
for dt in DType::ALL.iter().copied() {
let Some(spelling) = sk5_dtype_token(dt) else {
continue; };
spelled += 1;
let ops = [co(&[4096], dt), co(&[4096], dt)];
match derive(DeclaredFacts::none(), &ops, "cuda:sm89") {
None => not_derivable.push((dt, spelling)),
Some(token) => {
let got = token.split('|').nth(2).unwrap_or("<missing>");
if got != spelling {
mismatched.push((dt, spelling, got.to_string()));
}
}
}
}
assert_eq!(
spelled, 14,
"expected 14 sk5-spelled dtypes; got {spelled} — this test would be vacuous at 0",
);
let known_positional_exemptions = [DType::F8E8M0];
let unexpected_gaps: Vec<(DType, &'static str)> = not_derivable
.iter()
.copied()
.filter(|(dt, _)| !known_positional_exemptions.contains(dt))
.collect();
assert!(
unexpected_gaps.is_empty(),
"spelled but NOT derivable at operand 0, and NOT a known positional \
exemption: {unexpected_gaps:?} — a dtype with a seam token that no \
cell can carry is a real capability gap",
);
let exempted_but_derivable: Vec<DType> = known_positional_exemptions
.iter()
.copied()
.filter(|dt| !not_derivable.iter().any(|(nd, _)| nd == dt))
.collect();
assert!(
exempted_but_derivable.is_empty(),
"{exempted_but_derivable:?} are expected to DECLINE at operand 0 \
(KISS-CLASSIFY §3.2 positional restriction) but derived a token — \
the positional decline regressed",
);
assert!(
mismatched.is_empty(),
"derived token's dtype field disagrees with the spelling: {mismatched:?}",
);
}
fn gem_f32(m: i64, n: i64, k: i64) -> GemCell {
GemCell {
m,
n,
k,
batch: None,
weight_dtype: DType::F32,
acc_dtype: DType::F32,
out_dtype: DType::F32,
math_precision: GemMathPrecision::BitStable,
}
}
#[test]
fn fuel_derives_relu_add_sk5_token_byte_for_byte() {
let op = f32c(&[4096]);
let token = derive(
DeclaredFacts::none(),
&[op.clone(), op.clone(), op],
"cuda:sm89",
)
.expect("relu_add f32 must derive a token");
assert_eq!(
token,
"sk5|ew2|f32|cuda:sm89|ix32|grid|r1|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-"
);
}
#[test]
fn sk5_prefix_and_a_derived_class_code_on_every_token_class() {
const CLASS_CODES: [&str; 7] = ["ctr", "ndx", "rdc", "ew1", "ew2", "ew3", "oth"];
const RETIRED_FAMILY_CODES: [&str; 19] = [
"une", "bin", "ter", "scn", "red", "gem", "nrm", "cnv", "pol", "idx", "shp", "srt",
"fft", "lin", "rnd", "seg", "sft", "att", "los",
];
let op = f32c(&[4096]);
let facts = [
DeclaredFacts::none(),
DeclaredFacts::reduction(ReduceAxes::All),
DeclaredFacts::reduction(ReduceAxes::TrailingAxis),
DeclaredFacts::reduction(ReduceAxes::Keepdim(0x01)),
DeclaredFacts::contraction(gem_f32(8, 4096, 4096)),
DeclaredFacts::indexed(0),
];
for f in facts {
let token = derive(f, std::slice::from_ref(&op), "cuda:sm89")
.unwrap_or_else(|| panic!("{f:?} must derive"));
assert!(token.starts_with("sk5|"), "{token} lacks the sk5 prefix");
assert!(!token.contains("sk2") && !token.contains("sk4"), "{token}");
let class = token.split('|').nth(1).unwrap();
assert!(
CLASS_CODES.contains(&class),
"{token}: `{class}` not a class"
);
assert!(
!RETIRED_FAMILY_CODES.contains(&class),
"{token}: retired op-family code `{class}` in field 2"
);
}
}
#[test]
fn non_gem_tokens_tail_identical_to_the_sk2_era_battery() {
let cases: [(Vec<FdxOperandDesc>, &str, &str); 4] = [
(
vec![f32c(&[4096]), f32c(&[4096]), f32c(&[4096])],
"sk2|bin|f32|cuda:sm89|ix32|grid|r1|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-",
"sk5|ew2|f32|cuda:sm89|ix32|grid|r1|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-",
),
(
vec![f32c(&[7])],
"sk2|bin|f32|cuda:sm89|ix32|warp|r1|co/00/v1/da/f|-",
"sk5|ew1|f32|cuda:sm89|ix32|warp|r1|co/00/v1/da/f|-",
),
(
vec![co(&[4096], DType::I16)],
"sk2|bin|s16|cuda:sm89|ix32|grid|r1|co/00/v8/d16/f|-",
"sk5|ew1|i16|cuda:sm89|ix32|grid|r1|co/00/v8/d16/f|-",
),
(
vec![f32c(&[128, 256])],
"sk2|bin|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f|-",
"sk5|ew1|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f|-",
),
];
for (ops, sk2, expect) in cases {
let token = derive(DeclaredFacts::none(), &ops, "cuda:sm89").expect("derives");
assert_eq!(token, expect);
let tail = |t: &str| t.split('|').skip(3).collect::<Vec<_>>().join("|");
assert_eq!(
tail(&token),
tail(sk2),
"fields 4.. moved between the sk2-era battery and sk5"
);
let dt = |t: &str| t.split('|').nth(2).unwrap().to_string();
assert_eq!(dt(&token), dt(sk2).replace("s16", "i16"));
}
}
#[test]
fn kiss_a1_gem_skinny_decode_golden() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
let cell = DeclaredFacts::contraction(gem_f32(8, 4096, 4096));
let cuda = derive(cell, &ops, "cuda:sm89").expect("derives");
assert_eq!(
cuda,
"sk5|ctr|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-|ctll/d16/f32/f32/f32/st"
);
let vk = derive(cell, &ops, "vulkan:spirv1.6").expect("derives");
assert_eq!(
vk,
"sk5|ctr|f32|vulkan:spirv1.6|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-|ctll/d16/f32/f32/f32/st"
);
}
#[test]
fn namespaced_target_is_a_byte_exact_passthrough_for_any_vocabulary() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
let cell = DeclaredFacts::contraction(gem_f32(8, 4096, 4096));
let targets = [
"cuda:sm89",
"vulkan:spirv1.6",
"vulkan:sg64.ops-abr.arith-f16.cm-none",
"vulkan:sg64.ops-abr.arith-f16.cm-none.cv-none",
"zzfuture:A-b.C_9.MixedCase.trailing-Z",
];
for target in targets {
let token = derive(cell, &ops, target)
.unwrap_or_else(|| panic!("`{target}` is namespaced and must derive"));
let fields: Vec<&str> = token.split('|').collect();
assert!(
fields.len() > 3,
"`{target}`: token has {} fields, no target coordinate to check: {token}",
fields.len()
);
assert_eq!(
fields[3], target,
"target must survive byte-exact as field 4 (§6.8-0002); token: {token}"
);
}
assert_eq!(
derive(cell, &ops, "vulkansg64.ops-abr.arith-f16.cm-none"),
None,
"a non-namespaced target must decline, not pass through"
);
}
#[test]
fn malformed_namespaced_targets_decline() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
let cell = DeclaredFacts::contraction(gem_f32(8, 4096, 4096));
for bad in [
"a:b:c", ":x", "x:", "cuda:sm|89", "cuda:sm;89", "cuda:sm/89", "cuda: sm89", "cuda:sm89\t", ] {
assert_eq!(
derive(cell, &ops, bad),
None,
"malformed target `{bad:?}` must decline, not pass bare contains(':') check",
);
}
}
#[test]
fn an_oversized_token_declines_rather_than_emitting_unbounded() {
let ops = [f32c(&[4, 8])];
let cell = DeclaredFacts::none();
let huge_target = format!("cuda:{}", "a".repeat(MAX_TOKEN_BYTES));
assert_eq!(
derive(cell, &ops, &huge_target),
None,
"a token whose target alone exceeds MAX_TOKEN_BYTES must decline",
);
assert!(
derive(cell, &ops, "cuda:sm89").is_some(),
"positive control: the cell itself is derivable with a short target",
);
}
#[test]
fn ctr_batched_cell_golden() {
let ops = [f32c(&[256, 4096]), f32c(&[4096, 4096]), f32c(&[256, 4096])];
let batched = GemCell {
batch: Some(256),
..gem_f32(256, 4096, 4096)
};
let token =
derive(DeclaredFacts::contraction(batched), &ops, "cuda:sm90").expect("derives");
assert_eq!(
token,
"sk5|ctr|f32|cuda:sm90|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-|cmll/d16/bm/f32/f32/f32/st"
);
let plain = derive(
DeclaredFacts::contraction(gem_f32(256, 4096, 4096)),
&ops,
"cuda:sm90",
)
.expect("derives");
assert_eq!(plain, token.replace("/bm/", "/"));
}
#[test]
fn ctr_simt_f32_vs_tf32_distinct_by_mp() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
let simt = derive(
DeclaredFacts::contraction(gem_f32(8, 4096, 4096)),
&ops,
"cuda:sm90",
)
.expect("derives");
let tf32_cell = GemCell {
math_precision: GemMathPrecision::ReducedMantissa,
..gem_f32(8, 4096, 4096)
};
let tf32 =
derive(DeclaredFacts::contraction(tf32_cell), &ops, "cuda:sm90").expect("derives");
assert_eq!(
simt,
"sk5|ctr|f32|cuda:sm90|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-|ctll/d16/f32/f32/f32/st"
);
assert_eq!(
tf32,
"sk5|ctr|f32|cuda:sm90|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-|ctll/d16/f32/f32/f32/rm"
);
assert_ne!(simt, tf32, "SIMT-f32 and TF32 must not collide");
}
#[test]
fn ctr_mixed_precision_fp8_golden() {
let f8 = |dims: &[usize]| {
FdxOperandDesc::from_layout(
&Layout::new(
Shape::from_dims(dims),
Layout::contiguous(Shape::from_dims(dims))
.stride()
.iter()
.copied()
.collect::<StrideVec>(),
4,
),
DType::F8E4M3,
)
};
let ops = [f8(&[8, 4096]), f8(&[4096, 4096]), f8(&[8, 4096])];
let cell = GemCell {
weight_dtype: DType::F8E4M3,
acc_dtype: DType::F32,
out_dtype: DType::F16,
..gem_f32(8, 4096, 4096)
};
let token = derive(DeclaredFacts::contraction(cell), &ops, "cuda:sm90").expect("derives");
assert_eq!(
token,
"sk5|ctr|f8e4m3fn|cuda:sm90|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-|ctll/d16/f8e4m3fn/f32/f16/st"
);
let twin = GemCell {
out_dtype: DType::F32,
..cell
};
let twin_token =
derive(DeclaredFacts::contraction(twin), &ops, "cuda:sm90").expect("derives");
assert_ne!(
token, twin_token,
"mixed-precision FP8 cells must not collide under sk5"
);
}
#[test]
fn gem_declines_unmappable_or_invalid() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
let mx_weight = GemCell {
weight_dtype: DType::F4,
..gem_f32(8, 4096, 4096)
};
assert_eq!(
derive(DeclaredFacts::contraction(mx_weight), &ops, "cuda:sm89"),
None
);
let negative_m = gem_f32(-1, 4096, 4096);
assert_eq!(
derive(DeclaredFacts::contraction(negative_m), &ops, "cuda:sm89"),
None
);
let negative_batch = GemCell {
batch: Some(-2),
..gem_f32(8, 4096, 4096)
};
assert_eq!(
derive(
DeclaredFacts::contraction(negative_batch),
&ops,
"cuda:sm89"
),
None
);
}
#[test]
fn gem_size_class_ladder() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
for (m, n, k, expect) in [
(8i64, 9, 129, "ctsm/da"), (128, 2048, 2049, "csml/da"), (1, 129, 24, "ctms/d8"), (3000, 8, 4096, "cltl/d16"),
] {
let token = derive(
DeclaredFacts::contraction(gem_f32(m, n, k)),
&ops,
"cuda:sm89",
)
.expect("derives");
let field = token.rsplit('|').next().unwrap();
assert!(
field.starts_with(expect),
"gem({m},{n},{k}) contraction field {field} != {expect}…"
);
}
}
#[test]
fn retired_and_reserved_spellings_never_emitted() {
const ALL: [DType; 17] = [
DType::U8,
DType::I8,
DType::U32,
DType::I16,
DType::I32,
DType::I64,
DType::BF16,
DType::F16,
DType::F32,
DType::F64,
DType::F8E4M3,
DType::F8E5M2,
DType::F6E2M3,
DType::F6E3M2,
DType::F4,
DType::F8E8M0,
DType::F8E6M2,
];
let assert_clean = |token: &str| {
assert!(!token.contains("f32s"), "retired f32s spelling in {token}");
assert!(!token.contains("fnuz"), "reserved fnuz spelling in {token}");
assert!(
!token.contains("|e4m3|"),
"retired bare e4m3 (primary) in {token}"
);
assert!(
!token.contains("/e4m3/"),
"retired bare e4m3 (gem group) in {token}"
);
assert!(
!token.contains("|e4m3fn|"),
"retired unprefixed e4m3fn (primary) in {token}"
);
assert!(
!token.contains("/e4m3fn/"),
"retired unprefixed e4m3fn (gem group) in {token}"
);
assert!(
!token.contains("|e5m2|"),
"retired unprefixed e5m2 (primary) in {token}"
);
assert!(
!token.contains("/e5m2/"),
"retired unprefixed e5m2 (gem group) in {token}"
);
assert!(
!token.contains("|s8|") && !token.contains("/s8/"),
"retired sk3 s8 spelling in {token}"
);
assert!(
!token.contains("|s16|") && !token.contains("/s16/"),
"retired sk3 s16 spelling in {token}"
);
};
for dt in ALL {
if let Some(token) = derive(DeclaredFacts::none(), &[co(&[4096], dt)], "cuda:sm89") {
assert_clean(&token);
}
let cell = GemCell {
weight_dtype: dt,
acc_dtype: dt,
out_dtype: dt,
..gem_f32(8, 4096, 4096)
};
if let Some(token) = derive(
DeclaredFacts::contraction(cell),
&[
co(&[8, 4096], dt),
co(&[4096, 4096], dt),
co(&[8, 4096], dt),
],
"cuda:sm89",
) {
assert_clean(&token);
}
}
let token = derive(
DeclaredFacts::none(),
&[co(&[4096], DType::F8E4M3)],
"cuda:sm89",
)
.expect("f8e4m3fn derives");
assert!(token.starts_with("sk5|ew1|f8e4m3fn|"), "got {token}");
}
#[test]
fn mx_scale_and_reserved_dtypes_decline_at_operand_zero_for_distinct_reasons() {
assert_eq!(
derive(
DeclaredFacts::none(),
&[co(&[4096], DType::F8E8M0)],
"cuda:sm89"
),
None,
"F8E8M0 at operand 0 must typed-decline — a ratified MX scale, valid \
only as a sibling operand (KISS-CLASSIFY §3.2)",
);
assert!(
sk5_dtype_token(DType::F8E8M0).is_some(),
"F8E8M0 must remain a spellable token — the decline above is \
positional (operand-0 usage), not a vocabulary exclusion",
);
assert_eq!(
derive(
DeclaredFacts::none(),
&[co(&[4096], DType::F8E6M2)],
"cuda:sm89"
),
None,
"F8E6M2 is reserved at sk5 (KISS #517) and must typed-decline",
);
assert_eq!(
sk5_dtype_token(DType::F8E6M2),
None,
"F8E6M2 has no sk5 spelling: reserved, never emitted",
);
assert_eq!(
fuel_ir::sk4_token(DType::F8E6M2),
Some("f8e6m2"),
"fuel-ir's dtype table moved: the sk5 reservation is no longer \
deriver-local — reconcile with the GAP row"
);
for dt in [DType::F6E2M3, DType::F6E3M2, DType::F4] {
assert_eq!(
derive(DeclaredFacts::none(), &[co(&[4096], dt)], "cuda:sm89"),
None,
"{dt:?} is NOT in the §6.1 set and must still typed-decline",
);
assert!(
sk5_dtype_token(dt).is_none(),
"{dt:?} is a vocabulary exclusion, not a positional one",
);
}
}
#[test]
fn f8e6m2_is_reserved_at_sk5_in_every_keyed_position() {
let ops = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
for (slot, with_reserved) in [
(
"weight",
GemCell {
weight_dtype: DType::F8E6M2,
..gem_f32(8, 4096, 4096)
},
),
(
"accumulator",
GemCell {
acc_dtype: DType::F8E6M2,
..gem_f32(8, 4096, 4096)
},
),
(
"output",
GemCell {
out_dtype: DType::F8E6M2,
..gem_f32(8, 4096, 4096)
},
),
] {
assert!(
derive(
DeclaredFacts::contraction(gem_f32(8, 4096, 4096)),
&ops,
"cuda:sm90"
)
.is_some(),
"control: the f32 contraction must derive"
);
assert_eq!(
derive(DeclaredFacts::contraction(with_reserved), &ops, "cuda:sm90"),
None,
"f8e6m2 in the contraction {slot} coordinate must decline (reserved at sk5)",
);
}
let reduce_ops = [f32c(&[4, 8]), f32c(&[4, 1])];
assert!(
derive_acc_mp(
DeclaredFacts::reduction(ReduceAxes::TrailingAxis),
&reduce_ops,
"cuda:sm89",
Some(AccMp {
acc_dtype: DType::F64,
math_precision: GemMathPrecision::BitStable,
}),
)
.is_some(),
"control: an f64 accumulator derives"
);
assert_eq!(
derive_acc_mp(
DeclaredFacts::reduction(ReduceAxes::TrailingAxis),
&reduce_ops,
"cuda:sm89",
Some(AccMp {
acc_dtype: DType::F8E6M2,
math_precision: GemMathPrecision::BitStable,
}),
),
None,
"f8e6m2 in the (acc + mp) accumulator coordinate must decline",
);
}
#[test]
fn acc_mp_omitted_entirely_when_nothing_deviates() {
let ops = [co(&[4096], DType::F32)];
let plain = derive(DeclaredFacts::reduction(ReduceAxes::All), &ops, "cuda:sm89")
.expect("plain reduction derives");
let diagonal = derive_acc_mp(
DeclaredFacts::reduction(ReduceAxes::All),
&ops,
"cuda:sm89",
Some(AccMp {
acc_dtype: DType::F32,
math_precision: GemMathPrecision::BitStable,
}),
)
.expect("diagonal reduction derives");
assert_eq!(
diagonal, plain,
"a non-deviating (acc+mp) must produce a BYTE-IDENTICAL token; any difference means the field is not additive and every token a consumer already holds has silently changed meaning",
);
assert!(
!plain.ends_with("|-|-") && !diagonal.ends_with("|-|-"),
"a spurious trailing `|-` means the reduce field's MANDATORY convention was applied to an OMITTED-WHEN-ABSENT field: {diagonal}",
);
}
#[test]
fn acc_mp_emitted_with_both_slots_when_either_deviates() {
let ops = [co(&[4096], DType::F16)];
let derive = |a: AccMp| {
derive_acc_mp(
DeclaredFacts::reduction(ReduceAxes::All),
&ops,
"cuda:sm89",
Some(a),
)
.expect("derives")
};
let acc_only = derive(AccMp {
acc_dtype: DType::F32,
math_precision: GemMathPrecision::BitStable,
});
assert!(acc_only.ends_with("|f32/st"), "got {acc_only}");
let mp_only = derive(AccMp {
acc_dtype: DType::F16,
math_precision: GemMathPrecision::ReducedMantissa,
});
assert!(mp_only.ends_with("|f16/rm"), "got {mp_only}");
let diagonal = derive(AccMp {
acc_dtype: DType::F16,
math_precision: GemMathPrecision::BitStable,
});
assert_ne!(acc_only, diagonal);
assert_ne!(mp_only, diagonal);
assert_ne!(acc_only, mp_only);
assert_eq!(
acc_only.split('|').count(),
10,
"9 base fields + 1: {acc_only}"
);
}
#[test]
fn acc_mp_on_a_gem_cell_declines_rather_than_coexisting() {
let cell = gem_f32(8, 4096, 4096);
let ops = [
co(&[8, 4096], DType::F32),
co(&[4096, 4096], DType::F32),
co(&[8, 4096], DType::F32),
];
assert!(
derive(DeclaredFacts::contraction(cell), &ops, "cuda:sm89").is_some(),
"control: the gem cell must derive without an (acc+mp), or the decline below proves nothing about coexistence",
);
assert_eq!(
derive_acc_mp(
DeclaredFacts::contraction(cell),
&ops,
"cuda:sm89",
Some(AccMp {
acc_dtype: DType::F32,
math_precision: GemMathPrecision::ReducedMantissa
}),
),
None,
"a gem cell carrying BOTH precision fields must decline",
);
}
#[test]
fn declines_rather_than_guessing() {
let bad_dtype = co(&[4096], DType::F4);
assert_eq!(
derive(DeclaredFacts::none(), &[bad_dtype], "cuda:sm89"),
None
);
assert_eq!(
derive(DeclaredFacts::none(), &[f32c(&[4096])], "sm89"),
None
);
assert_eq!(derive(DeclaredFacts::none(), &[], "cuda:sm89"), None);
let nine = vec![f32c(&[4096]); 9];
assert_eq!(derive(DeclaredFacts::none(), &nine, "cuda:sm89"), None);
let mut broken = f32c(&[4096]);
broken.strides = vec![1, 1];
assert_eq!(derive(DeclaredFacts::none(), &[broken], "cuda:sm89"), None);
}
#[test]
fn reduction_vec_width_is_v1_when_innermost_axis_reduced() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::All),
&[f32c(&[4096])],
"cuda:sm89",
)
.expect("reduction must derive");
assert_eq!(
token,
"sk5|rdc|f32|cuda:sm89|ix32|grid|r1|co/00/v1/d16/f|rall"
);
}
#[test]
fn reduction_v1_gate_reads_the_innermost_mask_bit() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::Keepdim(0x01)),
&[f32c(&[128, 256]), f32c(&[1, 256])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|rdc|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f;co/01/v4/d16/f|x01"
);
}
#[test]
fn kiss_a1_reduction_trailing_axis_golden() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::TrailingAxis),
&[f32c(&[4, 8]), f32c(&[4, 1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|rdc|f32|cuda:sm89|ix32|warp|r2|co/00/v1/d8/f;co/02/v1/da/f|rlast"
);
}
#[test]
fn kiss_a1_reduction_all_axes_golden() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::All),
&[f32c(&[4, 8]), f32c(&[1, 1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|rdc|f32|cuda:sm89|ix32|warp|r2|co/00/v1/d8/f;co/03/v1/da/f|rall"
);
}
#[test]
fn kiss_a1_reduction_rank1_all_axes_golden() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::All),
&[f32c(&[8]), f32c(&[1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|rdc|f32|cuda:sm89|ix32|warp|r1|co/00/v1/d8/f;co/01/v1/da/f|rall"
);
}
#[test]
fn kiss_a1_reduction_keepdim_mask_golden() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::Keepdim(0x0a)),
&[f32c(&[2, 4, 3, 5]), f32c(&[2, 1, 3, 1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|rdc|f32|cuda:sm89|ix32|block|r4|co/00/v1/da/f;co/0a/v1/da/f|x0a"
);
}
#[test]
fn trailing_axis_at_rank1_is_rall_not_rlast() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::TrailingAxis),
&[f32c(&[8]), f32c(&[1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token, "sk5|rdc|f32|cuda:sm89|ix32|warp|r1|co/00/v1/d8/f;co/01/v1/da/f|rall",
"a rank-1 TrailingAxis reduction is the same operation as All at rank 1 \
and must derive the same reduce field",
);
}
#[test]
fn keepdim_spanning_every_axis_is_rall_not_the_bitmask() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::Keepdim(0x03)),
&[f32c(&[4, 8]), f32c(&[1, 1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token, "sk5|rdc|f32|cuda:sm89|ix32|warp|r2|co/00/v1/d8/f;co/03/v1/da/f|rall",
"Keepdim(0x03) at rank 2 reduces every axis, same operation as All",
);
}
#[test]
fn keepdim_empty_mask_is_dash_not_x00() {
let token = derive(
DeclaredFacts::reduction(ReduceAxes::Keepdim(0x00)),
&[f32c(&[4, 8]), f32c(&[4, 8])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token, "sk5|ew1|f32|cuda:sm89|ix32|warp|r2|co/00/v4/d8/f;co/00/v4/d8/f|-",
"an empty reduced-axis set is not a reduction (§6.5-0015 rule 3): the cell \
derives the class its shape gives it (`ew1`, one input pointwise), never \
`rdc` with a `-` reduce field (the reader's `RdcWithoutReduce` decline), and \
never x00. Vec width is v4 here, not v1, because nothing is reduced — the v1 \
gate is specifically for a reduced innermost axis, which an empty mask never is",
);
}
#[test]
fn rank0_all_and_keepdim_zero_derive_identical_tokens() {
let all_token = derive(
DeclaredFacts::reduction(ReduceAxes::All),
&[f32c(&[]), f32c(&[])],
"cuda:sm89",
)
.expect("derives");
let trailing_token = derive(
DeclaredFacts::reduction(ReduceAxes::TrailingAxis),
&[f32c(&[]), f32c(&[])],
"cuda:sm89",
)
.expect("derives");
let keepdim_token = derive(
DeclaredFacts::reduction(ReduceAxes::Keepdim(0x00)),
&[f32c(&[]), f32c(&[])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(all_token, keepdim_token, "All vs Keepdim(0) at rank 0");
assert_eq!(
trailing_token, keepdim_token,
"TrailingAxis vs Keepdim(0) at rank 0"
);
}
#[test]
fn reduce_field_keepdim_classification_is_exhaustive_over_rank_and_mask() {
for rank in 0usize..=8 {
let full_mask: u8 = if rank == 0 {
0
} else if rank >= 8 {
0xFF
} else {
(1u8 << rank) - 1
};
for m in 0u32..=0xFF {
let m = m as u8;
if m & !full_mask != 0 {
continue;
}
let field = reduce_field(ReduceAxes::Keepdim(m).mask(rank), rank);
if m == 0 {
assert_eq!(field, "-", "rank {rank} mask {m:#04x} must be '-'");
} else if rank > 0 && m == full_mask {
assert_eq!(field, "rall", "rank {rank} mask {m:#04x} must be 'rall'");
} else if rank > 0 && m == (1u8 << (rank - 1)) {
assert_eq!(field, "rlast", "rank {rank} mask {m:#04x} must be 'rlast'");
} else {
assert_eq!(
field,
format!("x{m:02x}"),
"rank {rank} mask {m:#04x} must be the literal bitmask",
);
}
}
}
}
#[test]
fn trailing_unit_axis_reads_rank_minus_1_inner() {
let token = derive(DeclaredFacts::none(), &[f32c(&[4, 1])], "cuda:sm89").expect("derives");
assert_eq!(token, "sk5|ew1|f32|cuda:sm89|ix32|warp|r2|co/00/v1/da/f|-");
}
#[test]
fn zero_extent_buckets_da() {
let token = derive(DeclaredFacts::none(), &[f32c(&[0])], "cuda:sm89")
.expect("derives (never panics)");
assert_eq!(token, "sk5|ew1|f32|cuda:sm89|ix32|warp|r1|co/00/v1/da/f|-");
}
#[test]
fn flipped_full_reverse_is_contiguous_per_abs_stride() {
let layout = Layout::new(
Shape::from(vec![4usize, 3]),
[-3isize, 1].into_iter().collect::<StrideVec>(),
9,
);
let desc = FdxOperandDesc::from_layout(&layout, DType::F32);
let token = derive(DeclaredFacts::none(), &[desc], "cuda:sm89").expect("derives");
assert_eq!(token, "sk5|ew1|f32|cuda:sm89|ix32|warp|r2|co/00/v1/da/r|-");
}
#[test]
fn alignment_zero_derives_v1() {
let mut desc = f32c(&[4096]);
desc.align_bytes = 0;
let token = derive(DeclaredFacts::none(), &[desc], "cuda:sm89").expect("derives");
assert_eq!(token, "sk5|ew1|f32|cuda:sm89|ix32|grid|r1|co/00/v1/d16/f|-");
}
#[test]
fn work_class_and_masks_use_the_iteration_frame() {
let token = derive(
DeclaredFacts::none(),
&[f32c(&[256]), f32c(&[128, 256]), f32c(&[128, 256])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|ew2|f32|cuda:sm89|ix32|grid|r2|co/01/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-"
);
}
#[test]
fn explicit_unit_broadcast_axis_matches_rank_deficient_own_axes_layout() {
let explicit = FdxOperandDesc::from_layout(
&Layout::new(
Shape::from(vec![1usize, 256]),
[0isize, 1].into_iter().collect::<StrideVec>(),
0,
),
DType::F32,
);
let token = derive(
DeclaredFacts::none(),
&[explicit, f32c(&[128, 256]), f32c(&[128, 256])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|ew2|f32|cuda:sm89|ix32|grid|r2|co/01/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-"
);
}
#[test]
fn frame_padded_dense_vs_strided_rank3_differ_by_own_layout() {
let dense = FdxOperandDesc::from_layout(
&Layout::new(
Shape::from(vec![4usize, 16, 64]),
[1024isize, 64, 1].into_iter().collect::<StrideVec>(),
0,
),
DType::F32,
);
let strided = FdxOperandDesc::from_layout(
&Layout::new(
Shape::from(vec![4usize, 16, 64]),
[2048isize, 128, 2].into_iter().collect::<StrideVec>(),
0,
),
DType::F32,
);
let frame_setter = f32c(&[8, 4, 16, 64]);
let dense_token = derive(
DeclaredFacts::none(),
&[dense, frame_setter.clone()],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
dense_token,
"sk5|ew1|f32|cuda:sm89|ix32|grid|r4|co/01/v4/d16/f;co/00/v4/d16/f|-"
);
let strided_token =
derive(DeclaredFacts::none(), &[strided, frame_setter], "cuda:sm89").expect("derives");
assert_eq!(
strided_token, "sk5|ew1|f32|cuda:sm89|ix32|grid|r4|st/01/v1/d16/f;co/00/v4/d16/f|-",
"a non-forward-unit own stride must derive `st`/`v1` even though \
the frame-padded extents are byte-identical to the dense case"
);
}
#[test]
fn kiss_a1_broadcast_operand_golden() {
let bcast = FdxOperandDesc::from_layout(
&Layout::new(
Shape::from(vec![128usize, 256]),
[0isize, 1].into_iter().collect::<StrideVec>(),
0,
),
DType::F32,
);
let token = derive(
DeclaredFacts::none(),
&[f32c(&[128, 256]), bcast, f32c(&[128, 256])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
token,
"sk5|ew2|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f;br/01/v1/d16/f;co/00/v4/d16/f|-"
);
}
#[test]
fn kiss_a1_bin_canonical_and_inplace_goldens() {
let op = f32c(&[128, 256]);
let canonical = derive(
DeclaredFacts::none(),
&[op.clone(), op.clone(), op.clone()],
"cuda:sm89",
)
.expect("derives");
assert_eq!(
canonical,
"sk5|ew2|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f;co/00/v4/d16/f|-"
);
let inplace = derive_structure_key_token(
&DeclaredFacts::none(),
&[
CellOperand::input(op.clone()),
CellOperand::in_place(op.clone()),
],
"cuda:sm89",
None,
)
.expect("derives");
assert_eq!(
inplace,
"sk5|ew2|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f|-"
);
let in_out =
derive(DeclaredFacts::none(), &[op.clone(), op], "cuda:sm89").expect("derives");
assert_eq!(
in_out,
"sk5|ew1|f32|cuda:sm89|ix32|grid|r2|co/00/v4/d16/f;co/00/v4/d16/f|-"
);
}
#[test]
fn f16_wide_vector_matches_kiss_a1_subkeys() {
let op = co(&[64, 128], DType::F16);
let token = derive(DeclaredFacts::none(), &[op.clone(), op], "cuda:sm89").expect("derives");
assert_eq!(
token,
"sk5|ew1|f16|cuda:sm89|ix32|grid|r2|co/00/v8/d16/f;co/00/v8/d16/f|-"
);
}
fn class_of(facts: DeclaredFacts, descs: &[FdxOperandDesc]) -> Option<CellClass> {
derive_cell_class(&facts, &canon(descs))
}
#[test]
fn declared_facts_are_tried_in_rule_order_before_the_shape_test() {
let same = [f32c(&[8, 16]), f32c(&[8, 16]), f32c(&[8, 16])];
let gem = gem_f32(8, 16, 16);
assert_eq!(
class_of(DeclaredFacts::contraction(gem), &same),
Some(CellClass::DenseContraction)
);
assert_eq!(
class_of(DeclaredFacts::indexed(1), &same),
Some(CellClass::IndexedAccess)
);
assert_eq!(
class_of(DeclaredFacts::reduction(ReduceAxes::All), &same),
Some(CellClass::Reduction)
);
assert_eq!(
class_of(DeclaredFacts::none(), &same),
Some(CellClass::ElementwiseTwo)
);
let mut f = DeclaredFacts::contraction(gem);
f.index_operand = Some(1);
f.reduced_axes = Some(ReduceAxes::All);
assert_eq!(class_of(f, &same), Some(CellClass::DenseContraction));
let mut f = DeclaredFacts::indexed(1);
f.reduced_axes = Some(ReduceAxes::All);
assert_eq!(class_of(f, &same), Some(CellClass::IndexedAccess));
}
#[test]
fn pointwise_cells_classify_by_input_count_and_only_one_to_three() {
let d = || f32c(&[8, 16]);
let ops = |n_in: usize, n_out: usize| -> Vec<CellOperand> {
(0..n_in)
.map(|_| CellOperand::input(d()))
.chain((0..n_out).map(|_| CellOperand::output(d())))
.collect()
};
let cls = |n_in, n_out| derive_cell_class(&DeclaredFacts::none(), &ops(n_in, n_out));
assert_eq!(cls(1, 1), Some(CellClass::ElementwiseOne));
assert_eq!(cls(2, 1), Some(CellClass::ElementwiseTwo));
assert_eq!(cls(3, 1), Some(CellClass::ElementwiseThree));
assert_eq!(cls(4, 1), Some(CellClass::Other), "four inputs derive oth");
assert_eq!(
cls(0, 2),
Some(CellClass::Other),
"no input operand derives oth"
);
assert_eq!(cls(2, 2), Some(CellClass::ElementwiseTwo));
}
#[test]
fn a_cell_that_is_not_pointwise_by_shape_derives_oth() {
let bad_input = [f32c(&[4, 3]), f32c(&[4, 8])];
assert_eq!(
class_of(DeclaredFacts::none(), &bad_input),
Some(CellClass::Other)
);
let short_output = [f32c(&[128, 256]), f32c(&[64, 256])];
assert_eq!(
class_of(DeclaredFacts::none(), &short_output),
Some(CellClass::Other)
);
let low_rank_output = [f32c(&[4, 8]), f32c(&[8])];
assert_eq!(
class_of(DeclaredFacts::none(), &low_rank_output),
Some(CellClass::Other),
"a written operand must have the iteration RANK, not just extents"
);
let ok = [f32c(&[4, 8]), f32c(&[4, 8])];
assert_eq!(
class_of(DeclaredFacts::none(), &ok),
Some(CellClass::ElementwiseOne)
);
}
#[test]
fn broadcast_inputs_are_pointwise_but_a_broadcast_written_operand_is_not() {
let bcast = [f32c(&[256]), f32c(&[1, 256]), f32c(&[128, 256])];
assert_eq!(
class_of(DeclaredFacts::none(), &bcast),
Some(CellClass::ElementwiseTwo)
);
let ops = [
CellOperand::input(f32c(&[128, 256])),
CellOperand::in_place(f32c(&[1, 256])),
];
assert_eq!(
derive_cell_class(&DeclaredFacts::none(), &ops),
Some(CellClass::Other)
);
let ok = [
CellOperand::input(f32c(&[128, 256])),
CellOperand::in_place(f32c(&[128, 256])),
];
assert_eq!(
derive_cell_class(&DeclaredFacts::none(), &ok),
Some(CellClass::ElementwiseTwo)
);
}
#[test]
fn a_cell_not_in_canonical_order_declines() {
let d = f32c(&[8, 16]);
let out_first = [
CellOperand::output(d.clone()),
CellOperand::input(d.clone()),
];
assert_eq!(derive_cell_class(&DeclaredFacts::none(), &out_first), None);
assert_eq!(
derive_structure_key_token(&DeclaredFacts::none(), &out_first, "cuda:sm89", None),
None
);
let in_first = [CellOperand::input(d.clone()), CellOperand::output(d)];
assert!(
derive_structure_key_token(&DeclaredFacts::none(), &in_first, "cuda:sm89", None)
.is_some()
);
}
#[test]
fn an_index_slot_must_name_a_read_operand() {
let d = f32c(&[8, 16]);
let ops = canon(&[d.clone(), d.clone(), d]); assert_eq!(
derive_cell_class(&DeclaredFacts::indexed(1), &ops),
Some(CellClass::IndexedAccess)
);
assert_eq!(
derive_cell_class(&DeclaredFacts::indexed(2), &ops),
None,
"slot 2 is the output"
);
assert_eq!(
derive_cell_class(&DeclaredFacts::indexed(3), &ops),
None,
"slot 3 is out of range"
);
}
#[test]
fn the_gated_fields_follow_the_derived_class_not_the_declared_facts() {
let ops = [f32c(&[8, 16]), f32c(&[16, 16]), f32c(&[8, 16])];
let mut f = DeclaredFacts::contraction(gem_f32(8, 16, 16));
f.reduced_axes = Some(ReduceAxes::All);
let ctr = derive(f, &ops, "cuda:sm89").expect("derives");
let fields: Vec<&str> = ctr.split('|').collect();
assert_eq!(fields[1], "ctr");
assert_eq!(fields[8], "-", "reduce field is `-` on a ctr cell: {ctr}");
assert!(
fields[9].starts_with("c"),
"contraction field present: {ctr}"
);
assert_eq!(fields.len(), 10);
let mut g = DeclaredFacts::indexed(0);
g.reduced_axes = Some(ReduceAxes::TrailingAxis);
let ndx = derive(g, &ops, "cuda:sm89").expect("derives");
let fields: Vec<&str> = ndx.split('|').collect();
assert_eq!(fields[1], "ndx");
assert_eq!(fields[8], "-", "reduce field is `-` on an ndx cell: {ndx}");
assert_eq!(fields.len(), 9, "no contraction field on a non-ctr cell");
assert!(
ctr.contains("/v4/"),
"a declared-but-unused reduced set leaked into the v1 gate: {ctr}"
);
}
#[test]
fn an_empty_reduced_set_never_derives_rdc_with_a_dash_reduce_field() {
for (name, axes, descs) in [
(
"keepdim(0) rank 2",
ReduceAxes::Keepdim(0),
vec![f32c(&[4, 8]), f32c(&[4, 8])],
),
("all at rank 0", ReduceAxes::All, vec![f32c(&[]), f32c(&[])]),
(
"trailing at rank 0",
ReduceAxes::TrailingAxis,
vec![f32c(&[]), f32c(&[])],
),
] {
let token = derive(DeclaredFacts::reduction(axes), &descs, "cuda:sm89")
.unwrap_or_else(|| panic!("{name} must derive"));
let class = token.split('|').nth(1).unwrap();
assert_ne!(
class, "rdc",
"{name}: derived rdc with an empty set: {token}"
);
assert_eq!(token.rsplit('|').next(), Some("-"), "{name}: {token}");
}
let control = derive(
DeclaredFacts::reduction(ReduceAxes::All),
&[f32c(&[4, 8]), f32c(&[1, 1])],
"cuda:sm89",
)
.expect("derives");
assert_eq!(control.split('|').nth(1), Some("rdc"), "{control}");
}
#[test]
fn a_reduction_output_that_is_not_keepdim_declines() {
let typed = |axes: ReduceAxes, descs: Vec<FdxOperandDesc>| {
derive_structure_key_token_typed(
&DeclaredFacts::reduction(axes),
&canon(&descs),
"cuda:sm89",
None,
)
};
assert_eq!(
typed(ReduceAxes::Keepdim(0b10), vec![f32c(&[4, 8]), f32c(&[4])]),
Err(DeriveDecline::CollapsedReduction)
);
assert_eq!(
typed(
ReduceAxes::Keepdim(0b10),
vec![f32c(&[4, 8]), f32c(&[4, 8])]
),
Err(DeriveDecline::CollapsedReduction)
);
assert_eq!(
typed(ReduceAxes::All, vec![f32c(&[4, 8]), f32c(&[])]),
Err(DeriveDecline::CollapsedReduction)
);
for (axes, descs) in [
(
ReduceAxes::Keepdim(0b10),
vec![f32c(&[4, 8]), f32c(&[4, 1])],
),
(ReduceAxes::All, vec![f32c(&[4, 8]), f32c(&[1, 1])]),
(ReduceAxes::TrailingAxis, vec![f32c(&[4, 8]), f32c(&[4, 1])]),
] {
assert!(typed(axes, descs).is_ok());
}
}
#[test]
fn a_reduced_axis_beyond_the_frame_rank_declines() {
let typed = |mask: u8| {
derive_structure_key_token_typed(
&DeclaredFacts::reduction(ReduceAxes::Keepdim(mask)),
&canon(&[f32c(&[4, 8]), f32c(&[4, 1])]),
"cuda:sm89",
None,
)
};
assert_eq!(typed(0b100), Err(DeriveDecline::ReducedAxisOutOfRange));
assert_eq!(typed(0b110), Err(DeriveDecline::ReducedAxisOutOfRange));
assert_eq!(typed(0x80), Err(DeriveDecline::ReducedAxisOutOfRange));
assert!(typed(0b10).is_ok());
let wide = [f32c(&[1; 8]), f32c(&[1; 8])];
assert!(
derive_structure_key_token_typed(
&DeclaredFacts::reduction(ReduceAxes::Keepdim(0x80)),
&canon(&wide),
"cuda:sm89",
None,
)
.is_ok()
);
}
#[test]
fn typed_declines_name_their_reason() {
let facts = DeclaredFacts::none();
let ok = canon(&[f32c(&[8]), f32c(&[8])]);
let typed = |ops: &[CellOperand], target: &str| {
derive_structure_key_token_typed(&facts, ops, target, None)
};
assert_eq!(typed(&[], "cuda:sm89"), Err(DeriveDecline::OperandCount));
assert_eq!(typed(&ok, "cuda"), Err(DeriveDecline::InvalidTarget));
let scale = canon(&[
FdxOperandDesc::from_layout(&Layout::contiguous(Shape::from_dims(&[8])), DType::F8E8M0),
f32c(&[8]),
]);
assert_eq!(
typed(&scale, "cuda:sm89"),
Err(DeriveDecline::ScaleDtypeAtOperand0)
);
let reserved = canon(&[
FdxOperandDesc::from_layout(&Layout::contiguous(Shape::from_dims(&[8])), DType::F8E6M2),
f32c(&[8]),
]);
assert_eq!(
typed(&reserved, "cuda:sm89"),
Err(DeriveDecline::ReservedDtype)
);
let mut bad = canon(&[f32c(&[8]), f32c(&[8])]);
bad.swap(0, 1);
assert_eq!(
typed(&bad, "cuda:sm89"),
Err(DeriveDecline::NotCanonicalOrder)
);
assert_eq!(
derive_structure_key_token(&facts, &ok, "cuda:sm89", None),
typed(&ok, "cuda:sm89").ok()
);
assert_eq!(
derive_structure_key_token(&facts, &bad, "cuda:sm89", None),
None
);
}
#[test]
fn dequant_form_is_the_optional_sixth_subkey_component() {
let descs = [f32c(&[8, 4096]), f32c(&[4096, 4096]), f32c(&[8, 4096])];
let facts = DeclaredFacts::contraction(gem_f32(8, 4096, 4096));
let build = |form: Option<DequantForm>| -> String {
let mut ops = canon(&descs);
if let Some(f) = form {
ops[0] = ops[0].clone().with_dequant_form(f);
}
derive_structure_key_token(&facts, &ops, "cuda:sm90", None).expect("derives")
};
let plain = build(None);
let linear = build(Some(DequantForm::Linear));
let codebook = build(Some(DequantForm::Codebook));
assert!(
linear.contains("|co/00/v4/d16/f/dl;co/00/v4/d16/f;co/00/v4/d16/f|"),
"{linear}"
);
assert!(
codebook.contains("|co/00/v4/d16/f/dc;co/00/v4/d16/f;co/00/v4/d16/f|"),
"{codebook}"
);
assert_eq!(linear.replacen("/dl;", ";", 1), plain);
assert_eq!(codebook.replacen("/dc;", ";", 1), plain);
assert_ne!(linear, codebook);
assert_ne!(linear, plain);
assert_ne!(codebook, plain);
}
#[test]
fn a_sub_key_never_exceeds_six_parts() {
let descs = [f32c(&[8, 16]), f32c(&[16, 16]), f32c(&[8, 16])];
let ops: Vec<CellOperand> = canon(&descs)
.into_iter()
.map(|o| o.with_dequant_form(DequantForm::Codebook))
.collect();
let token = derive_structure_key_token(
&DeclaredFacts::contraction(gem_f32(8, 16, 16)),
&ops,
"cuda:sm89",
None,
)
.expect("derives");
let sub_keys = token.split('|').nth(7).unwrap();
for sk in sub_keys.split(';') {
let parts = sk.split('/').count();
assert_eq!(parts, 6, "sub-key `{sk}` has {parts} parts: {token}");
assert!(sk.ends_with("/dc"), "{sk}");
}
}
#[test]
fn an_extent_one_axis_sets_its_mask_bit_whatever_stride_is_recorded() {
let unit = |lead_stride: isize| {
FdxOperandDesc::from_layout(
&Layout::new(
Shape::from(vec![1usize, 256]),
[lead_stride, 1].into_iter().collect::<StrideVec>(),
0,
),
DType::F32,
)
};
let frame_setter = f32c(&[128, 256]);
let key = |d: FdxOperandDesc| {
derive(
DeclaredFacts::none(),
&[d, frame_setter.clone(), frame_setter.clone()],
"cuda:sm89",
)
.expect("derives")
};
let zero = key(unit(0));
let nonzero = key(unit(256));
let odd = key(unit(7));
assert_eq!(
nonzero, zero,
"one tensor, one spelling: the recorded stride of a unit axis is arbitrary"
);
assert_eq!(odd, zero);
assert!(
zero.contains("|co/01/v4/d16/f;"),
"bit 0 must be set for the unit axis: {zero}"
);
}
#[test]
fn an_extent_one_axis_in_a_unit_frame_axis_sets_no_bit() {
let unit = FdxOperandDesc::from_layout(
&Layout::new(
Shape::from(vec![1usize, 256]),
[256isize, 1].into_iter().collect::<StrideVec>(),
0,
),
DType::F32,
);
let token =
derive(DeclaredFacts::none(), &[unit.clone(), unit], "cuda:sm89").expect("derives");
assert!(
token.contains("|co/00/v4/d16/f;co/00/v4/d16/f|"),
"frame extent 1 must leave the bit clear: {token}"
);
}
}