use baracuda_kernels_types::{
ArchSku, ElementKind, MAX_RANK, OpCategory, OperandDesc, structure_key_token,
};
use fuel_ir::DType;
use super::structure_key::{FdxOperandDesc, StructureKeyProvider, StructureKeyToken};
#[derive(Debug, Clone, Copy, Default)]
pub struct BaracudaStructureKeyProvider;
impl StructureKeyProvider for BaracudaStructureKeyProvider {
fn structure_key(
&self,
op_class: &str,
operands: &[FdxOperandDesc],
arch: &str,
) -> Option<StructureKeyToken> {
let op = map_op_category(op_class, operands.len())?;
let arch = map_arch_sku(arch)?;
let mut mapped = Vec::with_capacity(operands.len());
for od in operands {
mapped.push(map_operand(od)?);
}
Some(StructureKeyToken(structure_key_token(op, &mapped, arch)))
}
}
fn map_operand(od: &FdxOperandDesc) -> Option<OperandDesc> {
let dtype = map_element_kind(od.dtype)?;
let rank = od.shape.len();
if rank > MAX_RANK || od.strides.len() != rank {
return None;
}
Some(OperandDesc::new(
rank,
&od.shape,
&od.strides,
dtype,
od.align_bytes,
))
}
fn map_element_kind(dt: DType) -> Option<ElementKind> {
Some(match dt {
DType::U8 => ElementKind::U8,
DType::I8 => ElementKind::I8,
DType::I32 => ElementKind::I32,
DType::I64 => ElementKind::I64,
DType::BF16 => ElementKind::Bf16,
DType::F16 => ElementKind::F16,
DType::F32 => ElementKind::F32,
DType::F64 => ElementKind::F64,
DType::F8E4M3 => ElementKind::Fp8E4M3FN,
DType::F8E5M2 => ElementKind::Fp8E5M2,
DType::U32 => ElementKind::U32,
DType::Bool => ElementKind::Bool,
DType::I16 | DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 | DType::F8E6M2 => {
return None;
}
})
}
fn map_arch_sku(arch: &str) -> Option<ArchSku> {
let digits = arch
.strip_prefix("sm_")
.or_else(|| arch.strip_prefix("sm"))?;
Some(match digits {
"80" => ArchSku::Sm80,
"89" => ArchSku::Sm89,
"90" => ArchSku::Sm90,
"90a" => ArchSku::Sm90a,
_ => return None,
})
}
#[allow(dead_code)]
fn arch_sku_digits(sku: ArchSku) -> &'static str {
match sku {
ArchSku::Sm80 => "80",
ArchSku::Sm89 => "89",
ArchSku::Sm90 => "90",
ArchSku::Sm90a => "90a",
}
}
fn map_op_category(op_class: &str, n_inputs: usize) -> Option<OpCategory> {
let cat = match op_class {
"MatMul" | "FusedLinear" | "QMatMul" | "Nf4Matmul" => OpCategory::Gemm,
"Conv2D" | "ConvTranspose2D" | "CausalConv1d" => OpCategory::Convolution,
"FlashAttn" | "FlashAttnBackwardQ" | "FlashAttnBackwardK" | "FlashAttnBackwardV"
| "PagedAttn" | "Rope" => OpCategory::Attention,
"SoftmaxLastDim"
| "SoftmaxLastDimBackward"
| "LogSoftmaxLastDim"
| "LogSoftmaxLastDimBackward" => OpCategory::Softmax,
"RmsNormLastDim"
| "RmsNormLastDimBackward"
| "LayerNormLastDim"
| "LayerNormLastDimBackward" => OpCategory::Normalization,
"SumReduce"
| "MaxReduce"
| "MinReduce"
| "MeanReduce"
| "ReduceSumTo"
| "ReduceMaxTo"
| "ReduceMaxToBackward"
| "ArgMaxDim"
| "ArgMinDim" => OpCategory::Reduction,
"CumSum" | "SelectiveScan" | "SsdChunkScan" => OpCategory::Scan,
"IndexSelect" | "Gather" | "IndexAdd" | "ScatterAdd" | "MaskedFill" => OpCategory::Indexing,
"Flip" | "Roll" | "Pad" | "PadBackward" | "Triu" | "Tril" | "Concat" | "Copy"
| "WriteSlice" | "WriteSliceRotating" | "WriteSliceDoff" => OpCategory::ShapeLayout,
"FusedSoftmaxCrossEntropy" => OpCategory::Loss,
"Where" => OpCategory::TernaryElementwise,
"Cast" | "Affine" | "InplaceAffine" => OpCategory::UnaryElementwise,
other if other.ends_with("Elementwise") || other.ends_with("Inplace") => match n_inputs {
1 => OpCategory::UnaryElementwise,
2 => OpCategory::BinaryElementwise,
3 => OpCategory::TernaryElementwise,
_ => return None,
},
_ => return None,
};
Some(cat)
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::{Layout, Shape, StrideVec};
fn contig_of(dims: &[usize], dt: DType) -> FdxOperandDesc {
FdxOperandDesc::from_layout(&Layout::contiguous(Shape::from(dims.to_vec())), dt)
}
fn contig_f16(dims: &[usize]) -> FdxOperandDesc {
contig_of(dims, DType::F16)
}
fn flipped_f16(dims: &[usize]) -> FdxOperandDesc {
let mut stride: StrideVec = Shape::from(dims.to_vec()).stride_contiguous();
if let Some(last) = stride.last_mut() {
*last = -*last;
}
let offset = dims.last().copied().unwrap_or(1).saturating_sub(1);
let layout = Layout::new(Shape::from(dims.to_vec()), stride, offset);
FdxOperandDesc::from_layout(&layout, DType::F16)
}
#[test]
fn contiguous_f16_pair_yields_stable_nonempty_token() {
let p = BaracudaStructureKeyProvider;
let ops = [contig_f16(&[128, 256]), contig_f16(&[128, 256])];
let t1 = p
.structure_key("MatMul", &ops, "sm_89")
.expect("linked provider must yield a token");
assert!(!t1.0.is_empty(), "token must be non-empty");
let t2 = p
.structure_key("MatMul", &ops, "sm_89")
.expect("token on the second call too");
assert_eq!(t1, t2, "structure_key is deterministic");
}
#[test]
fn different_operand_structure_yields_different_token() {
let p = BaracudaStructureKeyProvider;
let contig = [contig_f16(&[128, 256]), contig_f16(&[128, 256])];
let flipped = [flipped_f16(&[128, 256]), contig_f16(&[128, 256])];
let a = p.structure_key("MatMul", &contig, "sm_89").unwrap();
let b = p.structure_key("MatMul", &flipped, "sm_89").unwrap();
assert_ne!(a, b, "a flipped operand must key differently");
}
#[test]
fn maps_fdx_operand_desc_to_baracuda_operand_desc_field_for_field() {
let od = contig_f16(&[8, 16]);
let mapped = map_operand(&od).expect("mappable");
assert_eq!(mapped.rank, 2);
assert_eq!(&mapped.shape[..2], &[8i64, 16]);
assert_eq!(&mapped.strides[..2], &[16i64, 1]);
assert_eq!(mapped.dtype, ElementKind::F16);
assert_eq!(mapped.align_bytes, od.align_bytes);
assert!(
mapped.symbolic.is_none(),
"v1 does not fabricate symbolic facts"
);
}
#[test]
fn a_u32_operand_now_derives_a_key_where_it_previously_derived_none() {
let p = BaracudaStructureKeyProvider;
let u32s = [
contig_of(&[128, 256], DType::U32),
contig_of(&[128, 256], DType::U32),
];
let t = p
.structure_key("MatMul", &u32s, "sm_89")
.expect("a u32 operand must derive a key — before GAP-171 this was None");
assert!(!t.0.is_empty(), "token must be non-empty");
let f16s = [
contig_of(&[128, 256], DType::F16),
contig_of(&[128, 256], DType::F16),
];
assert_ne!(
t,
p.structure_key("MatMul", &f16s, "sm_89")
.expect("f16 derives"),
"u32 must key DIFFERENTLY from f16, or the dtype is being dropped",
);
for dt in [DType::I16, DType::F8E8M0, DType::F8E6M2] {
assert_eq!(
map_element_kind(dt),
None,
"{dt:?} has no faithful ElementKind at alpha.78 and must stay declined",
);
}
}
#[test]
fn element_kind_mapping() {
assert_eq!(map_element_kind(DType::F16), Some(ElementKind::F16));
assert_eq!(map_element_kind(DType::BF16), Some(ElementKind::Bf16));
assert_eq!(map_element_kind(DType::I8), Some(ElementKind::I8));
assert_eq!(
map_element_kind(DType::F8E4M3),
Some(ElementKind::Fp8E4M3FN)
);
assert_eq!(map_element_kind(DType::F8E5M2), Some(ElementKind::Fp8E5M2));
assert_eq!(map_element_kind(DType::U32), Some(ElementKind::U32));
assert_eq!(map_element_kind(DType::F4), None);
}
#[test]
fn every_known_arch_sku_round_trips_through_its_canonical_digits() {
const KNOWN: &[ArchSku] = &[ArchSku::Sm80, ArchSku::Sm89, ArchSku::Sm90a];
for &sku in KNOWN {
let digits = arch_sku_digits(sku);
assert_eq!(
map_arch_sku(&format!("sm_{digits}")),
Some(sku),
"sm_{digits} must select {sku:?}"
);
assert_eq!(
map_arch_sku(&format!("sm{digits}")),
Some(sku),
"sm{digits} (bare form) must select {sku:?}"
);
}
let mut seen = std::collections::BTreeSet::new();
for &sku in KNOWN {
assert!(
seen.insert(arch_sku_digits(sku)),
"two ArchSku variants share canonical digits {:?}",
arch_sku_digits(sku)
);
}
assert_eq!(seen.len(), 3, "non-vacuity: the known-SKU set is not empty");
}
#[test]
fn sm90a_canonical_digits_are_90a_not_90() {
assert_eq!(arch_sku_digits(ArchSku::Sm90a), "90a");
}
#[test]
fn arch_and_op_class_mapping() {
assert_eq!(map_arch_sku("sm_80"), Some(ArchSku::Sm80));
assert_eq!(map_arch_sku("sm_89"), Some(ArchSku::Sm89));
assert_eq!(map_arch_sku("sm_90"), Some(ArchSku::Sm90));
assert_eq!(map_arch_sku("sm89"), Some(ArchSku::Sm89));
assert_eq!(map_arch_sku("cpu"), None);
assert!(matches!(
map_op_category("MatMul", 2),
Some(OpCategory::Gemm)
));
assert!(matches!(
map_op_category("AddElementwise", 2),
Some(OpCategory::BinaryElementwise)
));
assert!(matches!(
map_op_category("ReluElementwise", 1),
Some(OpCategory::UnaryElementwise)
));
assert!(matches!(
map_op_category("Where", 3),
Some(OpCategory::TernaryElementwise)
));
assert_eq!(map_op_category("TotallyUnknownOp", 1), None);
}
#[test]
fn unmappable_inputs_yield_none() {
let p = BaracudaStructureKeyProvider;
let f16 = [contig_f16(&[8, 16])];
assert!(
p.structure_key("TotallyUnknownOp", &f16, "sm_89").is_none(),
"unmapped op family ⇒ no key"
);
assert!(
p.structure_key("ReluElementwise", &f16, "cpu").is_none(),
"CPU arch has no Baracuda build matrix ⇒ no key"
);
let i16op = [contig_of(&[8, 16], DType::I16)];
assert!(
p.structure_key("ReluElementwise", &i16op, "sm_89")
.is_none(),
"unmappable dtype ⇒ no key"
);
}
#[test]
fn over_rank_operand_declines_never_panics() {
let p = BaracudaStructureKeyProvider;
let big = contig_f16(&[2, 2, 2, 2, 2, 2, 2, 2, 2]);
assert!(
p.structure_key("ReluElementwise", &[big], "sm_89")
.is_none(),
"rank > MAX_RANK must decline, not panic"
);
}
}